From 89d9a7649e7ac1390db11bb24f41333677c9f841 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 16:15:18 +0000 Subject: [PATCH] refactor(redis-store): migrate MCP credential refresh locks Co-authored-by: Yujong Lee --- .../redis_distributed_lock.py | 84 ++------- .../redis_refresh_coordinator.py | 7 +- .../runtime_refresh_coordinator.py | 10 +- .../redis/stores/mcp_credential_refresh.py | 46 +++++ .../test_redis_distributed_lock.py | 159 ++++++++---------- 5 files changed, 136 insertions(+), 170 deletions(-) create mode 100644 litellm/redis/stores/mcp_credential_refresh.py diff --git a/litellm/proxy/mcp_server/outbound_credentials/redis_distributed_lock.py b/litellm/proxy/mcp_server/outbound_credentials/redis_distributed_lock.py index 7262cb6f100..06d0248cf2c 100644 --- a/litellm/proxy/mcp_server/outbound_credentials/redis_distributed_lock.py +++ b/litellm/proxy/mcp_server/outbound_credentials/redis_distributed_lock.py @@ -1,94 +1,44 @@ -"""Concrete ``DistributedLock`` over a Redis client: ``SET NX PX`` / owner-only renew / delete. +"""Concrete ``DistributedLock`` over ``McpRefreshLockStore``: ``SET NX PX`` / owner-only renew / delete. The cross-replica lock the ``RedisRefreshCoordinator`` elects refreshers with. ``acquire`` is an atomic ``SET key token NX PX ttl`` (only the first caller wins; the entry self-expires so a crashed holder can't wedge refresh). ``extend`` renews the lease only when the token still matches, and ``release`` deletes the key only when it still holds this caller's token, so a holder whose lock already -PX-expired and was re-acquired by another worker cannot delete the new holder's lock. ``is_held`` is -``EXISTS``. Every key is run through the injected ``namespace_key`` before it reaches Redis, so lock -keys carry the same namespace as cache keys and cannot collide with another deployment sharing Redis. +PX-expired and was re-acquired by another worker cannot delete the new holder's lock. -The Redis client is injected (in production the async client from LiteLLM's ``RedisCache``), so the -lock is unit-testable with a fake. A transport error on ``acquire`` returns ``LockAcquisition.ERROR`` - -distinct from ``HELD`` - so the coordinator refreshes anyway instead of mistaking a dead backend for a -busy holder; a Redis blip degrades to an extra refresh, never a stale bearer. +A Redis failure on ``acquire`` returns ``LockAcquisition.ERROR`` - distinct from ``HELD`` - so the +coordinator refreshes anyway instead of mistaking a dead backend for a busy holder; a Redis blip +degrades to an extra refresh, never a stale bearer. """ from __future__ import annotations -from collections.abc import Callable -from dataclasses import KW_ONLY, dataclass -from typing import Final, Protocol +from dataclasses import dataclass +from datetime import timedelta +from typing import Final -from litellm._logging import verbose_logger from litellm.proxy.mcp_server.outbound_credentials.redis_refresh_coordinator import ( LockAcquisition, ) - -# Delete the key only if it still holds this caller's token, so a holder whose lock already expired -# (PX) and was re-acquired by another worker cannot delete the new holder's lock. -_RELEASE_IF_OWNER = "if redis.call('get', KEYS[1]) == ARGV[1] then return redis.call('del', KEYS[1]) else return 0 end" -_EXTEND_IF_OWNER: Final = ( - "if redis.call('get', KEYS[1]) == ARGV[1] then return redis.call('pexpire', KEYS[1], ARGV[2]) else return 0 end" -) - - -class RedisCommands(Protocol): - """The slice of the async Redis client this lock needs.""" - - async def set(self, name: str, value: str, *, nx: bool = False, px: int | None = None) -> object | None: ... - - async def eval(self, script: str, numkeys: int, *keys_and_args: str) -> object: ... - - async def exists(self, *names: str) -> int: ... +from litellm.redis.stores.mcp_credential_refresh import McpRefreshLockStore +from litellm.redis.transport.contract import RedisFailure @dataclass(frozen=True, slots=True) class RedisDistributedLock: - client: RedisCommands - _: KW_ONLY - namespace_key: Callable[[str], str] = lambda key: key + store: McpRefreshLockStore async def acquire(self, key: str, token: str, ttl_seconds: float) -> LockAcquisition: - try: - result: Final = await self.client.set(self.namespace_key(key), token, nx=True, px=int(ttl_seconds * 1000)) - # Degrade on any Redis client error: redis.exceptions narrows only via an import that - # is Unknown under basedpyright, and the lock must never crash the resolve path. - except Exception as exc: # noqa: BLE001 - verbose_logger.warning("RedisDistributedLock.acquire failed: %s", exc) + acquired: Final = await self.store.acquire(key, token, ttl=timedelta(seconds=ttl_seconds)) + if isinstance(acquired, RedisFailure): return LockAcquisition.ERROR - return LockAcquisition.ACQUIRED if result is not None else LockAcquisition.HELD + return LockAcquisition.ACQUIRED if acquired else LockAcquisition.HELD async def extend(self, key: str, token: str, ttl_seconds: float) -> bool: - try: - result: Final = await self.client.eval( - _EXTEND_IF_OWNER, - 1, - self.namespace_key(key), - token, - str(int(ttl_seconds * 1000)), - ) - # Degrade on any Redis client error: redis.exceptions narrows only via an import that - # is Unknown under basedpyright, and the lock must never crash the resolve path. - except Exception as exc: # noqa: BLE001 - verbose_logger.warning("RedisDistributedLock.extend failed: %s", exc) - return False - return result == 1 + return await self.store.extend(key, token, ttl=timedelta(seconds=ttl_seconds)) is True async def release(self, key: str, token: str) -> None: - try: - await self.client.eval(_RELEASE_IF_OWNER, 1, self.namespace_key(key), token) - # Degrade on any Redis client error: redis.exceptions narrows only via an import that - # is Unknown under basedpyright, and the lock must never crash the resolve path. - except Exception as exc: # noqa: BLE001 - verbose_logger.warning("RedisDistributedLock.release failed: %s", exc) + _ = await self.store.release(key, token) async def is_held(self, key: str) -> bool: - try: - return await self.client.exists(self.namespace_key(key)) > 0 - # Degrade on any Redis client error: redis.exceptions narrows only via an import that - # is Unknown under basedpyright, and the lock must never crash the resolve path. - except Exception as exc: # noqa: BLE001 - # On error, report "not held" so a waiter stops waiting and re-reads rather than blocking. - verbose_logger.warning("RedisDistributedLock.is_held failed: %s", exc) - return False + return await self.store.is_held(key) is True diff --git a/litellm/proxy/mcp_server/outbound_credentials/redis_refresh_coordinator.py b/litellm/proxy/mcp_server/outbound_credentials/redis_refresh_coordinator.py index 98446f2b66e..f31f775639d 100644 --- a/litellm/proxy/mcp_server/outbound_credentials/redis_refresh_coordinator.py +++ b/litellm/proxy/mcp_server/outbound_credentials/redis_refresh_coordinator.py @@ -31,6 +31,7 @@ from typing import Final, Protocol from litellm.proxy.mcp_server.outbound_credentials.oauth_token_store import ( OAuthToken, ) +from litellm.redis.stores.mcp_credential_refresh import refresh_lock_key class LockAcquisition(Enum): @@ -63,7 +64,6 @@ class DistributedLock(Protocol): class RedisRefreshCoordinator: lock: DistributedLock _: KW_ONLY - key_prefix: str = "mcp:refresh_lock:" lock_ttl_seconds: float = 10.0 # The holder renews its lease while a slow token endpoint runs, but only up to this budget; past it # it stops renewing and the lock lapses, so a hung refresh degrades to "maybe an extra refresh" @@ -79,9 +79,6 @@ class RedisRefreshCoordinator: clock: Callable[[], float] = time.monotonic new_token: Callable[[], str] = lambda: uuid.uuid4().hex - def _key(self, user_id: str, server_id: str) -> str: - return f"{self.key_prefix}{user_id}:{server_id}" - async def run( self, user_id: str, @@ -89,7 +86,7 @@ class RedisRefreshCoordinator: refresh: Callable[[], Awaitable[OAuthToken | None]], reread: Callable[[], Awaitable[OAuthToken | None]], ) -> OAuthToken | None: - key: Final = self._key(user_id, server_id) + key: Final = refresh_lock_key(user_id, server_id) token: Final = self.new_token() match await self.lock.acquire(key, token, self.lock_ttl_seconds): case LockAcquisition.ACQUIRED: diff --git a/litellm/proxy/mcp_server/outbound_credentials/runtime_refresh_coordinator.py b/litellm/proxy/mcp_server/outbound_credentials/runtime_refresh_coordinator.py index 9376d1bda50..dd1ec88091a 100644 --- a/litellm/proxy/mcp_server/outbound_credentials/runtime_refresh_coordinator.py +++ b/litellm/proxy/mcp_server/outbound_credentials/runtime_refresh_coordinator.py @@ -23,6 +23,7 @@ from litellm.proxy.mcp_server.outbound_credentials.redis_distributed_lock import from litellm.proxy.mcp_server.outbound_credentials.redis_refresh_coordinator import ( RedisRefreshCoordinator, ) +from litellm.redis.stores.mcp_credential_refresh import TransportMcpRefreshLockStore def runtime_refresh_coordinator() -> RefreshCoordinator | None: @@ -31,11 +32,4 @@ def runtime_refresh_coordinator() -> RefreshCoordinator | None: redis_cache: Final = user_api_key_cache.redis_cache if redis_cache is None: return None - # The Redis client from init_async_client() is only partially typed; the lock validates every - # reply it depends on, so the untyped boundary is contained here. - redis_client: Final = redis_cache.init_async_client() # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType] # litellm redis wrapper is untyped - lock: Final = RedisDistributedLock( - redis_client, # pyright: ignore[reportArgumentType,reportUnknownArgumentType] # litellm redis wrapper is untyped - namespace_key=redis_cache.check_and_fix_namespace, - ) - return RedisRefreshCoordinator(lock) + return RedisRefreshCoordinator(RedisDistributedLock(TransportMcpRefreshLockStore(redis_cache.transport()))) diff --git a/litellm/redis/stores/mcp_credential_refresh.py b/litellm/redis/stores/mcp_credential_refresh.py new file mode 100644 index 00000000000..e968c3634e6 --- /dev/null +++ b/litellm/redis/stores/mcp_credential_refresh.py @@ -0,0 +1,46 @@ +from dataclasses import dataclass +from datetime import timedelta +from typing import Final, Protocol + +from litellm.redis.transport.contract import RedisFailure, RedisTransport, ttl_ms + +_RELEASE_IF_OWNER: Final = ( + "if redis.call('get', KEYS[1]) == ARGV[1] then return redis.call('del', KEYS[1]) else return 0 end" +) +_EXTEND_IF_OWNER: Final = ( + "if redis.call('get', KEYS[1]) == ARGV[1] then return redis.call('pexpire', KEYS[1], ARGV[2]) else return 0 end" +) + + +def refresh_lock_key(user_id: str, server_id: str) -> str: + return f"mcp:refresh_lock:{user_id}:{server_id}" + + +class McpRefreshLockStore(Protocol): + async def acquire(self, key: str, owner: str, *, ttl: timedelta) -> bool | RedisFailure: ... + + async def extend(self, key: str, owner: str, *, ttl: timedelta) -> bool | RedisFailure: ... + + async def release(self, key: str, owner: str) -> bool | RedisFailure: ... + + async def is_held(self, key: str) -> bool | RedisFailure: ... + + +@dataclass(frozen=True, slots=True) +class TransportMcpRefreshLockStore: + transport: RedisTransport + + async def acquire(self, key: str, owner: str, *, ttl: timedelta) -> bool | RedisFailure: + return await self.transport.set(key, owner, ttl=ttl, nx=True) + + async def extend(self, key: str, owner: str, *, ttl: timedelta) -> bool | RedisFailure: + reply: Final = await self.transport.eval(_EXTEND_IF_OWNER, (key,), (owner, ttl_ms(ttl))) + return reply if isinstance(reply, RedisFailure) else reply == 1 + + async def release(self, key: str, owner: str) -> bool | RedisFailure: + reply: Final = await self.transport.eval(_RELEASE_IF_OWNER, (key,), (owner,)) + return reply if isinstance(reply, RedisFailure) else reply == 1 + + async def is_held(self, key: str) -> bool | RedisFailure: + reply: Final = await self.transport.get(key) + return reply if isinstance(reply, RedisFailure) else reply is not None diff --git a/tests/unit/proxy/mcp_server/outbound_credentials/test_redis_distributed_lock.py b/tests/unit/proxy/mcp_server/outbound_credentials/test_redis_distributed_lock.py index d50ef6d6ba6..c74d2048b63 100644 --- a/tests/unit/proxy/mcp_server/outbound_credentials/test_redis_distributed_lock.py +++ b/tests/unit/proxy/mcp_server/outbound_credentials/test_redis_distributed_lock.py @@ -1,5 +1,6 @@ """Tests for the Redis lock: acquire NX/PX with a token, compare-and-delete release, namespacing.""" +import fakeredis import pytest from litellm.proxy.mcp_server.outbound_credentials.redis_distributed_lock import ( @@ -8,119 +9,97 @@ from litellm.proxy.mcp_server.outbound_credentials.redis_distributed_lock import from litellm.proxy.mcp_server.outbound_credentials.redis_refresh_coordinator import ( LockAcquisition, ) +from litellm.redis.stores.mcp_credential_refresh import TransportMcpRefreshLockStore +from tests._support.fake_redis import fake_redis_cache -class _FakeRedis: - """Models just enough Redis to exercise SET NX (token store) and the compare-and-delete EVAL.""" +@pytest.fixture +def server() -> fakeredis.FakeServer: + return fakeredis.FakeServer() - def __init__(self, set_returns=True, exists_returns=1, raise_on=()): - self._set_returns = set_returns - self._exists_returns = exists_returns - self._raise_on = set(raise_on) - self.set_calls = [] - self.eval_calls = [] - self.expire_calls = [] - self.deleted = [] - self.store: dict = {} - async def set(self, name, value, *, nx=False, px=None): - if "set" in self._raise_on: - raise RuntimeError("redis down") - self.set_calls.append((name, value, nx, px)) - if self._set_returns: - self.store[name] = value - return self._set_returns +@pytest.fixture +def raw(server: fakeredis.FakeServer) -> fakeredis.FakeAsyncRedis: + return fakeredis.FakeAsyncRedis(server=server) - async def eval(self, script, numkeys, *keys_and_args): - if "eval" in self._raise_on: - raise RuntimeError("redis down") - self.eval_calls.append((numkeys, keys_and_args)) - key, token = keys_and_args[0], keys_and_args[1] - if len(keys_and_args) == 3: - ttl_ms = keys_and_args[2] - if self.store.get(key) == token: - self.expire_calls.append((key, ttl_ms)) - return 1 - return 0 - if self.store.get(key) == token: # compare-and-delete: only the owner deletes - del self.store[key] - self.deleted.append(key) - return 1 - return 0 - async def exists(self, *names): - if "exists" in self._raise_on: - raise RuntimeError("redis down") - return self._exists_returns +def _lock(server: fakeredis.FakeServer, namespace: str | None = None) -> RedisDistributedLock: + return RedisDistributedLock(TransportMcpRefreshLockStore(fake_redis_cache(server, namespace=namespace).transport())) @pytest.mark.asyncio -async def test_acquire_sets_token_with_nx_px_and_reports_acquired(): - redis = _FakeRedis(set_returns=True) - assert await RedisDistributedLock(redis).acquire("k", "tok-1", 10.0) is (LockAcquisition.ACQUIRED) - assert redis.set_calls == [("k", "tok-1", True, 10000)] # token value, NX, px in ms +async def test_acquire_stores_the_token_under_the_namespace_with_the_ttl( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + assert await _lock(server, namespace="ns").acquire("k", "tok-1", 10.0) is LockAcquisition.ACQUIRED + + assert await raw.keys("*") == [b"ns:k"] + assert await raw.get("ns:k") == b"tok-1" + assert 9_000 < await raw.pttl("ns:k") <= 10_000 @pytest.mark.asyncio -async def test_acquire_reports_held_when_key_already_held(): - # redis SET NX returns None when the key exists -> another worker holds it. - assert await RedisDistributedLock(_FakeRedis(set_returns=None)).acquire("k", "tok", 10.0) is LockAcquisition.HELD +async def test_acquire_reports_held_when_key_already_held(server: fakeredis.FakeServer) -> None: + lock = _lock(server) + assert await lock.acquire("k", "tok-a", 10.0) is LockAcquisition.ACQUIRED + assert await lock.acquire("k", "tok-b", 10.0) is LockAcquisition.HELD @pytest.mark.asyncio -async def test_acquire_reports_error_on_redis_error_distinct_from_held(): - # A dead backend must be distinguishable from a busy holder so the coordinator refreshes anyway. - assert await RedisDistributedLock(_FakeRedis(raise_on=["set"])).acquire("k", "tok", 10.0) is LockAcquisition.ERROR +async def test_acquire_reports_error_on_redis_error_distinct_from_held(server: fakeredis.FakeServer) -> None: + lock = _lock(server) + server.connected = False + assert await lock.acquire("k", "tok", 10.0) is LockAcquisition.ERROR @pytest.mark.asyncio -async def test_release_deletes_only_when_the_token_matches(): - # Regression: a holder whose lock PX-expired and was re-acquired by another worker must not be - # able to delete the new holder's lock. release with a stale token is a no-op. - redis = _FakeRedis() - lock = RedisDistributedLock(redis) - await lock.acquire("k", "owner-B", 10.0) # B currently holds the lock - await lock.release("k", "owner-A") # A's stale token - assert redis.deleted == [] and redis.store.get("k") == "owner-B" # B's lock survives - await lock.release("k", "owner-B") # the real owner releases - assert redis.deleted == ["k"] and "k" not in redis.store - - -@pytest.mark.asyncio -async def test_keys_are_namespaced_before_reaching_redis(): - redis = _FakeRedis() - lock = RedisDistributedLock(redis, namespace_key=lambda key: f"ns:{key}") - await lock.acquire("k", "tok", 10.0) - await lock.extend("k", "tok", 10.0) - await lock.release("k", "tok") - await lock.is_held("k") - assert redis.set_calls[0][0] == "ns:k" # acquire namespaced - assert redis.eval_calls[0][1][0] == "ns:k" # extend (EVAL KEYS[1]) namespaced - assert redis.eval_calls[1][1][0] == "ns:k" # release (EVAL KEYS[1]) namespaced - assert redis.deleted == ["ns:k"] - - -@pytest.mark.asyncio -async def test_extend_refreshes_ttl_only_when_the_token_matches(): - redis = _FakeRedis() - lock = RedisDistributedLock(redis) +async def test_release_deletes_only_when_the_token_matches( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + lock = _lock(server) await lock.acquire("k", "owner-B", 10.0) + await lock.release("k", "owner-A") + assert await raw.get("k") == b"owner-B" + await lock.release("k", "owner-B") + assert await raw.exists("k") == 0 + + +@pytest.mark.asyncio +async def test_a_lock_the_old_writer_took_is_held_renewed_and_released_by_its_owner( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + await raw.set("ns:k", "owner-old", nx=True, px=10_000) + lock = _lock(server, namespace="ns") + + assert await lock.is_held("k") is True + assert await lock.acquire("k", "owner-new", 10.0) is LockAcquisition.HELD + assert await lock.extend("k", "owner-old", 30.0) is True + assert 29_000 < await raw.pttl("ns:k") <= 30_000 + await lock.release("k", "owner-old") + assert await lock.is_held("k") is False + + +@pytest.mark.asyncio +async def test_extend_refreshes_ttl_only_when_the_token_matches( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + lock = _lock(server) + await lock.acquire("k", "owner-B", 1.0) assert await lock.extend("k", "owner-A", 10.0) is False + assert await raw.pttl("k") <= 1_000 assert await lock.extend("k", "owner-B", 10.0) is True - assert redis.expire_calls == [("k", "10000")] + assert 9_000 < await raw.pttl("k") <= 10_000 @pytest.mark.asyncio -async def test_extend_degrades_to_false_on_redis_error(): - assert await RedisDistributedLock(_FakeRedis(raise_on=["eval"])).extend("k", "tok", 10.0) is False +async def test_extend_release_and_is_held_degrade_on_redis_error(server: fakeredis.FakeServer) -> None: + lock = _lock(server) + await lock.acquire("k", "tok", 10.0) + server.connected = False + assert await lock.extend("k", "tok", 10.0) is False + await lock.release("k", "tok") + assert await lock.is_held("k") is False -@pytest.mark.asyncio -async def test_is_held_reflects_exists(): - assert await RedisDistributedLock(_FakeRedis(exists_returns=1)).is_held("k") is True - assert await RedisDistributedLock(_FakeRedis(exists_returns=0)).is_held("k") is False - - -@pytest.mark.asyncio -async def test_is_held_degrades_to_false_on_redis_error(): - assert await RedisDistributedLock(_FakeRedis(raise_on=["exists"])).is_held("k") is False + server.connected = True + assert await lock.is_held("k") is True