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

@@ -284,9 +284,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,48 @@ 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"),
+    )
+    # 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
+    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,66 @@
+"""Concrete ``DistributedLock`` over a Redis client: ``SET NX PX`` / ``DEL`` / ``EXISTS``.
+
+The cross-replica lock the ``RedisRefreshCoordinator`` elects refreshers with. ``acquire`` is an
+atomic ``SET key NX PX ttl`` (only the first caller wins; the entry self-expires so a crashed holder
+can't wedge refresh), ``release`` is ``DEL``, ``is_held`` is ``EXISTS``. 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 typing import Protocol
+
+from litellm._logging import verbose_logger
+from litellm.proxy._experimental.mcp_server.outbound_credentials.redis_refresh_coordinator import (
+    LockAcquisition,
+)
+
+
+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 delete(self, *names: str) -> int: ...
+
+    async def exists(self, *names: str) -> int: ...
+
+
+class RedisDistributedLock:
+    def __init__(self, client: RedisCommands) -> None:
+        self._client = client
+
+    async def acquire(self, key: str, ttl_seconds: float) -> LockAcquisition:
+        try:
+            result = await self._client.set(
+                key, "1", 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) -> None:
+        try:
+            await self._client.delete(key)
+        # 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(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,102 @@
+"""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
+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 NX PX ttl`` reported as a
+    ``LockAcquisition`` (won / held by another / backend error); ``release`` is ``DEL``; ``is_held``
+    is ``EXISTS`` (so a waiter can poll without taking the lock)."""
+
+    async def acquire(self, key: str, ttl_seconds: float) -> LockAcquisition: ...
+
+    async def release(self, key: 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,
+    ) -> 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
+
+    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)
+        match await self._lock.acquire(key, self._lock_ttl_seconds):
+            case LockAcquisition.ACQUIRED:
+                try:
+                    return await refresh()
+                finally:
+                    await self._lock.release(key)
+            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()

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

@@ -0,0 +1,37 @@
+"""Serialize + encrypt boundary for caching an OAuth token in a shared (Redis) cache.
+
+A cross-replica cache must serialize the token, and a plaintext bearer in Redis is a leak, so this
+encrypts the value (NaCl in production via the injected ``encrypt``, identity in tests). It caches
+**only** the ``access_token``: the hot path needs just the bearer, expiry is carried by the cache
+entry's TTL (set from the token's ``expires_at`` by the cache), and the long-lived refresh_token stays
+in the DB - the refresh path is always a cache miss that re-reads it - so it never reaches Redis. A
+decoded token therefore carries only the bearer (``expires_at`` and ``refresh_token`` both None); the
+TTL, not the value, bounds its life. An empty/undecryptable blob (e.g. master-key rotation) is a miss.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import (
+    OAuthToken,
+)
+
+
+class OAuthTokenCacheCodec:
+    def __init__(
+        self,
+        encrypt: Callable[[str], str],
+        decrypt: Callable[[str], str | None],
+    ) -> None:
+        self._encrypt = encrypt
+        self._decrypt = decrypt
+
+    def encode(self, token: OAuthToken) -> str:
+        return self._encrypt(token.access_token)
+
+    def decode(self, blob: str) -> OAuthToken | None:
+        access_token = self._decrypt(blob)
+        if not access_token:
+            return None
+        return OAuthToken(access_token=access_token, refresh_token=None)