feat(mcp): migrate authorization_code MCP to the v2 resolver (single-replica) [1/2]#31473
Conversation
Greptile SummaryThis PR migrates the
Confidence Score: 5/5Safe to merge. All three prior review findings are correctly resolved, and no new defects were introduced. The three issues flagged in the previous round — UTC-naive expiry interpretation, dropped grant scopes on refresh, and the delegate discovery path — are all addressed correctly. The UTC fix anchors timezone-naive datetimes to UTC before calling No files require special attention.
|
| Filename | Overview |
|---|---|
| litellm/proxy/_experimental/mcp_server/outbound_credentials/authz_code_refresher.py | New v2 refresher for authorization_code: correctly carries scopes and refresh_token forward when IdP omits them, matches v1 parity. Well-tested. |
| litellm/proxy/_experimental/mcp_server/outbound_credentials/v2_token_store.py | New per-user token read store with correct UTC-anchoring for timezone-naive expires_at, fixing the prior finding from the previous review round. |
| litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py | Composition root assembling Cached(Refreshing(V2PerUserTokenStore)) with injected DB/HTTP collaborators; lazy construction defers globals until runtime. Clean. |
| litellm/proxy/_experimental/mcp_server/outbound_credentials/presented_token_store.py | One-shot store for the create/test preview path; intentionally ignores the lookup key since it backs a single call. Correct by design. |
| litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py | authorization_code arm wired in, fail-closed null store default, has_user_token predicate added for preemptive-401 check. Store outage maps to None (challenge), not 500. |
| litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py | to_server_spec now maps oauth2+needs_user_oauth_token+non-delegate to AuthorizationCodeConfig; raise_user_oauth_challenge emits RFC 9728 resource_metadata 401. |
| litellm/proxy/_experimental/mcp_server/mcp_server_manager.py | _should_strip_caller_authorization extended for authorization_code; _create_mcp_client keeps v2 spec for authorization_code even when mcp_auth_header present; has_user_oauth_token exposed for discovery path. |
| litellm/proxy/_experimental/mcp_server/server.py | Discovery preemptive-401 now routes through has_user_oauth_token (v2 resolver) instead of the v1 DB check. Delegate path unchanged. Authorization stripping added to tools-list egress. |
| litellm/proxy/_experimental/mcp_server/rest_endpoints.py | Preview path injects a one-shot PresentedOAuthTokenStore via cred_provider instead of forwarding the caller header; Authorization stripped from effective_oauth2_headers for authorization_code preview. |
| litellm/proxy/_experimental/mcp_server/db.py | resolve_user_oauth_access_token and _remaining_token_seconds extracted as shared helpers; UTC anchoring present in _remaining_token_seconds. Parity with v2_token_store._iso_to_epoch. |
Reviews (3): Last reviewed commit: "fix(mcp): stop caller-supplied auth from..." | Re-trigger Greptile
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
PR overviewAll previously flagged issues have been addressed. No open security concerns remain on this pull request. Security reviewNo open security issues remain on this pull request. Fixed/addressed: 2 · PR risk: 0/10 |
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 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.
…seam tests) - ruff format per_user_oauth_store.py (clears the lint check) - v2_token_store._iso_to_epoch: anchor a tz-naive expiry to UTC before .timestamp(), matching v1's db.py _remaining_token_seconds (Greptile P1) so a non-UTC host doesn't read the expiry as local time and skew refresh timing - test_mcp_stale_session: repoint the 3 discovery tests off the removed v1 _get_user_oauth_extra_headers_from_db onto the v2 has_user_oauth_token seam; the delegate test now asserts the existence check is never consulted (delegate short-circuits to the resource_metadata 401 before any token lookup) - test_mcp_server_manager: repoint test_deferred_mode_uses_v1_auth_value at M2M (oauth2 client_credentials), which is still a deferred mode, since per-user oauth2 (authorization_code) now routes to the v2 resolver Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
74eb1e0 to
0f21603
Compare
…OAuth token A caller with a valid x-litellm-api-key could include their own "Authorization: Bearer <chosen>" header and have the proxy execute tools against that bearer instead of the user's stored OAuth credential. For a v2-migrated authorization_code server the caller's Authorization was seeded into extra_headers, and the graft's apply-if-absent then dropped the resolved per-user token in its favor. v1 prevented this by overwriting a stale client Authorization with the stored token; this restores that precedence on both egress paths (connect + call_tool). - _should_strip_caller_authorization: also strip for migrated per-user OAuth (authorization_code) servers - the v2 resolver injects the stored token, so a caller-forwarded Authorization must not be forwarded upstream. Delegate / pass-through (to_server_spec is None) keep forwarding the caller's bearer. - both seed sites (_prepare_mcp_server_headers, _call_regular_mcp_tool) drop only the Authorization from the caller's oauth2_headers (via _without_authorization), keeping any other forwarded header and any hook/static Authorization (which still wins, as in v1). - regression test for the call_tool path; updated the two tests that asserted the old (vulnerable) forwarding to assert the secure behavior. Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
0f21603 to
596abec
Compare
|
@greptile-apps rereview pls |
…fresh
When a refresh response omits `scope` (RFC 6749 §5.1, where omission means unchanged), the v2 refresher persisted scopes=None and overwrote the user's recorded grant. v1 carried the prior scopes forward via `or cred.get("scopes")`; the v2 path lost that because OAuthToken did not model scopes
OAuthToken now carries scopes, V2PerUserTokenStore populates them on read, and AuthorizationCodeRefresher carries them forward for both the persisted write and the returned/cached token, so repeated refreshes do not erode them. A present `scope` in the response still replaces the prior grant
Adds regression tests: a refresh omitting `scope` preserves the prior scopes, and a present `scope` overrides them
After to_server_spec maps oauth2 onto the v2 resolver, the interactive tools preview for an unsaved authorization_code server read the per-user token store, found nothing, and fail-closed with a 401, so the create/test tab could no longer list tools The preview now routes the just-authorized token (forwarded in oauth2_headers) through mcp_auth_header, so _create_mcp_client takes the per-request-override v1 path and uses it directly, matching v1's preview. Gated to the v2-mapped oauth2 case; M2M, delegate/passthrough, and token-exchange keep their existing preview path Adds tests: interactive oauth routes the forwarded token to mcp_auth_header, M2M and token-exchange do not
…ion_code tokens A caller-supplied per-request override (mcp_auth_header / x-mcp-auth / x-mcp-<alias>-authorization) disabled the v2 resolver in _create_mcp_client for any spec, so an authenticated user with a stored authorization_code token could force an arbitrary upstream bearer and bypass the stored credential and its save-time validation. _create_mcp_client now keeps the v2 spec for authorization_code and ignores the override; other modes keep the client-side-credentials override The create/test tools preview no longer relies on that override path. It resolves the just-authorized, not-yet-persisted token through the v2 resolver via a one-shot PresentedOAuthTokenStore passed as cred_provider - the same path runtime uses for the stored token - so preview and runtime resolve identically. This replaces the mcp_auth_header routing added earlier Adds tests: a caller override cannot bypass the v2 resolver for authorization_code; the interactive preview resolves via the presented store rather than a caller header; M2M and token-exchange build no presented provider
|
@greptileai review The latest commits address the prior findings. The authorization_code refresh now carries stored scopes forward when the IdP omits |
…uth store [2/2] (#31474) * feat(mcp): encrypt+serialize codec for caching OAuth tokens in Redis (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. * feat(mcp): DualCache-backed token cache backend (step 1b §1.5) 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. * feat(mcp): Redis SET NX PX refresh coordinator (step 1b §1.5) 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. * feat(mcp): Redis SET NX PX distributed lock (step 1b §1.5) 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. * feat(mcp): wire the cross-replica cache + coordinator into the per-user store (step 1b §1.5) Upgrade the composition root to use the DualCache-backed cache and SET NX PX refresh coordinator when Redis is wired, falling back to the foundation's in-process defaults on a single replica. Layers the cross-replica path on top of the single-replica dispatch store. Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]> * fix(mcp): refresh on lock-backend error instead of serving a stale token The cross-replica refresh coordinator elected refreshers with a boolean acquire: a Redis transport error was caught and returned as False, which is indistinguishable from "another worker holds the lock". On a total Redis outage every worker therefore took the wait-then-reread branch and served the still-expired token upstream (the upstream then 401s), even though the lock and coordinator docstrings claimed a Redis blip "degrades to an extra refresh". Make acquire tristate (LockAcquisition: ACQUIRED / HELD / ERROR) so the coordinator can tell a busy holder from a dead backend, and refresh anyway on ERROR. This single-flight lock is a load optimization, not a correctness mutex, so failing open is correct: it degrades a lock-backend outage to the no-coordinator behavior (an extra refresh), never a stale bearer. Add a regression test asserting an acquire error refreshes rather than re-reading the expired token, and update the docstrings to match. * style(mcp): wrap redis lock signatures at line-length 88 for CI ruff format * fix(mcp): a refresh loser surfaces None, not a stale token, when the winner failed The cross-replica coordinator's losers re-read the token the winner persisted. If the winner's refresh failed, the store still holds the expired token, so the loser re-read it and RefreshingTokenStore handed that expired bearer to the caller (the upstream then 401s) instead of the re-auth challenge the winner returned via None. Make the loser's re-read expiry-aware, mirroring refresh_latest_token: a re-read that is still expired surfaces None so the arm challenges. This only affects the loser path; the winner's freshly refreshed token is returned directly by the coordinator and is unaffected. --------- Co-authored-by: Claude Opus 4.8 (1M context) <[email protected]>
What this is
Split out of #31269. This is PR 1 of 2: 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, on a single replica. Builds on the OAuth token foundation #31275 (now merged intolitellm_internal_staging, so this PR is rebased onto staging).The cross-replica single-flight refresh (DualCache-backed cache +
SET NX PXcoordinator) is split into the follow-up #31474, stacked on this branch. This PR uses the foundation's in-process cache/coordinator defaults, which are correct for a single replica.Changes
call_toolegress, thetools/listconnection, and the preemptive discovery 401 route throughresolve_credentials/has_user_token;to_server_specmaps anoauth2server toAuthorizationCodeConfigwhen it holds per-user tokens.resource_metadatachallenge.V2PerUserTokenStore(read) andAuthorizationCodeRefresher(refresh_token grant + persist) replace v1'srefresh_user_oauth_token. The composition root assemblesCached(Refreshing(V2PerUserTokenStore))on the foundation's in-process defaults.V1PerUserTokenStoreadapter.Parity
Same cache key (
mcp:per_user_token:{user}:{server}), same NaCl encryption, same DualCache backend, same refresh-grant semantics (failed refresh = miss not 500, rotation carried forward) as v1.Review fixes
_iso_to_epochnow anchors a timezone-naiveexpires_atto UTC before.timestamp(), matching v1'sdb.py_remaining_token_seconds, so a non-UTC host doesn't read the expiry as local time and skew refresh timing._get_user_oauth_extra_headers_from_dbonto the v2has_user_oauth_tokenseam;test_deferred_mode_uses_v1_auth_valuewas repointed at M2M (oauth2client_credentials), which is still a deferred mode, since per-useroauth2now routes to v2.Type
🆕 New Feature
🤖 Generated with Claude Code