feat(mcp): proactive token refresh with self-cleaning single-flight

Add TokenRefresher (a mode-supplied seam: mint a fresh token from an expired one and persist it)
and RefreshingTokenStore: when the stored token is near expiry, the first caller refreshes while
concurrent callers await the same in-flight task and share its result, so the IdP is not
stampeded. The task self-cleans (a done-callback drops its entry), so the map is bounded by
in-flight refreshes rather than by distinct users/servers, and is detached from the caller so a
cancelled caller does not abort the refresh. An expired token the refresher cannot renew surfaces
as None so the arm challenges, never a stale bearer; it composes under CachedOAuthTokenStore.
OAuthToken's repr masks the access/refresh tokens so a stray log cannot leak them. Cross-replica
single-flight (Redis) and reactive-401 refresh are the later distributed hardening.

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

@@ -4,29 +4,40 @@ The resolver reads a user's token through the injected ``OAuthTokenStore`` seam;
 ``CachedOAuthTokenStore`` is an expiry-aware cache in front of it. ``TokenStoreUnavailable``
 signals an unreachable backing store, so an outage is never cached or read as "not authorized".
 
-Refresh (using ``refresh_token`` once the access token has expired) and distributed single-flight
-are the later hardening; this cache only avoids serving a token past its own expiry.
+``RefreshingTokenStore`` mints a fresh token through an injected ``TokenRefresher`` when the stored
+one is near expiry, under in-process per-(user, server) single-flight so concurrent callers share
+one refresh. Distributed (cross-replica) single-flight and reactive-401 refresh are the later
+hardening. The mode plugs in its own source and refresher; the cache, store seam, and refresh
+machinery are shared across the oauth2 modes (authorization_code / client_credentials /
+token_exchange).
 """
 
 from __future__ import annotations
 
+import asyncio
 import time
 from dataclasses import dataclass
 from typing import Callable, Dict, Optional, Protocol, Tuple
 
 
-@dataclass(frozen=True, slots=True)
+@dataclass(frozen=True, slots=True, repr=False)
 class OAuthToken:
     """A user's OAuth credential: the bearer value, when it expires, and how to refresh it.
 
     ``expires_at`` is epoch seconds (``None`` means no known expiry). ``refresh_token`` is kept for
