fix(mcp): preserve upstream OAuth 401 challenge for stored per-user MCP tokens - #31362
fix(mcp): preserve upstream OAuth 401 challenge for stored per-user MCP tokens#31362mateo-berri wants to merge 9 commits into
Conversation
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
Greptile SummaryFixes the gateway-managed per-user OAuth2 MCP path where a stored token present in the DB but revoked upstream caused the proxy to silently complete the handshake and return an empty tool list instead of surfacing the upstream
Confidence Score: 5/5Safe to merge; the new preflight probe is best-effort, gated on the already-authorized server list, and guarded by ten focused unit tests. The change is well-scoped: probes only fire for authorized servers, asyncio.gather keeps latency bounded, the stored DB token correctly takes precedence over the client header, and network errors fall through gracefully. No auth boundary regressions were found. No files require special attention.
|
| Filename | Overview |
|---|---|
| litellm/proxy/_experimental/mcp_server/server.py | Adds _get_gateway_oauth2_challenge helper, refactors _check_passthrough_upstream_auth to accept pre-resolved server list, and introduces _check_oauth2_upstream_auth that probes per-user OAuth2 servers in parallel; both transports share a single _get_allowed_mcp_servers lookup. Logic is correct: DB token takes precedence over client header, unauthorized probes are impossible, and challenge fallbacks for delegate/non-delegate servers are handled properly. |
| litellm/proxy/_experimental/mcp_server/mcp_server_manager.py | Adds resolve_user_oauth_authorization_header: routes through the v2 resolver for AuthorizationCodeConfig servers and falls back to the v1 DB lookup for delegate/unmigrated servers. Missing-token maps to None (never raises). Lazy circular import is an acceptable workaround. |
| litellm/proxy/_experimental/mcp_server/outbound_credentials/httpx_auth.py | Adds header_value() to StaticHeaderAuth to expose the credential outside the httpx auth flow while keeping it wrapped in SecretStr. |
| tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py | Adds TestCheckOauth2UpstreamAuth with 10 mocked scenarios covering all significant branches of the new preflight. |
| tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py | Adds three unit tests for resolve_user_oauth_authorization_header covering v2 resolver, missing-token-to-None, and v1 fallback paths. |
| tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_httpx_auth.py | Adds test confirming header_value() returns the raw credential string while it remains absent from repr. |
Reviews (8): Last reviewed commit: "refactor(mcp): resolve the oauth2 prefli..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Greptile SummaryThis PR fixes a gap in the managed-OAuth2 MCP path: when a stored per-user token is stale or revoked, the previous code only surfaced the 401 challenge for the missing-token case. The upstream
Confidence Score: 4/5Safe to merge for correctness; the new pre-flight probe is fail-open and only raises on a genuine upstream 401. Performance degrades linearly with the number of per-user OAuth2 servers a caller is authorized for. The fix is logically correct and the test coverage is thorough. The one concern is that _check_oauth2_upstream_auth probes upstream servers sequentially (one await per server) while the analogous _check_passthrough_upstream_auth gathers probes in parallel. For most deployments (1–2 servers) this difference is negligible, but at scale it turns a single probe round-trip into N serial round-trips on every request, each with a 5-second timeout ceiling. litellm/proxy/_experimental/mcp_server/server.py — specifically the for-loop probe structure in _check_oauth2_upstream_auth.
|
| Filename | Overview |
|---|---|
| litellm/proxy/_experimental/mcp_server/server.py | Adds _get_gateway_oauth2_challenge helper (refactors inline code) and _check_oauth2_upstream_auth (new pre-flight probe for stale per-user OAuth tokens). The new probe is invoked synchronously in the hot path but probes servers sequentially instead of in parallel like _check_passthrough_upstream_auth. |
| tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py | Adds 8 focused mock-only unit tests for _check_oauth2_upstream_auth and the handler integration. Covers stale token, missing header, DB-token preference, M2M skip, no-token-no-probe, and fallback challenge cases. No real network calls. |
Reviews (2): Last reviewed commit: "chore(lint): use builtin generics in _ch..." | Re-trigger Greptile
_check_oauth2_upstream_auth probed each per-user OAuth2 upstream sequentially, so N authorized servers cost N serial round-trips (each with a 5s timeout) on every request. Mirror _check_passthrough_upstream_auth and fan the per-server stored-token lookup plus initialize probe out via asyncio.gather, then surface the first genuine upstream 401 challenge in deterministic server order. Fail-open semantics are preserved: the probe and DB-token helpers each swallow transient errors internally, so a hiccup never blocks a valid request.
Directly exercise the primary fix scenario: a present-but-revoked per-user token stored in the DB (not the client Authorization header) must surface the upstream 401 challenge. The existing stale-token test patched the DB helper to None, so it only covered the client-header fallback; the DB-token precedence branch was only verified to fall through on 200. This adds the missing 401 sub-path so the stored-token revocation case can never regress.
|
Generated by Claude Code |
|
bugbot run Generated by Claude Code |
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using high effort and found 3 potential issues.
Autofix Details
Bugbot Autofix prepared fixes for all 3 issues found in the latest run.
- ✅ Fixed: Wrong delegate OAuth fallback challenge
- When an upstream 401 omits WWW-Authenticate, _check_oauth2_upstream_auth now emits the proxied resource_metadata challenge for delegate_auth_to_upstream servers and only falls back to the gateway authorization_uri for gateway-managed OAuth2 servers.
- ✅ Fixed: SSE handler skips OAuth probe
- handle_sse_mcp now invokes _check_oauth2_upstream_auth after the passthrough probe so stale stored per-user OAuth tokens surface as a 401 challenge on /sse instead of degrading to empty tool lists.
- ✅ Fixed: OAuth2 probe ignores upstream 403
- _check_oauth2_upstream_auth now raises HTTPException(403, 'Forbidden') on an upstream 403 mirroring _check_passthrough_upstream_auth, so a valid-but-unauthorized per-user OAuth2 token no longer silently degrades into an empty tool list.
You can send follow-ups to the cloud agent here.
- _check_oauth2_upstream_auth now offers the proxied resource_metadata challenge for delegate_auth_to_upstream servers when the upstream 401 omits WWW-Authenticate, instead of pointing clients at the gateway authorization server. - _check_oauth2_upstream_auth now raises 403 on an upstream forbidden, matching the passthrough preflight so a valid-but-unauthorized stored token does not silently degrade to an empty tool list. - handle_sse_mcp now runs _check_oauth2_upstream_auth so stale stored per-user OAuth tokens surface a 401 challenge on the /sse transport the same way they do on streamable HTTP.
|
|
Add regression coverage for the three bugbot fixes in fa76b92, and reflow one over-wrapped call in _check_oauth2_upstream_auth so the file stays ruff-format clean. - test_upstream_403_raises_forbidden: an upstream probe 403 raises HTTPException(403) without a re-auth challenge, matching the passthrough probe (a fresh token with the same scopes would loop). - test_delegate_auth_401_falls_back_to_passthrough_challenge: a delegate_auth_to_upstream server that 401s without a WWW-Authenticate header gets the proxied resource_metadata challenge, not the gateway authorization_uri, so the client re-authorizes against the upstream IdP. - test_sse_handler_propagates_stale_token_401: the /sse handler now runs the per-user OAuth2 preflight too, so a stale stored token surfaces the upstream 401 instead of degrading into an empty tool list. Each test fails if its branch is reverted (mutation-checked).
|
Generated by Claude Code |
|
bugbot run Generated by Claude Code |
There was a problem hiding this comment.
✅ Bugbot reviewed your changes and found no new issues!
Comment @cursor review or bugbot run to trigger another review on this PR
Reviewed by Cursor Bugbot for commit 0d709a0. Configure here.
|
Generated by Claude Code |
|
Generated by Claude Code |
|
bugbot run Generated by Claude Code |
There was a problem hiding this comment.
✅ Bugbot reviewed your changes and found no new issues!
Comment @cursor review or bugbot run to trigger another review on this PR
Reviewed by Cursor Bugbot for commit 3479ec8. Configure here.
…2 resolver The probe in _check_oauth2_upstream_auth previously read the stored per-user token via the v1 helper (resolve_user_oauth_access_token), a parallel read of the same DB rows the v2 resolver owns for authorization_code servers. The new manager seam resolve_user_oauth_authorization_header routes the probe through the v2 resolver (same store, cache, and refresh chain as the egress) for servers it owns and falls back to the v1 lookup for servers it does not (delegate/BYOK), so the probe can never test a token the call path would not send
|
Generated by Claude Code |
|
bugbot run Generated by Claude Code |
There was a problem hiding this comment.
✅ Bugbot reviewed your changes and found no new issues!
Comment @cursor review or bugbot run to trigger another review on this PR
Reviewed by Cursor Bugbot for commit e9dbf6d. Configure here.
|
Addressed in PR #32302 |
Pull request was closed
![veria-ai[bot]](https://avatars.githubusercontent.com/in/1723776?s=60&v=4){kind=small}
| server=server, | ||
| user_api_key_auth=user_api_key_auth, | ||
| ) | ||
| auth_header = stored_auth or (oauth2_headers or {}).get("Authorization") |
There was a problem hiding this comment.
High: Proxy API key forwarded to OAuth upstream
When clients authenticate normally through Authorization, process_mcp_request treats this value as their LiteLLM API key and also places it in oauth2_headers. If the resolver returns no stored upstream token, this fallback sends that proxy credential to server.url, allowing the upstream server to capture and reuse it. Keep the proxy admission credential separate from upstream OAuth credentials, and only use a caller-provided bearer here when the request authenticated through a distinct x-litellm-api-key header or the server explicitly delegates authentication upstream.
PR overviewThis PR updates MCP OAuth handling so per-user MCP token flows preserve and propagate the upstream OAuth 401 challenge behavior. The touched MCP server request path manages how client authentication and upstream OAuth headers are prepared when proxying MCP requests. One security issue remains open: the MCP proxy path can still reuse a client's LiteLLM Open issues (1)
Fixed/addressed: 0 · PR risk: 7/10 |
4 participants
Relevant issues
N/A
Linear ticket
N/A
Pre-Submission checklist
make test-unit@greptileaiand received a Confidence Score of at least 4/5Screenshots / Proof of Fix
Reproduced on latest
litellm_internal_staging: when a stored per-user OAuth token is stale or invalid, the managed-oauth2 gate only handled the missing-token case, so the upstream MCP401challenge was swallowed instead of being surfaced to the client. Upstream PR #27847 only covers the token-forwarding case (extra_headers: [Authorization]); this change also preserves the challenge for the per-user stored OAuth case (extra_headers: {})This run supersedes the earlier proof that used a local stub upstream. The upstream here is Linear's public MCP server, which emits a genuine RFC 6750 challenge for invalid tokens; a direct curl with a garbage token shows exactly what the gateway must preserve:
Live before/after on a local Postgres-backed proxy (fresh database, real migrations applied via
--use_v2_migration_resolver). Before ran the merge-base commit7f991481cc069d7a069a8a50c140bcfaec9a4e6con port 28641, after ran this PR's head3479ec82ce3cf1205fb0596b774b6a17a90220ceon port 29473, both withpython litellm/proxy/proxy_cli.py --config qa_config.yaml --port <port> --detailed_debug --use_v2_migration_resolverand this config:Setup:
curl http://localhost:<port>/user/newfor userqa31362-mcp-user,curl http://localhost:<port>/key/generatewith"object_permission": {"mcp_servers": ["9b493abaad9dab94382265614863a761"]}, then a stored per-user OAuth credential for that user and server written viastore_user_oauth_credential(access_token="qa31362-revoked-token", expires_in=3600), i.e. a token the gateway still considers valid but the upstream rejectsClient request, identical in both runs:
Before (
7f991481cc, port 28641): the handshake succeeds and the session silently degrades into an empty tool list, with nowww-authenticateanywhere, even though the proxy log shows Linear rejected the stored tokenthen
tools/liston that session:proxy log during that
tools/list:After (
3479ec82ce, port 29473): the same initialize request surfaces Linear's challenge, byte-identical to the direct-to-upstream response aboveType
🐛 Bug Fix
Changes
server.pysurfaces the upstream OAuth401challenge for the gateway-managed per-user stored-token path on both the Streamable HTTP and/ssetransports, so a stale or revoked token returns a re-authorization challenge instead of silently degrading into an empty tool list. The pre-flight probe mirrors_check_passthrough_upstream_auth: it fans the per-server DB-token lookup andinitializeprobe out withasyncio.gather, prefers the stored DB token over the caller'sAuthorizationheader, skips M2M (client_credentials) servers, raises403on a forbidden upstream, and falls back to the proxiedresource_metadatachallenge fordelegate_auth_to_upstreamservers rather than the gatewayauthorization_uri. Both handlers now resolve_get_allowed_mcp_serversonce per request and inject the result into_check_passthrough_upstream_authand_check_oauth2_upstream_auth, so the new probe adds no additional authorized-servers lookup on the hot path. Mutation-checked regression coverage lives intest_mcp_server.py. Part of upstreaming a downstream patch series maintained against 1.85.1Note
Medium Risk
Touches MCP transport auth preflight and OAuth token resolution on the connect path; behavior is scoped to authorized servers and heavily tested, but wrong challenge mapping could break client re-auth flows.
Overview
Fixes gateway-managed per-user OAuth2 MCP connections so a stale or revoked stored token returns an HTTP 401 with a proper
WWW-Authenticatechallenge instead of opening a session that later shows an empty tool list.Adds
_check_oauth2_upstream_authon Streamable HTTP and/sse, mirroring the existing pass-through preflight: it probes authorized upstreams in parallel with the credential the egress would use, prefers the stored DB token over the clientAuthorizationheader, skips M2M servers, maps upstream 403 to Forbidden (no re-auth hint), and chooses delegateresource_metadatavs gatewayauthorization_uriwhen the upstream omits a challenge.MCPServerManager.resolve_user_oauth_authorization_headeraligns the probe with the v2 credential resolver (v1 DB fallback when unmigrated).StaticHeaderAuth.header_value()exposes the secret for hand-built probe requests._get_gateway_oauth2_challengecentralizes gateway discovery challenges; both transports share one_get_allowed_mcp_serversresult for pass-through and OAuth2 prefights.Reviewed by Cursor Bugbot for commit e9dbf6d. Bugbot is set up for automated code reviews on this repo. Configure here.