From b8e43df44b69da56c4a603517cae2454fa5dda31 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 30 Oct 2026 10:57:13 -0700 Subject: [PATCH] refactor(redis-store): refresh locks take a user and server target, the store owns the key --- .../redis_distributed_lock.py | 18 +++--- .../redis_refresh_coordinator.py | 28 ++++----- .../redis/stores/mcp_credential_refresh.py | 34 +++++----- .../test_redis_distributed_lock.py | 63 ++++++++++--------- .../test_redis_refresh_coordinator.py | 13 ++-- 5 files changed, 83 insertions(+), 73 deletions(-) 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 06d0248cf2c..c85d4114dfe 100644 --- a/litellm/proxy/mcp_server/outbound_credentials/redis_distributed_lock.py +++ b/litellm/proxy/mcp_server/outbound_credentials/redis_distributed_lock.py @@ -20,7 +20,7 @@ from typing import Final from litellm.proxy.mcp_server.outbound_credentials.redis_refresh_coordinator import ( LockAcquisition, ) -from litellm.redis.stores.mcp_credential_refresh import McpRefreshLockStore +from litellm.redis.stores.mcp_credential_refresh import McpRefreshLockStore, RefreshLockTarget from litellm.redis.transport.contract import RedisFailure @@ -28,17 +28,17 @@ from litellm.redis.transport.contract import RedisFailure class RedisDistributedLock: store: McpRefreshLockStore - async def acquire(self, key: str, token: str, ttl_seconds: float) -> LockAcquisition: - acquired: Final = await self.store.acquire(key, token, ttl=timedelta(seconds=ttl_seconds)) + async def acquire(self, target: RefreshLockTarget, token: str, ttl_seconds: float) -> LockAcquisition: + acquired: Final = await self.store.acquire(target, token, ttl=timedelta(seconds=ttl_seconds)) if isinstance(acquired, RedisFailure): return LockAcquisition.ERROR return LockAcquisition.ACQUIRED if acquired else LockAcquisition.HELD - async def extend(self, key: str, token: str, ttl_seconds: float) -> bool: - return await self.store.extend(key, token, ttl=timedelta(seconds=ttl_seconds)) is True + async def extend(self, target: RefreshLockTarget, token: str, ttl_seconds: float) -> bool: + return await self.store.extend(target, token, ttl=timedelta(seconds=ttl_seconds)) is True - async def release(self, key: str, token: str) -> None: - _ = await self.store.release(key, token) + async def release(self, target: RefreshLockTarget, token: str) -> None: + _ = await self.store.release(target, token) - async def is_held(self, key: str) -> bool: - return await self.store.is_held(key) is True + async def is_held(self, target: RefreshLockTarget) -> bool: + return await self.store.is_held(target) 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 f31f775639d..ac72e58b901 100644 --- a/litellm/proxy/mcp_server/outbound_credentials/redis_refresh_coordinator.py +++ b/litellm/proxy/mcp_server/outbound_credentials/redis_refresh_coordinator.py @@ -31,7 +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 +from litellm.redis.stores.mcp_credential_refresh import RefreshLockTarget class LockAcquisition(Enum): @@ -51,13 +51,13 @@ class DistributedLock(Protocol): after PX-expiry); ``extend`` refreshes the ``PX`` lease only for the owner; ``is_held`` is ``EXISTS`` (so a waiter can poll without taking the lock).""" - async def acquire(self, key: str, token: str, ttl_seconds: float) -> LockAcquisition: ... + async def acquire(self, target: RefreshLockTarget, token: str, ttl_seconds: float) -> LockAcquisition: ... - async def extend(self, key: str, token: str, ttl_seconds: float) -> bool: ... + async def extend(self, target: RefreshLockTarget, token: str, ttl_seconds: float) -> bool: ... - async def release(self, key: str, token: str) -> None: ... + async def release(self, target: RefreshLockTarget, token: str) -> None: ... - async def is_held(self, key: str) -> bool: ... + async def is_held(self, target: RefreshLockTarget) -> bool: ... @dataclass(frozen=True, slots=True) @@ -86,11 +86,11 @@ class RedisRefreshCoordinator: refresh: Callable[[], Awaitable[OAuthToken | None]], reread: Callable[[], Awaitable[OAuthToken | None]], ) -> OAuthToken | None: - key: Final = refresh_lock_key(user_id, server_id) + target: Final = RefreshLockTarget(user_id, server_id) token: Final = self.new_token() - match await self.lock.acquire(key, token, self.lock_ttl_seconds): + match await self.lock.acquire(target, token, self.lock_ttl_seconds): case LockAcquisition.ACQUIRED: - return await self._refresh_with_lease_renewal(key, token, refresh) + return await self._refresh_with_lease_renewal(target, token, refresh) case LockAcquisition.ERROR: # No election happened (lock backend down), so waiting would just re-read the # still-expired token. Refresh anyway; worst case is an extra refresh, not a stale bearer. @@ -100,34 +100,34 @@ class RedisRefreshCoordinator: # the token it persisted - the winner wrote the fresh token to the store, so a plain # re-read sees it without us refreshing again. deadline: Final = self.clock() + self.wait_timeout_seconds - while self.clock() < deadline and await self.lock.is_held(key): + while self.clock() < deadline and await self.lock.is_held(target): await self.sleep(self.poll_interval_seconds) return await reread() async def _refresh_with_lease_renewal( self, - key: str, + target: RefreshLockTarget, token: str, refresh: Callable[[], Awaitable[OAuthToken | None]], ) -> OAuthToken | None: refresh_task: Final = asyncio.ensure_future(refresh()) - renewal_task: Final = asyncio.create_task(self._renew_lease_until_done(key, token, refresh_task)) + renewal_task: Final = asyncio.create_task(self._renew_lease_until_done(target, token, refresh_task)) try: return await refresh_task finally: renewal_task.cancel() with suppress(asyncio.CancelledError): await renewal_task - await self.lock.release(key, token) + await self.lock.release(target, token) async def _renew_lease_until_done( self, - key: str, + target: RefreshLockTarget, token: str, refresh_task: asyncio.Future[OAuthToken | None], ) -> None: budget_deadline: Final = self.clock() + self.refresh_budget_seconds while not refresh_task.done() and self.clock() < budget_deadline: await self.sleep(self.lock_ttl_seconds / 2) - if not refresh_task.done() and not await self.lock.extend(key, token, self.lock_ttl_seconds): + if not refresh_task.done() and not await self.lock.extend(target, token, self.lock_ttl_seconds): return diff --git a/litellm/redis/stores/mcp_credential_refresh.py b/litellm/redis/stores/mcp_credential_refresh.py index e968c3634e6..2950ab7f5d3 100644 --- a/litellm/redis/stores/mcp_credential_refresh.py +++ b/litellm/redis/stores/mcp_credential_refresh.py @@ -12,35 +12,41 @@ _EXTEND_IF_OWNER: Final = ( ) -def refresh_lock_key(user_id: str, server_id: str) -> str: - return f"mcp:refresh_lock:{user_id}:{server_id}" +@dataclass(frozen=True, slots=True) +class RefreshLockTarget: + user_id: str + server_id: str + + +def _lock_key(target: RefreshLockTarget) -> str: + return f"mcp:refresh_lock:{target.user_id}:{target.server_id}" class McpRefreshLockStore(Protocol): - async def acquire(self, key: str, owner: str, *, ttl: timedelta) -> bool | RedisFailure: ... + async def acquire(self, target: RefreshLockTarget, owner: str, *, ttl: timedelta) -> bool | RedisFailure: ... - async def extend(self, key: str, owner: str, *, ttl: timedelta) -> bool | RedisFailure: ... + async def extend(self, target: RefreshLockTarget, owner: str, *, ttl: timedelta) -> bool | RedisFailure: ... - async def release(self, key: str, owner: str) -> bool | RedisFailure: ... + async def release(self, target: RefreshLockTarget, owner: str) -> bool | RedisFailure: ... - async def is_held(self, key: str) -> bool | RedisFailure: ... + async def is_held(self, target: RefreshLockTarget) -> 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 acquire(self, target: RefreshLockTarget, owner: str, *, ttl: timedelta) -> bool | RedisFailure: + return await self.transport.set(_lock_key(target), 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))) + async def extend(self, target: RefreshLockTarget, owner: str, *, ttl: timedelta) -> bool | RedisFailure: + reply: Final = await self.transport.eval(_EXTEND_IF_OWNER, (_lock_key(target),), (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,)) + async def release(self, target: RefreshLockTarget, owner: str) -> bool | RedisFailure: + reply: Final = await self.transport.eval(_RELEASE_IF_OWNER, (_lock_key(target),), (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) + async def is_held(self, target: RefreshLockTarget) -> bool | RedisFailure: + reply: Final = await self.transport.get(_lock_key(target)) 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 c74d2048b63..239c8d31d02 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 @@ -9,9 +9,12 @@ 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 litellm.redis.stores.mcp_credential_refresh import RefreshLockTarget, TransportMcpRefreshLockStore from tests._support.fake_redis import fake_redis_cache +TARGET = RefreshLockTarget(user_id="u", server_id="s") +KEY = "mcp:refresh_lock:u:s" + @pytest.fixture def server() -> fakeredis.FakeServer: @@ -31,25 +34,25 @@ def _lock(server: fakeredis.FakeServer, namespace: str | None = None) -> RedisDi 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 _lock(server, namespace="ns").acquire(TARGET, "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 + assert await raw.keys("*") == [f"ns:{KEY}".encode()] + assert await raw.get(f"ns:{KEY}") == b"tok-1" + assert 9_000 < await raw.pttl(f"ns:{KEY}") <= 10_000 @pytest.mark.asyncio 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 + assert await lock.acquire(TARGET, "tok-a", 10.0) is LockAcquisition.ACQUIRED + assert await lock.acquire(TARGET, "tok-b", 10.0) is LockAcquisition.HELD @pytest.mark.asyncio 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 + assert await lock.acquire(TARGET, "tok", 10.0) is LockAcquisition.ERROR @pytest.mark.asyncio @@ -57,26 +60,26 @@ 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 + await lock.acquire(TARGET, "owner-B", 10.0) + await lock.release(TARGET, "owner-A") + assert await raw.get(KEY) == b"owner-B" + await lock.release(TARGET, "owner-B") + assert await raw.exists(KEY) == 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) + await raw.set(f"ns:{KEY}", "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 + assert await lock.is_held(TARGET) is True + assert await lock.acquire(TARGET, "owner-new", 10.0) is LockAcquisition.HELD + assert await lock.extend(TARGET, "owner-old", 30.0) is True + assert 29_000 < await raw.pttl(f"ns:{KEY}") <= 30_000 + await lock.release(TARGET, "owner-old") + assert await lock.is_held(TARGET) is False @pytest.mark.asyncio @@ -84,22 +87,22 @@ 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 9_000 < await raw.pttl("k") <= 10_000 + await lock.acquire(TARGET, "owner-B", 1.0) + assert await lock.extend(TARGET, "owner-A", 10.0) is False + assert await raw.pttl(KEY) <= 1_000 + assert await lock.extend(TARGET, "owner-B", 10.0) is True + assert 9_000 < await raw.pttl(KEY) <= 10_000 @pytest.mark.asyncio 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) + await lock.acquire(TARGET, "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 + assert await lock.extend(TARGET, "tok", 10.0) is False + await lock.release(TARGET, "tok") + assert await lock.is_held(TARGET) is False server.connected = True - assert await lock.is_held("k") is True + assert await lock.is_held(TARGET) is True diff --git a/tests/unit/proxy/mcp_server/outbound_credentials/test_redis_refresh_coordinator.py b/tests/unit/proxy/mcp_server/outbound_credentials/test_redis_refresh_coordinator.py index 18bd4a66b65..19e642df0c7 100644 --- a/tests/unit/proxy/mcp_server/outbound_credentials/test_redis_refresh_coordinator.py +++ b/tests/unit/proxy/mcp_server/outbound_credentials/test_redis_refresh_coordinator.py @@ -11,8 +11,9 @@ from litellm.proxy.mcp_server.outbound_credentials.redis_refresh_coordinator imp LockAcquisition, RedisRefreshCoordinator, ) +from litellm.redis.stores.mcp_credential_refresh import RefreshLockTarget -_KEY = "mcp:refresh_lock:u:s" +_TARGET = RefreshLockTarget(user_id="u", server_id="s") class _FakeLock: @@ -71,9 +72,9 @@ async def test_winner_acquires_and_releases_with_the_same_token_and_never_reread coord = RedisRefreshCoordinator(lock, new_token=lambda: "tok") result = await coord.run("u", "s", refresh, reread) assert result is refreshed - assert lock.acquire_calls == [(_KEY, "tok", 10.0)] + assert lock.acquire_calls == [(_TARGET, "tok", 10.0)] # released with the SAME token it acquired with, so it can only delete its own lock - assert lock.released == [(_KEY, "tok")] + assert lock.released == [(_TARGET, "tok")] assert reread_calls == [] @@ -90,7 +91,7 @@ async def test_winner_releases_even_when_refresh_raises(): coord = RedisRefreshCoordinator(lock, new_token=lambda: "tok") with pytest.raises(RuntimeError): await coord.run("u", "s", refresh, reread) - assert lock.released == [(_KEY, "tok")] # the lock is freed even on failure + assert lock.released == [(_TARGET, "tok")] # the lock is freed even on failure @pytest.mark.asyncio @@ -127,8 +128,8 @@ async def test_winner_renews_the_lock_while_refresh_runs(): result = await task assert result is not None and result.access_token == "new" - assert lock.extend_calls == [(_KEY, "tok", 10.0)] - assert lock.released == [(_KEY, "tok")] + assert lock.extend_calls == [(_TARGET, "tok", 10.0)] + assert lock.released == [(_TARGET, "tok")] @pytest.mark.asyncio