feat(mcp): make token_exchange (OBO) production-ready - discovery threading + audit hardening + RFC 9728 challenge

[tin-berri]

Relevant issues

Stacked on #31526 (the token_exchange / OBO arm). Base is that branch so the diff is just this work; it retargets to litellm_internal_staging after #31526 merges

Linear ticket

N/A

What this is

Makes token_exchange (OBO) production-ready on top of #31526. It began as the discovery fix (an OBO server's tools could not be listed through the aggregator) and then folded in a set of hardening fixes found by an adversarial audit of every OBO code path, credential shape, IdP failure mode, and third-party MCP client interaction. Each fix below is a separate commit with its own regression tests

1. Discovery threading (the original fix)

With #31526, OBO tool calls work, but an OBO server's tools could not be discovered. The tools/list path never threaded the caller's token, so every list hit the no-subject branch of the resolver. v1 masked this with a no-subject client_credentials fallback (discovery quietly used a service token); #31526 dropped that fallback per the v2 rule against falling through to a weaker source, so listing had no credential and the tools never appeared. Since an MCP client lists before it calls, the tools were invisible

The fix gives discovery the caller's own token: _get_tools_from_server extracts the inbound subject_token via the existing _extract_bearer_token and passes it into _create_mcp_client, gated on auth_type == oauth2_token_exchange so the bearer never leaks into other modes; server.py forwards oauth2_headers at the list call site. The list path's graceful degradation (a failed per-server fetch returns an empty list, never crashing the catalog) is preserved

2. Hardening fixes from the audit

a. A caller header can no longer bypass the exchange. The _create_mcp_client guard that ignores a per-request x-mcp-* override only protected authorization_code; on a token_exchange server the override dropped the v2 spec, skipped the RFC 8693 exchange, and forwarded the caller's raw header upstream. The guard now also covers token_exchange, so the exchange always runs and a caller cannot substitute an arbitrary upstream credential to defeat the audience-scoping

b. OBO now works for prompts and resources, not just tools. prompts/list, prompts/get, resources/list, resources/read, and resource-templates/list never threaded the subject token, so those operations failed closed (empty or 401) on an OBO server. They now thread the caller's bearer like the tools paths, via a shared _obo_subject_token helper gated on the mode

c. The resolver-owned credential is authoritative. A guardrail such as MCPJWTSigner, static_headers, or any injected Authorization could shadow the exchanged token, so the upstream received the injected JWT instead of the minted token and rejected it (no warning fired). For token_exchange and authorization_code, _create_mcp_client now drops the conflicting header and keeps the resolved token. No change for none / passthrough / static modes, where an injected Authorization still wins. This makes the signer-plus-OBO combination coherent instead of silently broken; the two are alternative identity-propagation strategies (LiteLLM-issued vs IdP-exchanged) that should never both reach one upstream

d. The OpenAPI path no longer leaks the raw subject token. The OpenAPI / local _request_extra_headers forwarder gated its strip on has_client_credentials only, so an OpenAPI-backed token_exchange server with extra_headers: [Authorization] forwarded the raw subject token upstream and never exchanged. It now uses the centralized _should_strip_caller_authorization, matching the managed paths

e. RFC 9728 challenge on unauthorized. OBO returned an opaque 401 Bearer error="invalid_request" with no discovery info, and any IdP exchange failure collapsed to a retryable 503. An OBO server now behaves like a standards-compliant OAuth resource server:

HTTP/1.1 401 Unauthorized
WWW-Authenticate: Bearer resource_metadata="/.well-known/oauth-protected-resource/mcp/<server>",
                  error="invalid_token",
                  error_description="Missing or invalid subject token; authenticate with the IdP and retry"

The challenge is emitted preemptively at connect (in the per-server preemptive-401), because the OBO tools are not discoverable without a subject and a tool-call-time 401 is wrapped into a JSON-RPC error so the WWW-Authenticate would be lost. The protected-resource metadata for an OBO server advertises the JWT-auth issuer(s) (JWT_ISSUER / litellm_jwtauth.issuers) as authorization_servers, which is the same IdP that issues and validates the subject token. So a spec-compliant MCP client discovers the IdP, SSOs, and retries with a fresh subject token, which LiteLLM then exchanges. An IdP 4xx (subject rejected) is a non-retryable 401 (the challenge) so a caller with a dead token re-authenticates instead of looping; 5xx and transport failures stay a retryable 503

