refactor(redis-store): migrate MCP credential refresh locks

Co-authored-by: Yujong Lee <yujong@berri.ai>

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

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:

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())))

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

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