Skip to content

Revert "feat(mcp): cross-replica single-flight refresh for the v2 per-user OAuth store [2/2]"#31492

Merged
tin-berri merged 1 commit into
litellm_mcp_v2_authz_code_dispatchfrom
revert-31474-litellm_mcp_v2_authz_code_xrepl
Jun 27, 2026
Merged

Revert "feat(mcp): cross-replica single-flight refresh for the v2 per-user OAuth store [2/2]"#31492
tin-berri merged 1 commit into
litellm_mcp_v2_authz_code_dispatchfrom
revert-31474-litellm_mcp_v2_authz_code_xrepl

Conversation

@tin-berri

Copy link
Copy Markdown
Contributor

Reverts #31474

@tin-berri tin-berri merged commit 5d19529 into litellm_mcp_v2_authz_code_dispatch Jun 27, 2026
34 of 36 checks passed
@tin-berri tin-berri deleted the revert-31474-litellm_mcp_v2_authz_code_xrepl branch June 27, 2026 03:56
@greptile-apps

greptile-apps Bot commented Jun 27, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR reverts the cross-replica OAuth token refresh layer. The main changes are:

  • Removes the Redis refresh coordinator and distributed lock implementation
  • Removes the DualCache-backed token backend and encrypted cache codec
  • Restores the per-user OAuth store to in-process cache and refresh coordination
  • Deletes tests for the reverted Redis and shared-cache pieces

Confidence Score: 4/5

Not safe to merge until the stale-token reread path is fixed.

The changed OAuth refresh path has a concrete stale-token failure mode, while the revert is otherwise narrow and covered by focused token-store context.

litellm/proxy/_experimental/mcp_server/outbound_credentials/oauth_token_store.py

T-Rex T-Rex Logs

What T-Rex did

  • Ran a focused async Python harness against the real oauth_token_store.py module with a fake expired token store and a coordinator that only invokes reread.
  • Compared cross-replica outputs from before and after the revert, using the cross-replica inspection script to verify the presence/absence of cross-replica files and the use of in-process defaults.

View all artifacts

T-Rex Ran code and verified through T-Rex

Reviews (1): Last reviewed commit: "Revert "feat(mcp): cross-replica single-..." | Re-trigger Greptile

server_id,
refresh=refresh_latest_token,
reread=reread_fresh_token,
reread=lambda: self._inner.fetch(user_id, server_id),

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Stale token re-read
reread now returns the raw inner-store value without applying the same expiry check used before refresh. When an injected coordinator loses its election and the winning refresh fails, the DB still contains the expired token, so this path returns a stale bearer instead of None; the resolver will send that expired access token upstream rather than challenging the user.

Artifacts

Repro: focused async harness for stale token reread

  • Contains supporting evidence from the run (text/x-python; charset=utf-8).

Repro: harness output showing expired token returned from fetch

  • Keeps the command output available without making the summary code-heavy.

View artifacts

T-Rex Ran code and verified through T-Rex

@codecov

codecov Bot commented Jun 27, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 0% with 2 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
...erver/outbound_credentials/per_user_oauth_store.py 0.00% 2 Missing ⚠️

📢 Thoughts on this report? Let us know!

tin-berri added a commit that referenced this pull request Jun 27, 2026
…replica) [1/2] (#31473)

* feat(mcp): implement the authorization_code resolver arm

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).

* feat(mcp): v1-backed OAuth token source for authorization_code

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.

* style(mcp): modern type annotations in the authorization_code arm and source

* refactor(mcp): share v1's OAuth egress core; make V1PerUserTokenStore 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).

* feat(mcp): route oauth2 per-user (authorization_code) servers through 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.

* feat(mcp): per-server fail-closed OAuth challenge at the v2 egress

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.

* feat(mcp): cut the call_tool egress over to v2 for authorization_code 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.

* feat(mcp): cut the tools/list connection over to v2 for authorization_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.

* feat(mcp): route the preemptive 401 existence check through the v2 resolver

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.

* refactor(mcp): extract the authorization_code arm into a helper

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.

* fix(mcp): emit the canonical WWW-Authenticate header name in the OAuth 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.

* feat(mcp): v2-native per-user token read store (step 1b inner store)

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.

* feat(mcp): v2-native authorization_code token refresher (step 1b)

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.

* feat(mcp): wire the v2-native per-user OAuth store into the resolver (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.

* refactor(mcp): delete the unwired V1PerUserTokenStore adapter (step 1b 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.

* fix(mcp): green CI for authz_code dispatch (format + UTC expiry + v2-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]>

* fix(mcp): caller Authorization must not override the stored per-user 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]>

* fix(mcp): preserve recorded OAuth scopes across authorization_code refresh

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

* fix(mcp): keep user token in authorization_code tools preview

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

* fix(mcp): stop caller-supplied auth from overriding stored authorization_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

* feat(mcp): cross-replica single-flight refresh for the v2 per-user OAuth 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]>

* Revert "feat(mcp): cross-replica single-flight refresh for the v2 per-user OA…" (#31492)

This reverts commit cd2fb6b.

---------

Co-authored-by: Claude Opus 4.8 (1M context) <[email protected]>
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