diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/oauth_token_store.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/oauth_token_store.py index 7cdeccaebc4..f1e4b17da29 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/oauth_token_store.py +++ b/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 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_oauth_token_store.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_oauth_token_store.py index 5c5017966cc..92f8a8fa1fa 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_oauth_token_store.py +++ b/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