feat(mcp): OAuth token store seam + expiry-aware cache for authorization_code

Lay the foundation for the authorization_code resolver arm: OAuthToken (access_token,
expires_at, refresh_token), the OAuthTokenStore Protocol seam, TokenStoreUnavailable for
outages, and CachedOAuthTokenStore, an expiry-aware cache that serves a token only while
unexpired, caches the "not authorized" None for a default TTL, and propagates a store outage
without caching it. Mirrors the BYOK store/cache pattern, adapted for tokens. Refresh and
distributed single-flight are deferred to the hardening step.

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)

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"