feat(mcp): migrate authorization_code MCP to the v2 resolver and a v2-native token store - #31269
feat(mcp): migrate authorization_code MCP to the v2 resolver and a v2-native token store#31269tin-berri wants to merge 29 commits into
Conversation
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
0c7fdde to
33ce19a
Compare
57bbf5a to
8cceb45
Compare
33ce19a to
0e93dac
Compare
0e93dac to
f275aac
Compare
8e3f073 to
acfee1e
Compare
3a2dca6 to
13ea490
Compare
6b2a061 to
c7cfd84
Compare
7f2e4e2 to
d85da65
Compare
Greptile SummaryThis PR migrates the
Confidence Score: 4/5The change is safe to merge. The v2 path is additive and gated by The architecture is sound: the strangler pattern is correctly applied, the
|
| Filename | Overview |
|---|---|
| litellm/proxy/_experimental/mcp_server/db.py | Adds resolve_user_oauth_access_token (shared token-resolution core) and _remaining_token_seconds helper. Minor: prefetched_creds typed as dict[str, object] instead of dict[str, Any], mismatching callers. |
| litellm/proxy/_experimental/mcp_server/mcp_server_manager.py | Wires V1PerUserTokenStore into UpstreamCredentialProvider; adds raise_user_oauth_challenge on unauthorized; adds to_server_spec guard in _resolve_oauth2_headers_for_tool_call. Logic is sound. |
| litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py | Adds MCPAuth.oauth2 branch to to_server_spec and raise_user_oauth_challenge. Minor: www-authenticate header is lowercase while RFC 6750 and sibling raise_public use WWW-Authenticate. |
| litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py | Activates the AuthorizationCodeConfig arm; _NullOAuthTokenStore default ensures fail-closed behavior. Clean implementation. |
| litellm/proxy/_experimental/mcp_server/outbound_credentials/v1_token_store.py | New strangler adapter bridging v2 OAuthTokenStore seam to v1's token resolution. Well-structured. |
| litellm/proxy/_experimental/mcp_server/server.py | Refactors _get_user_oauth_extra_headers_from_db to thin wrapper; adds migrated_to_v2 guard on tools/list path. Correct. |
Reviews (1): Last reviewed commit: "feat(mcp): cut the tools/list connection..." | Re-trigger Greptile
The unauthorized case becomes a structured Unauthorized (detail + optional WWW-Authenticate header + optional structured body) instead of a bare string, and raise_public emits the header and body when present. This lets a mode reproduce a rich 401 challenge (e.g. BYOK's provisioning prompt) through the generic resolver edge. of_unauthorized's new params are keyword-only and default to None, so existing callers and the summary string are unchanged.
…t flat CredError's unauthorized payload was a pydantic BaseModel, whose base resolves as unknown in this repo's basedpyright (every model in the file trips reportUntypedBaseClass plus an unknown model_config), so the tagged-union case read as unknown and the public edge's challenge access added reportUnknownMemberType errors over the per-rule ceiling. A frozen dataclass is fully typed here, so error.unauthorized resolves directly with no cast or accessor and the per-rule basedpyright counts match base.
…ion_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.
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.
…ports in the token modules
… refresh_token docstring
…aching) CachedOAuthTokenStore no longer caches the "not authorized" None result; every miss re-reads the inner store. v1's per-user token cache never caches misses, so a token written by the OAuth flow is visible on the next request without an invalidation hook, and uniformly across replicas since the in-process cache holds no stale None to clear. invalidate() now only covers rotation or revocation of a cached token. Negative caching (with distributed invalidation) can return later if a slow DB-backed v2-native source makes per-miss reads expensive.
The proactive token-refresh / cache-expiry buffer defaulted to 30s, which is an outlier among OAuth clients. Spring Security uses 60s as both its JWT clock-skew tolerance and its refresh buffer, and 60s sits inside RFC 7519's "a few minutes" leeway while preserving nearly all of a typical token's life; 30s was untested, so pin the default with two boundary-probe regression tests.
81c7402 to
54414ff
Compare
41371d9 to
47c4536
Compare
The refresh seam took only the OAuthToken, but a refresher needs the server's config (token endpoint, client credentials, scopes) to run the grant and the (user_id, server_id) key to persist the minted token, neither of which is derivable from the token. Widen TokenRefresher.refresh to (user_id, server_id, token) and pass them through from RefreshingTokenStore so each stacked mode PR plugs into the final seam rather than forcing a later signature change across the stack.
47c4536 to
32e3736
Compare
…replica token caching) Make CachedOAuthTokenStore's storage and RefreshingTokenStore's single-flight injectable so a cross-replica deployment can back them with Redis without touching the resolver. The defaults preserve today's behavior exactly: InMemoryTokenCacheBackend (the bounded per-process dict) and InProcessRefreshCoordinator (the asyncio single-flight). A distributed deployment injects a shared DualCache-backed backend and a SET NX PX coordinator. invalidate() is now async (the backend may be). The cache stores via the backend with a TTL derived from the token's expiry; the coordinator threads a reread callback for the cross-replica case (losers re-read the persisted token) that the in-process default ignores.
Resolve a user's authorization_code token through the injected OAuthTokenStore: present -> Authorization: Bearer <access_token>; absent -> the RFC 9728 WWW-Authenticate OAuth challenge; store unavailable -> the same challenge (not a 500), since a transient outage is not a definite absence. UpstreamCredentialProvider gains the oauth_token_store collaborator (fail-closed null default); per-subject isolation comes from keying the fetch on subject_id. Not live until to_server_spec maps authorization_code and a v1-backed token source is wired (next steps).
V1PerUserTokenStore reads the user's stored access token through v1's mcp_per_user_token_cache (Redis-backed, encrypted) and wraps it in an OAuthToken. v1 holds only the access token (its cache TTL is the lifetime), so no expires_at/refresh_token yet; the v2 cache holds it for its default TTL and the OAuth challenge drives re-auth once v1's cache drops it. Additive: nothing wires it yet, so no behavior change. Step 1b swaps it for a v2-native token store behind the OAuthTokenStore seam.
… refresh-capable Extract v1's per-user OAuth egress (Redis cache, else DB read with the refresh_token grant, then re-cache) from _get_user_oauth_extra_headers_from_db into resolve_user_oauth_access_token in db.py; the v1 header builder is now a thin wrapper over it and its callers are unchanged. V1PerUserTokenStore (the v2 OAuthTokenStore adapter) resolves through that same core via an injected server lookup, so the authorization_code arm injects exactly the token v1 would, with the same silent refresh, rather than a Redis-only read that can never refresh. One resolution implementation, two thin adapters (header dict and OAuthToken). Behavior-preserving: the existing v1 egress tests pass unchanged, and the arm is not wired into the live path yet (that lands with to_server_spec + the manager).
… the v2 resolver to_server_spec maps an oauth2 server to AuthorizationCodeConfig when it relies on per-user tokens (needs_user_oauth_token and not delegate_auth_to_upstream); client_credentials (M2M), delegated upstream OAuth, token exchange, and SigV4 still defer to v1. The manager injects V1PerUserTokenStore (resolving through v1's shared egress core) into the credential provider. The v2 path is live but still defers to a token v1 places in extra_headers; the cutover that makes v1 step aside lands next, alongside the unified challenge.
When an authorization_code server has no usable per-user token, the arm returns a semantic
unauthorized and the graft builds the 401 where the full MCPServer is in hand: a relative,
per-server RFC 9728 resource_metadata pointer (/.well-known/oauth-protected-resource/mcp/{name})
that names the server's own authorization server, instead of the resolver's earlier root pointer
which resolved to the gateway's generic PRM. Relative, so it is correct behind a reverse proxy
without request context. The listing-phase 401 still emits the RFC 8414 authorization_uri form;
both now target the same server, so the remaining difference is cosmetic and unifies in a later PR.
… servers _resolve_oauth2_headers_for_tool_call steps aside (builds no header) when to_server_spec maps the server, so the v2 resolver drives the token-present case instead of being shadowed by a token v1 places in extra_headers. Non-migrated oauth2 (delegate, client_credentials) and BYOK still build their header on v1. With this, v2 owns the authorization_code egress end to end: inject the refreshed per-user token when present, raise the per-server fail-closed 401 when absent.
…_code servers The listing connection's per-user OAuth header is no longer built by v1 for migrated servers; the v2 resolver drives it at connect time, ending the double-resolution where v1 built the token into extra_headers and the v2 graft then deferred to it. Safe because the preemptive 401 (in the streamable-http and SSE handlers) already challenges a missing token before the listing connection runs, so the connection is only reached with a token present. Non-migrated oauth2 (delegate) and the rest still build their header on v1. With this, resolve_credentials' result is honored on every authorization_code upstream path: tool calls and listing.
…solver The discovery-phase 401 no longer calls v1's _get_user_oauth_extra_headers_from_db to decide whether a migrated server has a token; it asks the v2 resolver via a new has_user_oauth_token manager method (to_server_spec + to_subject + resolve_credentials, Ok means a token exists). With this, every authorization_code resolution runs through the v2 resolver: the call_tool egress, the listing connection, and the discovery challenge. Delegate servers short-circuit before the check (the client completes PKCE with the upstream). The challenge itself still emits the RFC 8414 authorization_uri form; the format unification stays a follow-up.
Mirror the api_key arm's structure: the inline AuthorizationCodeConfig body moves into _authorization_code(subject, server), keeping resolve_credentials a flat one-line-per-arm dispatch. The helper is annotated with the concrete StaticHeaderAuth it returns rather than the abstract httpx.Auth (which api_key uses) because a new method carrying the unresolved httpx.Auth return would add reportUnknownMemberType; the concrete type is both precise and budget-neutral.
…h challenge raise_user_oauth_challenge emitted the header lowercase while the sibling raise_public and every resource_metadata (RFC 9728) emitter use the canonical WWW-Authenticate; align it. HTTP header names are case-insensitive on the wire so this is cosmetic for compliant clients, but it keeps the challenge builders consistent and matches RFC 6750.
Reads the user's persisted authorization_code credential and returns a typed OAuthToken (access token, epoch expiry, refresh token), validating the decoded blob at this boundary so no Any leaks past it. The raw inner store that RefreshingTokenStore/CachedOAuthTokenStore wrap; the DB read + decode collaborator is injected so it stays testable. Not yet wired - V1PerUserTokenStore is still the composition-root store until the refresher and cross-worker cache land.
The refresh_token grant for the authorization_code mode: POSTs the RFC 6749 refresh_token grant to the server's token endpoint, persists the rotated triple, and returns the new typed OAuthToken for RefreshingTokenStore to cache. HTTP post and persist are injected so the grant + response parsing are testable without a live IdP/DB. Also extends the TokenRefresher seam with (user_id, server_id), which the foundation's refresh(token) lacked but the grant (server config) and persist (key) need.
…(step 1b §1.5) The serialize+encrypt boundary a cross-replica cache needs: a plaintext bearer in Redis is a leak, so encode() encrypts (NaCl in prod via the injected encrypt, identity in tests). Caches only access_token and expires_at, never the refresh_token - the hot path needs just the bearer, and the long-lived refresh_token stays in the DB (the refresh path is always a cache miss), matching v1. A decoded token always has refresh_token=None. Undecryptable (key rotation) or corrupt entries read as a miss.
32e3736 to
9de7157
Compare
The cross-replica TokenCacheBackend implementation that plugs into the foundation's CachedOAuthTokenStore seam: encrypts+serializes the token via the codec and stores it in LiteLLM's shared DualCache under the same per-(user,server) key v1 used, so workers share one refresh and a token cached by v1 or v2 is readable by the other across the cutover. Cache and codec are injected; a non-positive TTL (already-expired token) is not cached, and a missing/corrupt entry reads as a miss.
The cross-replica RefreshCoordinator that plugs into the foundation's RefreshingTokenStore seam: a SET NX PX lock elects one worker to refresh per (user, server) while the rest wait for it and re-read the token it persisted, so a rotating refresh_token is used once across the fleet, not once per worker. The lock self-expires (PX) so a crashed holder can't wedge refresh; a loser falls back to a bounded re-read and the surrounding store re-checks expiry next fetch, so a crash self-heals. The lock (a thin Redis SET NX/DEL/EXISTS wrapper in prod) is injected, so the single-flight logic is testable without Redis.
The concrete DistributedLock the RedisRefreshCoordinator elects refreshers with: acquire is an atomic SET key NX PX ttl (first caller wins, entry self-expires so a crashed holder can't wedge refresh), release is DEL, is_held is EXISTS. The async Redis client is injected (the client from LiteLLM's RedisCache in prod), so it is unit-testable with a fake. Any Redis error degrades to not-acquired / not-held so a cache blip causes an extra refresh, never a crash on the resolve path.
…(step 1b piece 4) Assemble Cached(Refreshing(V2PerUserTokenStore)) at the composition root and replace V1PerUserTokenStore in mcp_server_manager. The chain is built lazily on first fetch (its cache/DB/ Redis collaborators are LiteLLM globals not ready at import); when Redis is wired it uses the cross-replica path (DualCache cache + SET NX PX coordinator), else the in-process defaults. The DB read, refresh-grant POST, and persist acquire their globals per call like v1. authorization_code resolution now reads/refreshes through the v2-native lifecycle, not v1's core.
…b piece 5) Piece 4 replaced V1PerUserTokenStore with the v2-native chain at the composition root, leaving the adapter with no callers, so remove it and its test. The shared v1 read/refresh core (resolve_user_oauth_access_token and friends) stays - delegate's egress in server.py still uses it - and comes out with the delegate migration.
Relevant issues
The
authorization_code(per-user OAuth 3LO) MCP path, migrated end to end onto the v2 outbound-credential resolver and a v2-native token lifecycle. Stacked on the OAuth token foundation (#31275), which carries the shared seams (the token model, theOAuthTokenStoreandTokenRefresherseams, and the injectable cache backend and refresh coordinator); this diff is the authorization_code-specific code plus the cross-replica implementations of those seamsLinear ticket
N/A (MCP V2 migration)
Pre-Submission checklist
make test-unitType
🆕 New Feature
Changes
Two halves ship together, so authorization_code is fully v2-native with no v1 fallback left in its path
The resolver dispatch. Every per-user OAuth upstream call now resolves through the v2 resolver: the call_tool egress, the
tools/listconnection, and the preemptive discovery 401 route throughresolve_credentialsandhas_user_token, andto_server_specmaps anoauth2server toAuthorizationCodeConfigwhen it holds per-user tokens (needs_user_oauth_tokenand notdelegate_auth_to_upstream). When the user has not authorized, the arm fails closed and the egress raises a per-server RFC 9728resource_metadatachallenge, in the canonicalWWW-Authenticateform that the other resource_metadata emitters useThe v2-native token lifecycle. The resolver no longer reads through v1's core.
V2PerUserTokenStorereads the persisted credential and validates it into a typedOAuthTokenat the boundary;AuthorizationCodeRefresherruns the refresh_token grant and persists the rotated triple; both sit under the foundation'sRefreshingTokenStoreandCachedOAuthTokenStore. The chain is assembled at the composition root and built lazily on the first request, since its cache, DB, and Redis collaborators are runtime globals not ready at importCross-replica safety, matching v1. v1 cached per-user tokens in a shared
DualCacheand serialized refreshes with a Postgres advisory lock, so a rotating refresh_token is used once across workers rather than once per worker. The v2 path keeps that guarantee:DualCacheTokenCacheBackendcaches the access token (encrypted with the same NaCl helper and key as v1, so a token cached by either is readable by the other across the cutover; the long-lived refresh_token stays in the DB), andRedisRefreshCoordinatorwithRedisDistributedLockelect one refresher per (user, server) viaSET NX PX. With no Redis configured, the foundation's in-process defaults apply, which is correct for a single replica. Proactive (near-expiry) refresh ships here; reactive-401 refresh is deferred to the egress transport that sees the upstream 401What is removed, what stays.
V1PerUserTokenStore, the strangler adapter, is deleted now that the v2 chain is wired. The shared v1 read/refresh core (resolve_user_oauth_access_tokenand friends) stays because delegate's egress still calls it; it exits with the delegate migrationScreenshots / Proof of Fix
To be added: curl against a live proxy hitting a real authorization_code MCP server. A tool call resolving the stored bearer, then a near-expiry call showing the refresh_token grant round-trip the IdP and the rotated token persist, with the request still succeeding