feat(mcp): cross-replica single-flight refresh for the v2 per-user OAuth store [2/2]

litellm/proxy/_experimental/mcp_server/outbound_credentials/dual_cache_token_backend.py

@@ -0,0 +1,65 @@
+"""Cross-replica ``TokenCacheBackend``: stores the token in LiteLLM's shared ``DualCache``.
+
+Plugs into the foundation's ``CachedOAuthTokenStore`` via the ``TokenCacheBackend`` seam. The token is
+encrypted + serialized by the injected codec and written under a per-``(user, server)`` key with the
+given TTL, so every worker reads one refresh rather than each re-reading and re-refreshing - matching
+v1's ``MCPPerUserTokenCache`` (same NaCl encryption and key, so a token cached by either is readable by
+the other across the cutover). A missing or undecryptable entry reads as a miss.
+"""
+
+from __future__ import annotations
+
+from typing import Protocol
+
+from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import (
+    OAuthToken,
+)
+from litellm.proxy._experimental.mcp_server.outbound_credentials.token_cache_codec import (
+    OAuthTokenCacheCodec,
+)
+
+
+class AsyncCache(Protocol):
+    """The slice of LiteLLM's ``DualCache`` this backend needs (Redis-backed, shared across workers)."""
+
+    async def async_get_cache(self, key: str) -> object | None: ...
+
+    async def async_set_cache(
+        self, key: str, value: str, ttl: float | None = None
+    ) -> None: ...
+
+    async def async_delete_cache(self, key: str) -> None: ...
+
+
+class DualCacheTokenCacheBackend:
+    def __init__(
+        self,
+        cache: AsyncCache,
+        codec: OAuthTokenCacheCodec,
+        *,
+        key_prefix: str = "mcp:per_user_token:",
+    ) -> None:
+        self._cache = cache
+        self._codec = codec
+        self._key_prefix = key_prefix
+
+    def _key(self, user_id: str, server_id: str) -> str:
+        return f"{self._key_prefix}{user_id}:{server_id}"
+
+    async def get(self, user_id: str, server_id: str) -> OAuthToken | None:
+        blob = await self._cache.async_get_cache(self._key(user_id, server_id))
+        return self._codec.decode(blob) if isinstance(blob, str) else None
+
+    async def set(
+        self, user_id: str, server_id: str, token: OAuthToken, ttl_seconds: float
+    ) -> None:
+        if ttl_seconds <= 0:
+            return
+        await self._cache.async_set_cache(
+            self._key(user_id, server_id),
+            self._codec.encode(token),
+            ttl=ttl_seconds,
+        )
+
+    async def delete(self, user_id: str, server_id: str) -> None:
+        await self._cache.async_delete_cache(self._key(user_id, server_id))

litellm/proxy/_experimental/mcp_server/outbound_credentials/oauth_token_store.py

@@ -289,9 +289,18 @@ async def refresh_latest_token() -> OAuthToken | None:
                 return latest_token
             return await self._refresher.refresh(user_id, server_id, latest_token)
 
+        async def reread_fresh_token() -> OAuthToken | None:
+            # A loser re-reads what the winner persisted. If the winner's refresh failed, the store
+            # still holds the expired token; surface None (-> challenge) like the winner did rather
+            # than the stale bearer the upstream would 401.
+            latest_token = await self._inner.fetch(user_id, server_id)
+            if latest_token is None or self._is_expired(latest_token):
+                return None
+            return latest_token
+
         return await self._coordinator.run(
             user_id,
             server_id,
             refresh=refresh_latest_token,
-            reread=lambda: self._inner.fetch(user_id, server_id),
+            reread=reread_fresh_token,
         )

litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py

@@ -1,11 +1,11 @@
 """Composition root for the v2-native authorization_code per-user OAuth token store (step 1b).
 
 Assembles ``Cached(Refreshing(V2PerUserTokenStore))`` and replaces ``V1PerUserTokenStore`` in the
-resolver. The runtime collaborators (DB, HTTP) are LiteLLM globals not ready at import time, so the
-chain is built lazily on first use. The cache and refresh coordinator use the foundation's in-process
-defaults (correct for a single replica); the cross-replica path is layered on separately. The DB
-read/refresh-grant/persist collaborators acquire their globals per call, mirroring v1's lazy-import
-pattern.
+resolver. The runtime collaborators (DB, HTTP, the shared cache, Redis) are LiteLLM globals not ready
+at import time, so the chain is built lazily on first use. When Redis is wired it uses the
+cross-replica path (DualCache-backed cache + ``SET NX PX`` coordinator); otherwise it falls back to
+the foundation's in-process defaults (correct for a single replica). The DB read/refresh-grant/persist
+collaborators acquire their globals per call, mirroring v1's lazy-import pattern.
 """
 
 from __future__ import annotations