3. Pure challenge edge (follow-up refactor)

The two challenge builders in the outbound_credentials adapter read SERVER_ROOT_PATH ambiently via get_server_root_path(), a hidden environment read inside a module that is supposed to be a pure edge. That coupling made the preemptive-challenge test order-dependent under xdist (a sibling test sets SERVER_ROOT_PATH at import without cleanup, leaking the prefix into the challenge URL). The root path is now resolved at the imperative-shell call sites and injected keyword-only, so raise_user_oauth_challenge and raise_token_exchange_challenge are pure functions of their inputs; the duplicated resource_metadata path construction collapses into one oauth_protected_resource_path helper. The adapter tests pass root_path as a real value instead of monkeypatching the environment, and the stale-session preemptive test asserts structural invariants rather than the exact prefixed URL so a leaked env var can no longer break it.

Two adjacent cleanups rode along: _create_mcp_client dropped back under the strict complexity ceiling by extracting the v2 credential resolution into _resolve_v2_auth, and the OBO protected-resource-metadata branch (which had shipped without coverage) moved into _obo_protected_resource_response with five new tests covering the issuer branch end to end

4. Contract-audit follow-ups (cache identity, retry, logging)

A pass over the OBO behavior contract surfaced three more gaps that land here. The exchanged-token cache key now folds in the caller's tenant alongside the subject token and exchange config, so two tenants presenting the same opaque token can never collide on one cache entry; isolation is structural rather than a side effect of subject-token uniqueness. The tool-call path gained a single reactive retry: when an upstream rejects the injected token with a 401/403, the gateway invalidates the cached exchange, re-mints once through the IdP, and retries the call exactly once before surfacing the upstream error, so a token revoked or rotated upstream mid-TTL self-heals without an infinite loop. This is gated strictly to oauth2_token_exchange, so passthrough, authorization_code, client_credentials, api_key, and none keep their existing single-call behavior. The exchanger also now emits the v1-parity log lines it had dropped (an attempt line with server, endpoint, and audience, a success line, and a cache-hit line), while still never logging the form, subject token, secret, or minted token.

One contract item is deliberately left for later. The exchange does not capture or use an RFC 8693 refresh_token: every expiry re-runs the exchange rather than refreshing, matching v1 and avoiding speculative refresh plumbing for IdPs that mostly do not issue OBO refresh tokens. The cache plus single-flight already keep the IdP round-trips off the hot path, so re-exchange-on-expiry is the intended lifetime here. RFC 9728 to RFC 8414 discovery of the token endpoint (so the endpoint need not be configured) is tracked as a separate follow-up rather than bundled here

Token endpoint resolution is now strictly config-driven and fails closed. A token_exchange server uses only an explicitly configured token_exchange_endpoint/token_url; the gateway never guesses an IdP or falls back to a weaker source. An OBO server that carries client credentials but no endpoint is owned by the v2 arm rather than deferred to v1, so a missing endpoint returns a clean 412 (precondition) before any upstream or IdP call and the caller's subject token is never sent anywhere. The no-subject case keeps its 401 RFC 9728 challenge, a rejected subject stays 401, and an unreachable IdP stays 503.

5. List-time 401 no longer masked (contract follow-up)

The list path still had one masking gap. When the resolver raised its 401 challenge while building the client (a present-but-rejected subject token, or an authorization_code server with no stored token), _get_tools_from_server caught it in the broad except Exception: return [] and the server showed zero tools with no challenge, so a client could not tell "no tools" from "authenticate and retry". The preemptive connect-time challenge only covers the no-subject case; a subject the IdP rejects fails later, at exchange time, inside that try

The fix routes a v2 resolver auth failure through the same MCPUpstreamAuthError channel pass-through already uses. _get_tools_from_server now converts a 401/403 HTTPException into MCPUpstreamAuthError, preserving the WWW-Authenticate header, so a single-server route surfaces the challenge (the client re-authenticates) while the multi-server aggregator keeps absorbing it into an empty contribution for that one server. A non-auth error (412 no endpoint, 503 IdP down) stays absorbed as before, so one misconfigured server cannot blank the whole catalog. This is mode-agnostic, so authorization_code list discovery gets the same un-masking for free

