Skip to content

feat(mcp): migrate authorization_code MCP to the v2 resolver and a v2-native token store - #31269

Open
tin-berri wants to merge 29 commits into
litellm_internal_stagingfrom
litellm_mcp_v2_authz_code
Open

feat(mcp): migrate authorization_code MCP to the v2 resolver and a v2-native token store#31269
tin-berri wants to merge 29 commits into
litellm_internal_stagingfrom
litellm_mcp_v2_authz_code

Conversation

@tin-berri

@tin-berri tin-berri commented Jun 25, 2026

Copy link
Copy Markdown
Contributor

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, the OAuthTokenStore and TokenRefresher seams, and the injectable cache backend and refresh coordinator); this diff is the authorization_code-specific code plus the cross-replica implementations of those seams

Linear ticket

N/A (MCP V2 migration)

Pre-Submission checklist

  • I have added meaningful tests
  • My PR passes all unit tests on make test-unit
  • My PR's scope is as isolated as possible; it only solves 1 specific problem
  • I have requested a Greptile review and received a Confidence Score of at least 4/5

Type

🆕 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/list connection, and the preemptive discovery 401 route through resolve_credentials and has_user_token, and to_server_spec maps an oauth2 server to AuthorizationCodeConfig when it holds per-user tokens (needs_user_oauth_token and not delegate_auth_to_upstream). When the user has not authorized, the arm fails closed and the egress raises a per-server RFC 9728 resource_metadata challenge, in the canonical WWW-Authenticate form that the other resource_metadata emitters use

The v2-native token lifecycle. The resolver no longer reads through v1's core. V2PerUserTokenStore reads the persisted credential and validates it into a typed OAuthToken at the boundary; AuthorizationCodeRefresher runs the refresh_token grant and persists the rotated triple; both sit under the foundation's RefreshingTokenStore and CachedOAuthTokenStore. 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 import

Cross-replica safety, matching v1. v1 cached per-user tokens in a shared DualCache and 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: DualCacheTokenCacheBackend caches 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), and RedisRefreshCoordinator with RedisDistributedLock elect one refresher per (user, server) via SET 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 401

What 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_token and friends) stays because delegate's egress still calls it; it exits with the delegate migration

Screenshots / 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

@codecov

codecov Bot commented Jun 25, 2026

Copy link
Copy Markdown

@tin-berri
tin-berri force-pushed the litellm_mcp_v2_authz_code branch 2 times, most recently from 0c7fdde to 33ce19a Compare June 25, 2026 04:25
@tin-berri
tin-berri changed the base branch from litellm_internal_staging to litellm_mcp_v2_oauth_token_foundation June 25, 2026 04:25
@tin-berri
tin-berri force-pushed the litellm_mcp_v2_oauth_token_foundation branch from 57bbf5a to 8cceb45 Compare June 25, 2026 04:36
@tin-berri
tin-berri force-pushed the litellm_mcp_v2_authz_code branch from 33ce19a to 0e93dac Compare June 25, 2026 04:36
@tin-berri tin-berri changed the title feat(mcp): authorization_code foundation - 401 challenge + OAuth token store/cache feat(mcp): authorization_code resolver arm + v1-backed token source Jun 25, 2026
@tin-berri
tin-berri force-pushed the litellm_mcp_v2_authz_code branch from 0e93dac to f275aac Compare June 25, 2026 05:57
@tin-berri
tin-berri force-pushed the litellm_mcp_v2_oauth_token_foundation branch from 8e3f073 to acfee1e Compare June 25, 2026 06:01
@tin-berri
tin-berri force-pushed the litellm_mcp_v2_authz_code branch from 3a2dca6 to 13ea490 Compare June 25, 2026 06:08
@tin-berri
tin-berri force-pushed the litellm_mcp_v2_oauth_token_foundation branch from 6b2a061 to c7cfd84 Compare June 25, 2026 16:40
@tin-berri
tin-berri force-pushed the litellm_mcp_v2_authz_code branch 2 times, most recently from 7f2e4e2 to d85da65 Compare June 25, 2026 17:01
@tin-berri tin-berri changed the title feat(mcp): authorization_code resolver arm + v1-backed token source feat(mcp): migrate the authorization_code MCP egress to the v2 resolver Jun 25, 2026
@tin-berri
tin-berri marked this pull request as ready for review June 25, 2026 22:41
@greptile-apps

greptile-apps Bot commented Jun 25, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR migrates the authorization_code (per-user OAuth 3LO) MCP egress to the v2 credential resolver. It extracts the shared token-resolution core (resolve_user_oauth_access_token) from the v1 header builder, wires a V1PerUserTokenStore adapter into the UpstreamCredentialProvider, and activates the AuthorizationCodeConfig resolver arm — replacing the previous _not_implemented stub.

  • to_server_spec now maps per-user oauth2 servers to AuthorizationCodeConfig; M2M (client_credentials), delegated, token-exchange, and SigV4 servers stay on v1. Both egress paths (call_tool and tools/list) gain a to_server_spec guard that steps aside so the v2 resolver owns the token instead of being shadowed by a v1-built header.
  • raise_user_oauth_challenge emits a per-server RFC 9728 resource_metadata challenge (relative URL, correct behind a reverse proxy) instead of the root protected-resource metadata used previously.
  • Test coverage is added for all new paths with mock-only tests; the modified existing test is adjusted to reflect the behavioral split between migrated and non-migrated servers.

Confidence Score: 4/5

The change is safe to merge. The v2 path is additive and gated by to_server_spec; non-migrated servers are entirely unaffected.

The architecture is sound: the strangler pattern is correctly applied, the _NullOAuthTokenStore default is fail-closed, the migrated_to_v2 guards prevent double-resolution on both egress paths, and test coverage is comprehensive. Two minor issues were found: raise_user_oauth_challenge emits www-authenticate in lowercase while its sibling raise_public uses WWW-Authenticate (RFC 6750 specifies mixed case), and prefetched_creds in resolve_user_oauth_access_token is annotated dict[str, object] rather than dict[str, Any], misrepresenting actual callers.

adapter.py for the header-case inconsistency in raise_user_oauth_challenge; db.py for the type annotation on prefetched_creds.

Important Files Changed

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

Comment thread litellm/proxy/_experimental/mcp_server/db.py
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.
…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.
@tin-berri
tin-berri force-pushed the litellm_mcp_v2_oauth_token_foundation branch from 81c7402 to 54414ff Compare June 26, 2026 00:23
@tin-berri
tin-berri force-pushed the litellm_mcp_v2_authz_code branch from 41371d9 to 47c4536 Compare June 26, 2026 01:10
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.
@tin-berri
tin-berri force-pushed the litellm_mcp_v2_authz_code branch from 47c4536 to 32e3736 Compare June 26, 2026 04:46
…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.
tin-berri added 14 commits June 25, 2026 22:19
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.
@tin-berri
tin-berri force-pushed the litellm_mcp_v2_authz_code branch from 32e3736 to 9de7157 Compare June 26, 2026 05:20
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.
@tin-berri tin-berri changed the title feat(mcp): migrate the authorization_code MCP egress to the v2 resolver feat(mcp): migrate authorization_code MCP to the v2 resolver and a v2-native token store Jun 26, 2026
Base automatically changed from litellm_mcp_v2_oauth_token_foundation to litellm_internal_staging June 27, 2026 01:27
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant