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 new file mode 100644 index 00000000000..7cdeccaebc4 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/oauth_token_store.py @@ -0,0 +1,101 @@ +"""Per-user OAuth token store for the ``authorization_code`` mode. + +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. +""" + +from __future__ import annotations + +import time +from dataclasses import dataclass +from typing import Callable, Dict, Optional, Protocol, Tuple + + +@dataclass(frozen=True, slots=True) +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. + """ + + access_token: str + expires_at: Optional[float] = None + refresh_token: Optional[str] = None + + +class TokenStoreUnavailable(Exception): + """Raised by ``fetch`` when the backing token store is unreachable (e.g. the DB is down). + + Distinct from returning ``None`` for "the user has not authorized this server": a read-through + cache skips caching the failure, and the resolver maps it to its fail-closed status rather than + treating an outage as a definite absence. + """ + + +class OAuthTokenStore(Protocol): + """Per-user OAuth token lookup for the ``authorization_code`` mode. + + Returns the user's token for an upstream, or ``None`` when they have not completed the OAuth + flow (the arm turns that into a 401 challenge). The ``(user_id, server_id)`` pair fully scopes + the lookup, so an implementation must never return one subject's token to another. Raises + ``TokenStoreUnavailable`` when the backing store is unreachable, so an outage is never cached or + read as a definite absence. + """ + + async def fetch(self, user_id: str, server_id: str) -> Optional[OAuthToken]: ... + + +class CachedOAuthTokenStore: + """Expiry-aware cache over an ``OAuthTokenStore``. + + A cached token is served only while it is unexpired (minus ``expiry_skew_seconds``); past that + the inner store is read again. Tokens with no known expiry, and the ``None`` "not authorized" + result, are held for ``default_ttl_seconds`` so the store is not hit on every call. The clock is + injected (wall-clock, since ``expires_at`` is epoch) so expiry is deterministic in tests, and a + store outage (``TokenStoreUnavailable``) propagates without being cached. + """ + + def __init__( + self, + inner: OAuthTokenStore, + *, + default_ttl_seconds: float, + expiry_skew_seconds: float = 30.0, + max_size: int = 4096, + clock: Callable[[], float] = time.time, + ) -> None: + self._inner = inner + self._default_ttl_seconds = default_ttl_seconds + self._expiry_skew_seconds = expiry_skew_seconds + self._max_size = max_size + self._clock = clock + self._cache: Dict[Tuple[str, str], Tuple[Optional[OAuthToken], float]] = {} + + def _valid_until(self, token: Optional[OAuthToken]) -> float: + if token is not None and token.expires_at is not None: + return token.expires_at - self._expiry_skew_seconds + return self._clock() + self._default_ttl_seconds + + async def fetch(self, user_id: str, server_id: str) -> Optional[OAuthToken]: + key = (user_id, server_id) + hit = self._cache.get(key) + if hit is not None: + token, valid_until = hit + if self._clock() < valid_until: + return token + + token = await self._inner.fetch(user_id, server_id) + if len(self._cache) >= self._max_size: + self._cache.clear() + self._cache[key] = (token, self._valid_until(token)) + return token + + def invalidate(self, user_id: str, server_id: str) -> None: + """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) 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 new file mode 100644 index 00000000000..5c5017966cc --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_oauth_token_store.py @@ -0,0 +1,123 @@ +"""Tests for the v2 OAuth token cache (CachedOAuthTokenStore).""" + +from typing import Dict, List, Optional, Tuple + +import pytest + +from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import ( + CachedOAuthTokenStore, + OAuthToken, + TokenStoreUnavailable, +) + + +class _FakeStore: + """An OAuthTokenStore that records calls and returns canned tokens.""" + + def __init__(self, values: Dict[Tuple[str, str], Optional[OAuthToken]]) -> None: + self._values = values + self.calls: List[Tuple[str, str]] = [] + + async def fetch(self, user_id: str, server_id: str) -> Optional[OAuthToken]: + self.calls.append((user_id, server_id)) + return self._values.get((user_id, server_id)) + + +class _Clock: + def __init__(self, t: float = 1000.0) -> None: + self.t = t + + def __call__(self) -> float: + return self.t + + +async def test_serves_token_until_its_expiry(): + token = OAuthToken(access_token="at", expires_at=1100.0) + inner = _FakeStore({("u", "s"): token}) + clock = _Clock(1000.0) + store = CachedOAuthTokenStore( + inner, default_ttl_seconds=60, expiry_skew_seconds=30, clock=clock + ) + + assert await store.fetch("u", "s") is token + clock.t = 1060.0 # still before expiry - skew (1100 - 30 = 1070) + assert await store.fetch("u", "s") is token + assert inner.calls == [("u", "s")] # served from cache, store hit once + + +async def test_refetches_once_token_has_expired(): + token = OAuthToken(access_token="at", expires_at=1100.0) + inner = _FakeStore({("u", "s"): token}) + clock = _Clock(1000.0) + store = CachedOAuthTokenStore( + inner, default_ttl_seconds=60, expiry_skew_seconds=30, clock=clock + ) + + await store.fetch("u", "s") + clock.t = 1080.0 # past expiry - skew (1070) + await store.fetch("u", "s") + assert len(inner.calls) == 2 # re-read after the cached token expired + + +async def test_caches_not_authorized_none_for_default_ttl(): + inner = _FakeStore({}) # user has not authorized + clock = _Clock(1000.0) + store = CachedOAuthTokenStore(inner, default_ttl_seconds=60, clock=clock) + + assert await store.fetch("u", "s") is None + clock.t = 1059.0 + assert await store.fetch("u", "s") is None + assert inner.calls == [("u", "s")] # None cached for the TTL window + + +async def test_default_ttl_applies_to_tokens_without_expiry(): + token = OAuthToken(access_token="at", expires_at=None) + inner = _FakeStore({("u", "s"): token}) + clock = _Clock(1000.0) + store = CachedOAuthTokenStore(inner, default_ttl_seconds=60, clock=clock) + + await store.fetch("u", "s") + clock.t = 1061.0 + await store.fetch("u", "s") + assert len(inner.calls) == 2 # no-expiry token re-read after the default TTL + + +async def test_invalidate_forces_refetch(): + inner = _FakeStore({}) + store = CachedOAuthTokenStore(inner, default_ttl_seconds=60, clock=_Clock()) + + assert await store.fetch("u", "s") is None + inner._values[("u", "s")] = OAuthToken(access_token="fresh") + store.invalidate("u", "s") + result = await store.fetch("u", "s") + assert result is not None and result.access_token == "fresh" + + +async def test_store_unavailable_is_not_cached(): + class _FailingStore: + def __init__(self) -> None: + self.calls = 0 + + async def fetch(self, user_id: str, server_id: str) -> Optional[OAuthToken]: + self.calls += 1 + raise TokenStoreUnavailable("down") + + inner = _FailingStore() + store = CachedOAuthTokenStore(inner, default_ttl_seconds=60, clock=_Clock()) + + for _ in range(2): + with pytest.raises(TokenStoreUnavailable): + await store.fetch("u", "s") + assert inner.calls == 2 # outage re-attempted, not cached + + +async def test_isolates_by_subject(): + a = OAuthToken(access_token="a") + b = OAuthToken(access_token="b") + inner = _FakeStore({("u1", "s"): a, ("u2", "s"): b}) + store = CachedOAuthTokenStore(inner, default_ttl_seconds=60, clock=_Clock()) + + first = await store.fetch("u1", "s") + 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"