Alongside it, the exchanger now logs a warning when it refuses a non-Bearer token_type from the IdP (the fail-closed check itself lands in #31526; this adds the observability line now that this branch carries the exchanger's logging)

Auth model recap

A caller can present credentials two ways. Canonical single-JWT: Authorization: Bearer <IdP JWT> with no x-litellm-api-key, where the same JWT authenticates to LiteLLM via JWT auth and is the subject that gets exchanged (requires JWT auth configured against that IdP). Two-header: x-litellm-api-key: Bearer <litellm-key> for LiteLLM auth plus Authorization: Bearer <subject> for the exchange. In both, the subject always comes from Authorization, never from x-litellm-api-key, and only the exchanged token reaches the upstream

Screenshots / Proof of Fix

Verified end to end against a real Keycloak IdP (standing in for Okta), not a mock exchanger. The setup is Keycloak 26.2 with standard RFC 8693 token exchange enabled (realm litellm, user alice, a confidential litellm-exchange client the gateway authenticates the exchange with, and upstream-api as the exchange audience), a dedicated Postgres, a mock upstream MCP server that logs every Authorization it receives, and the proxy from this branch run with --use_v2_migration_resolver, JWT auth enabled, and the enterprise license.

The flow proven is exactly the OBO contract: SSO with Keycloak to obtain an <IdP JWT>, send it to LiteLLM, LiteLLM exchanges it at Keycloak's token endpoint for an upstream-scoped access token, and that minted token (never the JWT) is what reaches the upstream. The subject JWT and the token the upstream actually received, decoded side by side:

subject  (from Keycloak): azp=mcp-subject     aud=litellm-exchange  jti=onrtro:9266...
upstream (forwarded)    : azp=litellm-exchange aud=upstream-api      jti=ntrtte:b034...

It is a different token, re-scoped to the upstream audience, and the raw subject string never appears in the upstream's request log.

A full battery covering both credential shapes and every fail-closed path:

1. SSO: mint alice's IdP JWT from Keycloak (password grant)        PASS
2. Two-header flow (x-litellm-api-key + Authorization=subject)
   - tools listed, tool call returned the upstream result          PASS
   - upstream got the EXCHANGED token (jti != subject)             PASS
   - exchanged token is audience-scoped to upstream-api            PASS
   - raw subject NEVER forwarded upstream                          PASS
3. Single-JWT flow (only Authorization=<IdP JWT>, no api-key)
   - JWT authenticated to LiteLLM AND the call succeeded           PASS
   - the same JWT was exchanged; upstream got the minted token     PASS
   - raw subject JWT NEVER forwarded upstream                      PASS
4. No subject -> 401 + RFC 9728 resource_metadata challenge        PASS
5. token_exchange server with no endpoint -> 412 precondition      PASS
6. Invalid subject -> IdP 4xx -> non-retryable 401 (not 503)       PASS
7. Repeated calls reuse one exchanged token (cache, one IdP hit)   PASS

RESULT: 15 passed, 0 failed

The no-subject challenge, verbatim:

POST /mcp/obo_demo   (x-litellm-api-key only, no Authorization)
HTTP/1.1 401 Unauthorized
www-authenticate: Bearer resource_metadata="/.well-known/oauth-protected-resource/mcp/obo_demo",
                  error="invalid_token", error_description="Missing or invalid subject token; authenticate with the IdP and retry"

The two fail-closed call paths surface in the MCP tool result: a token_exchange server with no configured endpoint returns precondition required: token exchange endpoint is not configured for this server (412, with no IdP guessing and no fallback to a weaker source), and an invalid subject that Keycloak rejects with a 4xx returns Unauthorized (a non-retryable 401, not a retryable 503). Each behavior above also ships with unit regression tests that fail on the pre-fix code

Type

New Feature

🐛 Bug Fix

Changes

See the numbered sections above. Touched: mcp_server_manager.py, server.py, discoverable_endpoints.py, and the outbound_credentials adapter / exchanger / provider, with tests across test_mcp_server_manager.py, test_adapter.py, test_token_exchanger.py, test_resolver.py, test_discoverable_endpoints.py, test_mcp_stale_session.py, and the experimental_mcp_client client tests

Deliberately not implemented

The single JWT (SSO) shape, where the caller's IdP JWT both authenticates to LiteLLM and serves as the exchange subject, is the intended mode. The two-header shape (LiteLLM api key plus a separate Authorization subject) remains reachable but is not advertised, and in that shape the gateway does not pre-validate the subject's signature or expiry before the exchange; the IdP is the authority that accepts or rejects it. Delegation via actor_token (RFC 8693 section 4.1) is intentionally absent; the arm implements impersonation only

[tin-berri]

@greptileai

[greptile-apps]

Greptile Summary

This PR makes the OBO / token_exchange MCP auth mode production-ready by threading the caller's subject token through discovery (tools/list, prompts, resources), hardening the credential resolution path against bypasses and injection conflicts, adding an RFC 9728 challenge for missing/rejected subject tokens, and improving the exchange's reliability contract (tenant-isolated cache keys, a single reactive retry on upstream 401, and proper IdP error classification).

• Discovery threading: _get_tools_from_server now extracts and forwards the caller's bearer for oauth2_token_exchange servers only, unmasking the earlier "zero tools visible" bug; prompts, resources, and resource-templates get the same treatment via _obo_subject_token.
• Credential hardening: _create_mcp_client now blocks x-mcp-* caller overrides for TokenExchangeConfig (not just AuthorizationCodeConfig), _resolve_v2_auth drops a conflicting Authorization from extra_headers so a signer/static-header cannot shadow the resolver-owned token, and the OpenAPI forwarder now uses the centralized _should_strip_caller_authorization instead of its own has_client_credentials guard.
• Exchange robustness: IdP 4xx responses are classified into caller-fault (SubjectTokenRejected → 401 OBO challenge) vs gateway-fault (TokenExchangeClientError → 500), cache keys fold in tenant_id to prevent cross-tenant collisions, an invalidate path plus one reactive retry handle mid-TTL token rotation, and _ttl_seconds is corrected to never cache a token beyond its own remaining lifetime.

Confidence Score: 4/5

Safe to merge with the caveat that OBO tool calls bypass the per-server concurrency semaphore, which matters whenever max_concurrency is configured for a token_exchange server.

The OBO retry path (_obo_call_tool_with_retry) does not acquire the per-server concurrency semaphore that all non-OBO tool calls hold, so any max_concurrency limit configured on an OBO server is silently ignored for every OBO tool call. All other changes — credential hardening, challenge routing, cache isolation, IdP error classification, and the _ttl_seconds correction — look correct and are well-covered by the new tests.

litellm/proxy/_experimental/mcp_server/mcp_server_manager.py — the _obo_call_tool_with_retry method and its call site need the concurrency semaphore applied.

Important Files Changed

Filename	Overview
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py	Large refactor for OBO discovery threading, retry-on-401, and challenge routing; the OBO retry path (_obo_call_tool_with_retry) does not acquire the per-server concurrency semaphore that the non-OBO path uses.
litellm/proxy/_experimental/mcp_server/outbound_credentials/token_exchanger.py	Adds SubjectTokenRejected/TokenExchangeClientError exception types, tenant isolation in cache key, invalidate, logging, float expires_in handling, and a corrected _ttl_seconds that floors at the token's own remaining lifetime.
litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py	Extracts oauth_protected_resource_path helper, makes raise_user_oauth_challenge and new raise_token_exchange_challenge pure (root_path injected), and widens _token_exchange_spec to own servers with credentials but no endpoint.
litellm/proxy/_experimental/mcp_server/outbound_credentials/token_exchange_provider.py	Adds SubjectTokenRejected / TokenExchangeClientError classification for IdP 4xx responses; gateway-fault OAuth error codes (invalid_client, etc.) map to misconfigured, caller fault maps to unauthorized.
litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py	Adds _obo_protected_resource_response and _jwt_auth_issuers for OBO PRM, and fixes the _build_oauth_protected_resource_response order so OBO servers get their own metadata before falling through to the gateway default.
litellm/proxy/_experimental/mcp_server/server.py	Forwards oauth2_headers to _get_tools_from_server, uses centralized _should_strip_caller_authorization for OpenAPI path, and adds preemptive RFC 9728 challenge for OBO servers with no subject token at connect time.
litellm/experimental_mcp_client/client.py	Adds raise_on_error parameter to call_tool and extracts error_tool_result static helper to avoid code duplication.
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py	New tests for OBO discovery threading, 401 surfacing, token-thread mode gating, and retry behavior; existing tests refactored for readability only.
tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_exchanger.py	New tests for IdP rejection → unauthorized, gateway fault → misconfigured, transport failure → 503, tenant isolation in cache, invalidate targeting, and logging of bad token_type.
tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py	Tests refactored to inject root_path directly instead of monkeypatching; new tests for oauth_protected_resource_path and raise_token_exchange_challenge.
tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py	Mostly line-length reformatting of existing tests; no changes to assertions or mock behavior.
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py	Reformatted assertions plus a new test for the preemptive OBO challenge at connect time.
tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py	Adds invalidate_credentials tests and updates _FakeExchanger to track tenant_id in call records.
tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_exchange_provider.py	New file with tests for IdP 4xx classification (SubjectTokenRejected vs TokenExchangeClientError vs None for 5xx).

Comments Outside Diff (1)

1. litellm/proxy/_experimental/mcp_server/mcp_server_manager.py, line 3368-3392 (link)

   OBO tool calls bypass the per-server concurrency semaphore

   The non-OBO branch wraps client.call_tool inside async with self._limit_outbound_concurrency(mcp_server), but _obo_call_tool_with_retry (both the initial call and the single retry) never acquires that semaphore. If an operator has set max_concurrency on an OBO server, concurrent OBO tool calls are not rate-limited and can flood the upstream, defeating the guard entirely.

_{Reviews (3): Last reviewed commit: "style(mcp): PEP 604 union in the OBO ret..." | Re-trigger Greptile}

[greptile-apps]

Greptile Summary

This PR threads caller OAuth context into MCP tool discovery. The main changes are:

• Pass oauth2_headers from the MCP list-tools aggregator into per-server discovery.
• Extract the inbound bearer token for oauth2_token_exchange servers and pass it as subject_token to MCP client creation.
• Keep other auth modes from receiving the caller bearer token.
• Add tests for token-exchange discovery, non-OBO no-leak behavior, and background refresh without headers.

Confidence Score: 5/5

The changes are narrow, scoped to OAuth token-exchange discovery, and covered by targeted tests for the new threading behavior and non-leak cases.

The implementation matches the described call-path behavior, preserves existing graceful degradation for listing, and includes tests for the relevant auth-mode boundaries.

T-Rex Logs

What T-Rex did

• Ran a focused runtime test exercising MCPServerManager._get_tools_from_server and observed the head state increased tool_count to 1 with tool_names including 'obo_tool' and the overall validation passed.
• Validated that the head now supports oauth2_headers and that non-OBO scenarios completed with HTTP 200 OK while still capturing subject_token=[REDACTED].
• Verified the new in-process contract boundary by exercising MCPServerManager._get_tools_from_server with optional oauth2_headers=None, observed a simulated failure returning [] with exit code 0, and noted that no HTTP endpoints were started for this scoped validation.

_{Ran code and verified through T-Rex}

_{Reviews (1): Last reviewed commit: "feat(mcp): thread the caller token into ..." | Re-trigger Greptile}

[codecov]

Codecov Report

❌ Patch coverage is 84.21053% with 24 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
..._experimental/mcp_server/discoverable_endpoints.py	50.00%	13 Missing ⚠️
...er/outbound_credentials/token_exchange_provider.py	81.81%	4 Missing ⚠️
litellm/experimental_mcp_client/client.py	50.00%	3 Missing ⚠️
...oxy/_experimental/mcp_server/mcp_server_manager.py	96.15%	2 Missing ⚠️
...mental/mcp_server/outbound_credentials/resolver.py	83.33%	1 Missing ⚠️
litellm/proxy/_experimental/mcp_server/server.py	80.00%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

[mateo-berri]

Is the veria concern legit?

[mateo-berri]

@greptileai

[mateo-berri]

The non-OBO branch wraps client.call_tool inside async with self._limit_outbound_concurrency(mcp_server), but _obo_call_tool_with_retry (both the initial call and the single retry) never acquires that semaphore. If an operator has set max_concurrency on an OBO server, concurrent OBO tool calls are not rate-limited and can flood the upstream, defeating the guard entirely.

Is this a legit concern?

[tin-berri]

The non-OBO branch wraps client.call_tool inside async with self._limit_outbound_concurrency(mcp_server), but _obo_call_tool_with_retry (both the initial call and the single retry) never acquires that semaphore. If an operator has set max_concurrency on an OBO server, concurrent OBO tool calls are not rate-limited and can flood the upstream, defeating the guard entirely.

Is this a legit concern?

no

[mateo-berri]

LGTM; thanks!