-    the later refresh step; it is never minted into a header directly.
+    the later refresh step; it is never minted into a header directly. ``repr`` masks both secrets
+    so a stray log line cannot leak them (the values are still plain ``str`` for the header path,
+    since ``SecretStr`` resolves as unknown under this repo's basedpyright).
     """
 
     access_token: str
     expires_at: Optional[float] = None
     refresh_token: Optional[str] = None
 
+    def __repr__(self) -> str:
+        has_refresh = self.refresh_token is not None
+        return f"OAuthToken(access_token=***, expires_at={self.expires_at!r}, has_refresh_token={has_refresh})"
+
 
 class TokenStoreUnavailable(Exception):
     """Raised by ``fetch`` when the backing token store is unreachable (e.g. the DB is down).
@@ -50,6 +61,18 @@ class OAuthTokenStore(Protocol):
     async def fetch(self, user_id: str, server_id: str) -> Optional[OAuthToken]: ...
 
 
+class TokenRefresher(Protocol):
+    """Mints a fresh token from an expired one and persists it, returning the new token.
+
+    The action is mode-specific: the ``authorization_code`` refresh_token grant, the
+    ``client_credentials`` grant, or an RFC 8693 re-exchange. Returns ``None`` when it cannot
+    refresh (e.g. no ``refresh_token``), which the caller turns into a 401 challenge. It must
+    persist the new token so later requests (and the surrounding cache) read it without refreshing.
+    """
+
+    async def refresh(self, token: OAuthToken) -> Optional[OAuthToken]: ...
+
+
 class CachedOAuthTokenStore:
     """Expiry-aware cache over an ``OAuthTokenStore``.
 
@@ -99,3 +122,63 @@ class CachedOAuthTokenStore:
         """Drop a cached entry after the user (re)authorizes or revokes, so a stale token or a
         stale "not authorized" None cannot mask the change."""
         self._cache.pop((user_id, server_id), None)
+
+
+class RefreshingTokenStore:
+    """An ``OAuthTokenStore`` that proactively refreshes a near-expiry token.
+
+    Reads from an inner store; if the token is within ``expiry_skew_seconds`` of expiry, it mints a
+    fresh one via the injected ``TokenRefresher`` under per-(user, server) single-flight: the first
+    caller refreshes while concurrent callers await the same in-flight future and share its result,
+    instead of stampeding the IdP. The refresher persists the new token so later requests (and the
+    surrounding cache) read it without refreshing again. An expired token the refresher cannot renew
+    (``None``) is surfaced as ``None`` so the arm challenges, never a stale bearer.
+
+    Single-flight here is in-process (one event loop). Cross-replica single-flight (Redis SET NX)
+    and reactive-401 refresh are the later distributed hardening. Composes under
+    ``CachedOAuthTokenStore`` so the refreshed token is cached until its own expiry.
+    """
+
+    def __init__(
+        self,
+        inner: OAuthTokenStore,
+        refresher: TokenRefresher,
+        *,
+        expiry_skew_seconds: float = 30.0,
+        clock: Callable[[], float] = time.time,
+    ) -> None:
+        self._inner = inner
+        self._refresher = refresher
+        self._expiry_skew_seconds = expiry_skew_seconds
+        self._clock = clock
+        # In-flight refreshes, one future per (user, server). Entries exist only while a refresh
+        # is running (removed in `finally`), so the map is bounded by concurrency, not by the
+        # number of distinct users/servers ever seen.
+        self._inflight: Dict[Tuple[str, str], asyncio.Future[Optional[OAuthToken]]] = {}
+
+    def _is_expired(self, token: OAuthToken) -> bool:
+        return (
+            token.expires_at is not None
+            and self._clock() >= token.expires_at - self._expiry_skew_seconds
+        )
+
+    async def fetch(self, user_id: str, server_id: str) -> Optional[OAuthToken]:
+        token = await self._inner.fetch(user_id, server_id)
+        if token is None or not self._is_expired(token):
+            return token
+        return await self._refresh_single_flight(user_id, server_id, token)
+
+    async def _refresh_single_flight(
+        self, user_id: str, server_id: str, token: OAuthToken
+    ) -> Optional[OAuthToken]:
+        key = (user_id, server_id)
+        task = self._inflight.get(key)
+        if task is None:
+            # First caller starts the refresh; concurrent callers await the same task and share its
+            # result (or exception). The done-callback removes the entry, so the map self-cleans and
+            # is bounded by in-flight refreshes, not by the number of distinct users/servers. The
+            # task is detached from the caller, so a cancelled caller does not abort the refresh.
+            task = asyncio.ensure_future(self._refresher.refresh(token))
+            self._inflight[key] = task
+            task.add_done_callback(lambda _t, k=key: self._inflight.pop(k, None))
+        return await task

tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_oauth_token_store.py

@@ -1,5 +1,6 @@
-"""Tests for the v2 OAuth token cache (CachedOAuthTokenStore)."""
+"""Tests for the v2 OAuth token cache and refresh (CachedOAuthTokenStore, RefreshingTokenStore)."""
 
+import asyncio
 from typing import Dict, List, Optional, Tuple
 
 import pytest
@@ -7,6 +8,7 @@ import pytest
 from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import (
     CachedOAuthTokenStore,
     OAuthToken,
+    RefreshingTokenStore,
     TokenStoreUnavailable,
 )
 
@@ -121,3 +123,106 @@ async def test_isolates_by_subject():
     second = await store.fetch("u2", "s")
     assert first is not None and first.access_token == "a"
     assert second is not None and second.access_token == "b"
+
+
+class _RefreshablePair:
+    """A store + refresher pair that simulates persistence: refresh() updates what fetch returns,
+    and yields once so concurrent callers actually contend on the single-flight lock."""
+
+    def __init__(self, initial: Optional[OAuthToken]) -> None:
+        self._current = initial
+        self.fetch_calls = 0
+        self.refresh_calls = 0
+
+    async def fetch(self, user_id: str, server_id: str) -> Optional[OAuthToken]:
+        self.fetch_calls += 1
+        return self._current
+
+    async def refresh(self, token: OAuthToken) -> Optional[OAuthToken]:
+        self.refresh_calls += 1
+        await asyncio.sleep(
+            0
+        )  # yield so other concurrent callers reach the lock and wait
+        self._current = OAuthToken(access_token="refreshed", expires_at=9999.0)
+        return self._current
+
+
+async def test_refreshing_passes_through_a_fresh_token():
+    pair = _RefreshablePair(OAuthToken(access_token="ok", expires_at=9999.0))
+    store = RefreshingTokenStore(
+        pair, pair, expiry_skew_seconds=30, clock=_Clock(1000.0)
+    )
+
+    token = await store.fetch("u", "s")
+    assert token is not None and token.access_token == "ok"
+    assert pair.refresh_calls == 0  # not near expiry -> no refresh
+
+
+async def test_refreshing_mints_a_fresh_token_when_expired():
+    pair = _RefreshablePair(OAuthToken(access_token="old", expires_at=900.0))
+    store = RefreshingTokenStore(
+        pair, pair, expiry_skew_seconds=30, clock=_Clock(1000.0)
+    )
+
+    token = await store.fetch("u", "s")
+    assert token is not None and token.access_token == "refreshed"
+    assert pair.refresh_calls == 1
+
+
+async def test_refreshing_returns_none_when_it_cannot_refresh():
+    class _NoRefresh:
+        async def fetch(self, user_id: str, server_id: str) -> Optional[OAuthToken]:
+            return OAuthToken(access_token="old", expires_at=900.0)
+
+        async def refresh(self, token: OAuthToken) -> Optional[OAuthToken]:
+            return None  # e.g. no refresh_token
+
+    src = _NoRefresh()
+    store = RefreshingTokenStore(src, src, expiry_skew_seconds=30, clock=_Clock(1000.0))
+    # expired and unrefreshable -> None (the arm challenges), never a stale bearer
+    assert await store.fetch("u", "s") is None
+
+
+async def test_refreshing_is_single_flight_under_concurrency():
+    pair = _RefreshablePair(OAuthToken(access_token="old", expires_at=900.0))
+    store = RefreshingTokenStore(
+        pair, pair, expiry_skew_seconds=30, clock=_Clock(1000.0)
+    )
+
+    results = await asyncio.gather(*[store.fetch("u", "s") for _ in range(5)])
+    assert pair.refresh_calls == 1  # one refresh shared across 5 concurrent callers
+    assert all(r is not None and r.access_token == "refreshed" for r in results)
+
+
+async def test_refresh_failure_is_shared_by_joiners_not_re_run():
+    class _FailingRefresher:
+        def __init__(self) -> None:
+            self.calls = 0
+
+        async def fetch(self, user_id: str, server_id: str) -> Optional[OAuthToken]:
+            return OAuthToken(access_token="old", expires_at=900.0)
+
+        async def refresh(self, token: OAuthToken) -> Optional[OAuthToken]:
+            self.calls += 1
+            await asyncio.sleep(0)  # let the concurrent callers join the same task
+            raise RuntimeError("refresh boom")
+
+    src = _FailingRefresher()
+    store = RefreshingTokenStore(src, src, expiry_skew_seconds=30, clock=_Clock(1000.0))
+
+    results = await asyncio.gather(
+        *[store.fetch("u", "s") for _ in range(3)], return_exceptions=True
+    )
+    assert src.calls == 1  # single-flight: one attempt, the failure is shared
+    assert all(isinstance(r, RuntimeError) for r in results)
+
+
+def test_oauth_token_repr_masks_the_secrets():
+    token = OAuthToken(
+        access_token="super-secret", expires_at=123.0, refresh_token="rt-secret"
+    )
+    rendered = repr(token)
+    assert "super-secret" not in rendered
+    assert "rt-secret" not in rendered
+    assert "access_token=***" in rendered
+    assert "has_refresh_token=True" in rendered