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,74 @@
+"""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 dataclasses import KW_ONLY, dataclass
+from typing import Protocol
+
+from litellm._logging import verbose_logger
+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: ...
+
+
+@dataclass(frozen=True, slots=True)
+class DualCacheTokenCacheBackend:
+    """Every method degrades a cache or codec failure to its safe value - ``get`` to a miss
+    (``None``), ``set``/``delete`` to a no-op - so a Redis outage or an undecryptable entry reads as a
+    cache miss rather than a request error, matching v1 and this layer's "boundary failure = miss"
+    contract. The guarantee holds here regardless of whether the injected cache/codec also swallow.
+    """
+
+    cache: AsyncCache
+    codec: OAuthTokenCacheCodec
+    _: KW_ONLY
+    key_prefix: str = "mcp:per_user_token:"
+
+    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:
+        try:
+            blob = await self.cache.async_get_cache(self._key(user_id, server_id))
+            return self.codec.decode(blob) if isinstance(blob, str) else None
+        except Exception as exc:  # noqa: BLE001
+            verbose_logger.debug("MCP per-user token cache get failed (miss): %s", exc)
+            return None
+
+    async def set(self, user_id: str, server_id: str, token: OAuthToken, ttl_seconds: float) -> None:
+        if ttl_seconds <= 0:
+            return
+        try:
+            await self.cache.async_set_cache(
+                self._key(user_id, server_id),
+                self.codec.encode(token),
+                ttl=ttl_seconds,
+            )
+        except Exception as exc:  # noqa: BLE001
+            verbose_logger.debug("MCP per-user token cache set failed (ignored): %s", exc)
+
+    async def delete(self, user_id: str, server_id: str) -> None:
+        try:
+            await self.cache.async_delete_cache(self._key(user_id, server_id))
+        except Exception as exc:  # noqa: BLE001
+            verbose_logger.debug("MCP per-user token cache delete failed (ignored): %s", exc)

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

@@ -272,9 +272,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,28 +1,45 @@
 """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
 
+import asyncio
 from collections.abc import Callable
 from typing import TYPE_CHECKING
 
 from litellm._logging import verbose_logger
 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,
+    OAuthTokenStore,
+    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,
 )
@@ -34,6 +51,7 @@
 _DEFAULT_TTL_SECONDS = 300.0
 
 ServerLookup = Callable[[str], "MCPServer | None"]
+StoreBuilder = Callable[[ServerLookup], tuple[OAuthTokenStore, bool]]
 
 
 async def _read_credential(user_id: str, server_id: str) -> dict[str, object] | None:
@@ -97,29 +115,111 @@ async def _post_token_endpoint(url: str, form: dict[str, str]) -> dict[str, obje
         return body  # pyright: ignore
 
 
+def _redis_cache_is_available() -> bool:
+    from litellm.proxy.proxy_server import user_api_key_cache  # noqa: PLC0415
+
+    return user_api_key_cache.redis_cache is not None
+
+
+def _runtime_backend_and_coordinator() -> tuple[TokenCacheBackend | None, RefreshCoordinator | None, bool]:
+    """The cross-replica cache + coordinator when Redis is wired, else ``(None, None, False)`` 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, False
+    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, True
+
+
+def _build_per_user_oauth_token_store(
+    server_lookup: ServerLookup,
+) -> tuple[CachedOAuthTokenStore, bool]:
+    backend, coordinator, uses_redis = _runtime_backend_and_coordinator()
+    refresher = AuthorizationCodeRefresher(server_lookup, _post_token_endpoint, _persist_credential)
+    refreshing = RefreshingTokenStore(V2PerUserTokenStore(_read_credential), refresher, coordinator=coordinator)
+    return CachedOAuthTokenStore(refreshing, default_ttl_seconds=_DEFAULT_TTL_SECONDS, backend=backend), uses_redis
+
+
 def build_per_user_oauth_token_store(
     server_lookup: ServerLookup,
 ) -> CachedOAuthTokenStore:
-    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)
+    store, _uses_redis = _build_per_user_oauth_token_store(server_lookup)
+    return store
 
 
 class LazyPerUserOAuthTokenStore:
     """``OAuthTokenStore`` that builds the v2-native chain on first ``fetch``.
 
     The chain's cache/lock collaborators are LiteLLM runtime globals not available when the resolver
     is constructed at import time, so construction is deferred to the first request (by when they are
-    wired). Built once, then reused.
+    wired). A no-Redis chain is replaced once Redis becomes available.
     """
 
-    def __init__(self, server_lookup: ServerLookup) -> None:
+    def __init__(
+        self,
+        server_lookup: ServerLookup,
+        *,
+        store_builder: StoreBuilder = _build_per_user_oauth_token_store,
+        redis_available: Callable[[], bool] = _redis_cache_is_available,
+    ) -> None:
         self._server_lookup = server_lookup
-        self._store: CachedOAuthTokenStore | None = None
+        self._store_builder = store_builder
+        self._redis_available = redis_available
+        self._store: OAuthTokenStore | None = None
+        self._uses_redis = False
+        self._fetch_lock = asyncio.Condition()
+        self._local_fetches = 0
 
     async def fetch(self, user_id: str, server_id: str) -> OAuthToken | None:
-        if self._store is None:
-            self._store = build_per_user_oauth_token_store(self._server_lookup)
-        return await self._store.fetch(user_id, server_id)
+        if self._uses_redis:
+            store = self._store
+            if store is not None:
+                return await store.fetch(user_id, server_id)
+
+        store, uses_redis = await self._store_for_fetch()
+        try:
+            return await store.fetch(user_id, server_id)
+        finally:
+            if not uses_redis:
+                await self._finish_local_fetch()
+
+    async def _store_for_fetch(self) -> tuple[OAuthTokenStore, bool]:
+        async with self._fetch_lock:
+            while (
+                self._store is not None and not self._uses_redis and self._redis_available() and self._local_fetches > 0
+            ):
+                await self._fetch_lock.wait()
+            store = self._store
+            if store is None or (not self._uses_redis and self._redis_available()):
+                store, self._uses_redis = self._store_builder(self._server_lookup)
+                self._store = store
+            uses_redis = self._uses_redis
+            if not uses_redis:
+                self._local_fetches += 1
+            return store, uses_redis
+
+    async def _finish_local_fetch(self) -> None:
+        async with self._fetch_lock:
+            self._local_fetches -= 1
+            if self._local_fetches == 0:
+                self._fetch_lock.notify_all()

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

@@ -0,0 +1,94 @@
+"""Concrete ``DistributedLock`` over a Redis client: ``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.
+
+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 dataclasses import KW_ONLY, dataclass
+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"
+_EXTEND_IF_OWNER = (
+    "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: ...
+
+
+@dataclass(frozen=True, slots=True)
+class RedisDistributedLock:
+    client: RedisCommands
+    _: KW_ONLY
+    namespace_key: Callable[[str], str] = lambda key: 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 extend(self, key: str, token: str, ttl_seconds: float) -> bool:
+        try:
+            result = 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
+
+    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