@@ -17,12 +17,27 @@
 from litellm.proxy._experimental.mcp_server.outbound_credentials.authz_code_refresher import (
     AuthorizationCodeRefresher,
 )
+from litellm.proxy._experimental.mcp_server.outbound_credentials.dual_cache_token_backend import (
+    AsyncCache,
+    DualCacheTokenCacheBackend,
+)
 from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import (
     CachedOAuthTokenStore,
     OAuthToken,
+    RefreshCoordinator,
     RefreshingTokenStore,
+    TokenCacheBackend,
     TokenStoreUnavailable,
 )
+from litellm.proxy._experimental.mcp_server.outbound_credentials.redis_distributed_lock import (
+    RedisDistributedLock,
+)
+from litellm.proxy._experimental.mcp_server.outbound_credentials.redis_refresh_coordinator import (
+    RedisRefreshCoordinator,
+)
+from litellm.proxy._experimental.mcp_server.outbound_credentials.token_cache_codec import (
+    OAuthTokenCacheCodec,
+)
 from litellm.proxy._experimental.mcp_server.outbound_credentials.v2_token_store import (
     V2PerUserTokenStore,
 )
@@ -99,16 +114,53 @@ async def _post_token_endpoint(
         return body  # pyright: ignore
 
 
+def _runtime_backend_and_coordinator() -> tuple[
+    TokenCacheBackend | None, RefreshCoordinator | None
+]:
+    """The cross-replica cache + coordinator when Redis is wired, else ``(None, None)`` so the
+    foundation's in-process defaults are used (a single replica needs no shared cache or lock).
+    """
+    from litellm.proxy.common_utils.encrypt_decrypt_utils import (  # noqa: PLC0415
+        decrypt_value_helper,
+        encrypt_value_helper,
+    )
+    from litellm.proxy.proxy_server import user_api_key_cache  # noqa: PLC0415
+
+    redis_cache = user_api_key_cache.redis_cache
+    if redis_cache is None:
+        return None, None
+    codec = OAuthTokenCacheCodec(
+        encrypt_value_helper,
+        lambda blob: decrypt_value_helper(
+            blob, "mcp_per_user_token", exception_type="debug"
+        ),
+    )
+    # user_api_key_cache satisfies the AsyncCache slice (DualCache types ttl via **kwargs) and the
+    # Redis client from init_async_client() is partially typed - both are untyped-boundary casts.
+    cache: AsyncCache = user_api_key_cache  # pyright: ignore
+    redis_client = redis_cache.init_async_client()  # pyright: ignore
+    lock = RedisDistributedLock(
+        redis_client,  # pyright: ignore
+        namespace_key=redis_cache.check_and_fix_namespace,
+    )
+    backend = DualCacheTokenCacheBackend(cache, codec)
+    coordinator = RedisRefreshCoordinator(lock)
+    return backend, coordinator
+
+
 def build_per_user_oauth_token_store(
     server_lookup: ServerLookup,
 ) -> CachedOAuthTokenStore:
+    backend, coordinator = _runtime_backend_and_coordinator()
     refresher = AuthorizationCodeRefresher(
         server_lookup, _post_token_endpoint, _persist_credential
     )
-    # Cache and refresh coordinator use the foundation's in-process defaults (a single replica needs
-    # no shared cache or lock); the cross-replica path is layered on separately.
-    refreshing = RefreshingTokenStore(V2PerUserTokenStore(_read_credential), refresher)
-    return CachedOAuthTokenStore(refreshing, default_ttl_seconds=_DEFAULT_TTL_SECONDS)
+    refreshing = RefreshingTokenStore(
+        V2PerUserTokenStore(_read_credential), refresher, coordinator=coordinator
+    )
+    return CachedOAuthTokenStore(
+        refreshing, default_ttl_seconds=_DEFAULT_TTL_SECONDS, backend=backend
+    )
 
 
 class LazyPerUserOAuthTokenStore:

litellm/proxy/_experimental/mcp_server/outbound_credentials/redis_distributed_lock.py

@@ -0,0 +1,89 @@
+"""Concrete ``DistributedLock`` over a Redis client: ``SET NX PX`` / compare-and-delete / ``EXISTS``.
+
+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). ``release`` deletes the key only when it still holds this caller's token
+(a compare-and-delete Lua script), 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 a deployment's lock keys carry the
+same namespace its cache keys do and cannot collide with another deployment sharing the Redis backend.
+
+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.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Protocol
+
+from litellm._logging import verbose_logger
+from litellm.proxy._experimental.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"
+)
+
+
+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: ...
+
+
+class RedisDistributedLock:
+    def __init__(
+        self,
+        client: RedisCommands,
+        *,
+        namespace_key: Callable[[str], str] = lambda key: key,
+    ) -> None:
+        self._client = client
+        self._namespace_key = namespace_key
+
+    async def acquire(
+        self, key: str, token: str, ttl_seconds: float
+    ) -> LockAcquisition:
+        try:
+            result = 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)
+            return LockAcquisition.ERROR
+        return LockAcquisition.ACQUIRED if result is not None else LockAcquisition.HELD
+
+    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)
+
+    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

litellm/proxy/_experimental/mcp_server/outbound_credentials/redis_refresh_coordinator.py

@@ -0,0 +1,109 @@
+"""Cross-replica ``RefreshCoordinator``: one refresh per ``(user, server)`` across all workers.
+
+Plugs into the foundation's ``RefreshingTokenStore`` via the ``RefreshCoordinator`` seam. A ``SET NX
+PX`` lock elects one worker to run the refresh while the rest wait for it and re-read the token it
+persisted - so a rotating refresh_token is used once across the fleet, not once per worker. The lock
+auto-expires (``PX``), so a crashed holder can't wedge refresh; a loser that times out (or whose
+holder crashed mid-refresh) falls back to a re-read, and the surrounding store re-checks expiry on the
+next fetch, so a crash self-heals rather than serving stale forever. Reading needs no lock, so losers
+don't serialize behind each other. The lock is injected (a thin Redis ``SET NX``/``DEL``/``EXISTS``
+wrapper in production, a fake in tests).
+
+The lock is a single-flight optimization, not a correctness mutex, so it fails open: when the lock
+backend is unreachable, ``acquire`` reports ``ERROR`` (distinct from ``HELD``) and this coordinator
+refreshes anyway rather than wait on a holder that may not exist and then serve a still-expired token.
+That degrades a Redis outage to the no-coordinator behavior (each worker may refresh), never a stale
+bearer the upstream would 401.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import time
+import uuid
+from collections.abc import Awaitable, Callable
+from enum import Enum
+from typing import Protocol
+
+from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import (
+    OAuthToken,
+)
+
+
+class LockAcquisition(Enum):
+    """Outcome of a best-effort ``acquire``. ``ERROR`` is kept distinct from ``HELD`` so a caller can
+    tell "someone else is refreshing" (wait and re-read) from "the lock backend is down" (no election
+    happened, so refresh anyway) instead of conflating both into a single ``False``."""
+
+    ACQUIRED = "acquired"  # won the election; this worker refreshes
+    HELD = "held"  # another worker holds it; wait then re-read
+    ERROR = "error"  # lock backend unreachable; holder unknown, so refresh anyway
+
+
+class DistributedLock(Protocol):
+    """A best-effort cross-replica lock. ``acquire`` is ``SET key token NX PX ttl`` reported as a
+    ``LockAcquisition`` (won / held by another / backend error); ``release`` deletes the key only if
+    it still holds this caller's ``token`` (so it cannot delete a lock another worker re-acquired
+    after PX-expiry); ``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 release(self, key: str, token: str) -> None: ...
+
+    async def is_held(self, key: str) -> bool: ...
+
+
+class RedisRefreshCoordinator:
+    def __init__(
+        self,
+        lock: DistributedLock,
+        *,
+        key_prefix: str = "mcp:refresh_lock:",
+        lock_ttl_seconds: float = 10.0,
+        wait_timeout_seconds: float = 10.0,
+        poll_interval_seconds: float = 0.05,
+        sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
+        clock: Callable[[], float] = time.monotonic,
+        new_token: Callable[[], str] = lambda: uuid.uuid4().hex,
+    ) -> None:
+        self._lock = lock
+        self._key_prefix = key_prefix
+        self._lock_ttl_seconds = lock_ttl_seconds
+        self._wait_timeout_seconds = wait_timeout_seconds
+        self._poll_interval_seconds = poll_interval_seconds
+        self._sleep = sleep
+        self._clock = clock
+        self._new_token = new_token
+
+    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,
+        server_id: str,
+        refresh: Callable[[], Awaitable[OAuthToken | None]],
+        reread: Callable[[], Awaitable[OAuthToken | None]],
+    ) -> OAuthToken | None:
+        key = self._key(user_id, server_id)
+        token = self._new_token()
+        match await self._lock.acquire(key, token, self._lock_ttl_seconds):
+            case LockAcquisition.ACQUIRED:
+                try:
+                    return await refresh()
+                finally:
+                    await self._lock.release(key, token)
+            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.
+                return await refresh()
+            case LockAcquisition.HELD:
+                # Another worker holds the lock; wait for it to finish (release or PX-expiry), then read
+                # 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 = self._clock() + self._wait_timeout_seconds
+                while self._clock() < deadline and await self._lock.is_held(key):
+                    await self._sleep(self._poll_interval_seconds)
+                return await reread()