Skip to content

fix(mcp): forward short OAuth state upstream, keep session in a cookie - #32146

Merged
tin-berri merged 2 commits into
litellm_internal_stagingfrom
litellm_mcp_short_oauth_state
Jul 6, 2026
Merged

fix(mcp): forward short OAuth state upstream, keep session in a cookie#32146
tin-berri merged 2 commits into
litellm_internal_stagingfrom
litellm_mcp_short_oauth_state

Conversation

@tin-berri

@tin-berri tin-berri commented Jul 4, 2026

Copy link
Copy Markdown
Contributor

Relevant issues

Linear ticket

Resolves LIT-4197

Pre-Submission checklist

Please complete all items before asking a LiteLLM maintainer to review your PR

  • I have added meaningful tests
  • My PR passes all CI/CD checks (e.g., lint, format, unit tests)
  • My PR's scope is as isolated as possible; it only solves 1 specific problem
  • I have requested a Greptile review by commenting @greptileai and received a Confidence Score of at least 4/5 before requesting a maintainer review

Screenshots / Proof of Fix

Live proxy on localhost:4000 with one auth_type: oauth2 MCP server whose authorization_url points at an upstream IdP. The request below is the exact interactive /authorize an MCP client (mcp-inspector, Claude Code, Cursor, the LiteLLM UI connect button) makes, with a native loopback redirect_uri.

Before the fix, LiteLLM forwards its own long encrypted session as the upstream state:

$ curl -sS -D - -o /dev/null -G "http://localhost:4000/test_oauth_mcp/authorize" \
    --data-urlencode "response_type=code" \
    --data-urlencode "client_id=upstream-client-id-123" \
    --data-urlencode "code_challenge=gtt1Df7li2G3EF8y3Yzh4oP7SxKyTtKcbgzrA5qqpCs" \
    --data-urlencode "code_challenge_method=S256" \
    --data-urlencode "redirect_uri=http://127.0.0.1:6274/oauth/callback/debug" \
    --data-urlencode "state=ee230e3dfd4f19c7441941684f39c8a4e0e2c3c61a088e33403df5662b4047b8" \
    --data-urlencode "scope=mcp"

HTTP/1.1 307 Temporary Redirect
location: https://idp.example.com/oauth/authorize?client_id=upstream-client-id-123&redirect_uri=http%3A%2F%2Flocalhost%3A4000%2Fcallback&state=aytF9XRabR8VARNc...(468 chars)...&response_type=code&scope=mcp&code_challenge=gtt1Df7li2G3EF8y3Yzh4oP7SxKyTtKcbgzrA5qqpCs&code_challenge_method=S256

# upstream `state` is 468 chars and there is no Set-Cookie; strict IdPs reject it as "state parameter too long"

After the fix, the same request forwards a short handle upstream and moves the encrypted session into a per-flow HttpOnly cookie:

$ curl ... (identical /authorize request)

HTTP/1.1 307 Temporary Redirect
location: https://idp.example.com/oauth/authorize?client_id=upstream-client-id-123&redirect_uri=http%3A%2F%2Flocalhost%3A4000%2Fcallback&state=9bXVl-kzaK_5Z4C1...(43 chars)...&response_type=code&scope=mcp&code_challenge=gtt1Df7li2G3EF8y3Yzh4oP7SxKyTtKcbgzrA5qqpCs&code_challenge_method=S256
set-cookie: mcp_oauth_state_9bXVl-kzaK_5Z4C1...=oMom-SmuBH-7aK5W...(encrypted session)...; HttpOnly; Max-Age=600; Path=/; SameSite=lax

# upstream `state` is now 43 chars; `code_challenge` is still forwarded (PKCE stays a transparent pass-through)

The IdP redirects back to /callback with the handle; the browser replays the cookie, so LiteLLM recovers the session and returns the client's own original state plus the code, then clears the one-time cookie:

$ curl -sS -D - -o /dev/null -H "Cookie: mcp_oauth_state_9bXVl-kzaK_5Z4C1...=oMom-SmuBH-7aK5W..." \
    "http://localhost:4000/callback?code=upstream-auth-code-xyz&state=9bXVl-kzaK_5Z4C1..."

HTTP/1.1 302 Found
location: http://127.0.0.1:6274/oauth/callback/debug?code=upstream-auth-code-xyz&state=ee230e3dfd4f19c7441941684f39c8a4e0e2c3c61a088e33403df5662b4047b8
set-cookie: mcp_oauth_state_9bXVl-kzaK_5Z4C1...=""; expires=...; Max-Age=0; Path=/; SameSite=lax

# the client gets its OWN 64-char state back (not our handle), the auth code, and the cookie is deleted

Flows started before the change (or in flight across a rolling deploy) still carry the encrypted blob directly in state with no cookie; /callback falls back to decoding it, so they keep working:

$ curl -sS -D - -o /dev/null \
    "http://localhost:4000/callback?code=legacy-code-abc&state=<legacy 468-char encrypted state, no cookie>"

HTTP/1.1 302 Found
location: http://127.0.0.1:6274/oauth/callback/debug?code=legacy-code-abc&state=ee230e3dfd4f19c7441941684f39c8a4e0e2c3c61a088e33403df5662b4047b8

Type

🐛 Bug Fix

Changes

When LiteLLM proxies an interactive authorization_code OAuth flow to an upstream MCP authorization server, it needs to remember two things for /callback: the client's original state (to echo back so the client's CSRF check passes) and the client's redirect_uri (to know where to send the browser). The upstream only reflects the one state value it is given, so the old code packed the whole session (base_url, original state, PKCE fields, client redirect_uri) into an encrypted blob and sent that blob upstream as state. That blob runs a few hundred characters, and some authorization servers reject an over-long state, which is the "state parameter too long" failure

authorize_with_server now forwards a short random handle as the upstream state and stores the existing encrypted session in a per-flow HttpOnly, SameSite=lax cookie keyed by that handle. The browser carries the cookie across the upstream round trip, so /callback recovers the session with no server-side store, which keeps the flow correct across proxy replicas the same way the stateless-state design did. The handle is server-generated rather than the client's own state so it is unique per flow (no cross-flow cookie collisions) and unguessable, and the client still receives its own original state back at its redirect_uri because /callback restores it from the cookie

/callback reads the session from the cookie when present and falls back to decoding state directly otherwise, so states minted before this change keep working. PKCE is unchanged; code_challenge and code_verifier are still forwarded to the upstream, and they were never read back on /callback, so nothing depends on them surviving inside state

The change is confined to authorize_with_server and the shared /callback, so it covers both the discoverable /authorize (used by MCP clients) and the UI /server/oauth/{server_id}/authorize flow. Pass-through servers (auth_type none, upstream-delegated discovery), client_credentials, static-header auth, and the BYOK authorization-server endpoints do not route through these functions and are untouched


Note

Medium Risk
Touches OAuth authorize/callback and redirect/cookie handling on a security-sensitive path; behavior is backward-compatible but cookie loss or cross-site edge cases could break interactive MCP login.

Overview
Fixes LIT-4197: strict upstream IdPs were rejecting LiteLLM’s MCP OAuth proxy because the upstream state carried a long encrypted session blob.

/authorize (authorize_with_server) now sends a short random handle as upstream state and stores the same encrypted session in a per-flow HttpOnly, SameSite=lax cookie (mcp_oauth_state_<handle>, 10‑minute TTL). The browser carries that cookie through the IdP redirect, so no server-side session store is needed and multi-replica proxies still work.

/callback resolves the session via _resolve_encoded_oauth_state (cookie when present, otherwise the legacy value in state for in-flight or pre-deploy flows), then clears the one-time cookie on success, IdP error propagation, and failure paths.

Tests cover the full authorize→callback round trip, IdP error handling with cookie cleanup, and mock request cookies.

Reviewed by Cursor Bugbot for commit 444446e. Bugbot is set up for automated code reviews on this repo. Configure here.

@tin-berri

Copy link
Copy Markdown
Contributor Author

@greptileai

@greptile-apps

greptile-apps Bot commented Jul 4, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

Fixes LIT-4197: strict upstream IdPs were rejecting LiteLLM's MCP OAuth proxy because the state parameter carried a long (~468-char) encrypted session blob. The fix forwards a short 43-char random handle upstream and stores the encrypted session in a per-flow HttpOnly, SameSite=lax cookie (mcp_oauth_state_<handle>, 10-minute TTL), which the browser carries across the IdP redirect so no server-side session store is needed.

  • authorize_with_server now generates a secrets.token_urlsafe(32) relay handle as the upstream state, sets the encrypted session in a cookie, and returns the redirect — PKCE fields are still forwarded transparently.
  • /callback resolves the session via _resolve_encoded_oauth_state (cookie when present, legacy direct-state decoding when absent for backward compatibility), then clears the one-time cookie on success, IdP-error propagation, and generic exception paths.
  • Two new mock tests cover the full authorize→callback round trip and the IdP error + cookie-cleanup path.

Confidence Score: 5/5

Safe to merge — the change is isolated to the MCP OAuth proxy authorize/callback path, is backward-compatible for in-flight flows, and is covered by new mock tests that verify both the cookie round trip and error-path cookie cleanup.

The relay-state design is stateless and correct across replicas. Cookie attributes (HttpOnly, SameSite=lax, Max-Age=600, path/secure derived consistently for both set and delete) are appropriate for the cross-IdP-redirect use case. The fallback path for legacy encrypted states is safe. The one remaining gap — the one-time cookie is not cleared when an HTTPException is re-raised on the success path — is a minor UX leak (orphaned cookie expires in 600 s) that was already identified in prior review threads.

No files require special attention.

Important Files Changed

Filename Overview
litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py Core change: introduces cookie-based relay state for OAuth authorize/callback. Cookie helpers are correct (path/secure/httponly/samesite handled consistently between set and delete). Backward-compat fallback in _resolve_encoded_oauth_state works. Cookie is cleared on all callback return paths except the re-raised HTTPException (minor UX gap, previously flagged).
tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py Adds req.cookies={} to the shared mock fixture and two new async tests covering the full authorize→callback cookie round trip and the IdP error + cookie-cleanup path. No real network calls; all HTTP is mocked. Existing tests are unmodified.

Reviews (5): Last reviewed commit: "test(mcp): cover /callback error path co..." | Re-trigger Greptile

Comment thread litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py Outdated
@codecov

codecov Bot commented Jul 4, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

@greptile-apps

This comment was marked as outdated.

@tin-berri
tin-berri force-pushed the litellm_mcp_short_oauth_state branch from 71510f9 to 6ce6fd9 Compare July 4, 2026 20:04
@tin-berri

Copy link
Copy Markdown
Contributor Author

Addressed the review feedback in the latest commit

The /callback delete_cookie now mirrors the set_cookie attributes (path, secure, httponly, samesite) through a shared _oauth_state_cookie_path_and_secure helper, so the cookie is reliably expired on HTTPS deployments rather than left to age out

The one-time cookie is now also cleared on the /callback error-path and success-path exception branches, not just the happy path, so a failed decode or an invalid redirect no longer leaves an orphaned cookie

The regression test now asserts the callback response expires the cookie (Max-Age=0, empty value); a mutation that drops the clear fails the test

@greptileai

@tin-berri

Copy link
Copy Markdown
Contributor Author

@greptileai review

Some upstream authorization servers reject the OAuth authorize request with
"state parameter too long" because LiteLLM replaced the client's short state
with its own long encrypted session blob (base_url, original state, PKCE, client
redirect_uri) and sent that upstream as state.

Forward a short random handle as the upstream state instead, and carry the
encrypted session in a per-flow HttpOnly, SameSite=lax cookie bound to that
handle. The browser replays the cookie on /callback, so the session is recovered
without any server-side store and the client still gets its own original state
back. /callback falls back to decoding state directly when no cookie is present,
so flows in flight across a deploy keep working.

Resolves LIT-4197
@tin-berri
tin-berri force-pushed the litellm_mcp_short_oauth_state branch from 6ce6fd9 to 5c65f1c Compare July 4, 2026 20:32
The happy-path regression test already asserts the short-handle -> cookie round
trip. Add a focused test for the IdP-error branch of /callback: it must recover
the client's original state from the per-flow cookie (not the short handle),
propagate the error to the client's redirect_uri, and expire the one-time
cookie. Fails if the error path stops reading or clearing the cookie.
@tin-berri

Copy link
Copy Markdown
Contributor Author

bugbot run

@cursor cursor Bot left a comment

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.

Cursor Bugbot has reviewed your changes using high effort and found 1 potential issue.

Fix All in Cursor

❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.

Reviewed by Cursor Bugbot for commit 444446e. Configure here.

return RedirectResponse(url=complete_returned_url, status_code=302)
response = RedirectResponse(url=complete_returned_url, status_code=302)
_clear_oauth_state_cookie(response, request, state)
return response

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.

HTTPException skips cookie cleanup

Medium Severity

On the successful /callback path, when _get_validated_client_redirect_uri raises HTTPException, the handler re-raises without calling _clear_oauth_state_cookie. Other failure branches in the same handler clear the one-time mcp_oauth_state_* cookie, so the encrypted OAuth session can remain in the browser for the full Max-Age.

Fix in Cursor Fix in Web

Reviewed by Cursor Bugbot for commit 444446e. Configure here.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Accurate read of that branch, and it is intentional; the exposure is low enough that reworking it would cost more than it saves

The only thing that raises HTTPException inside that try is _get_validated_client_redirect_uri, so the cookie survives only when the decoded client_redirect_uri fails the sink-side VERIA-57 trust check. The re-raise is deliberate; it surfaces that as a 400 rather than the generic authentication-incomplete fallback, and it is pinned by the existing VERIA-57 regression tests, which assert pytest.raises(HTTPException) with status_code == 400. Converting the branch to a returned response so it can clear the cookie would break that contract and route around the proxy's HTTPException handler

The surviving cookie is also inert. It carries the encrypted {original_state, client_redirect_uri, code_challenge, ...} blob with no tokens and no authorization code; it is HttpOnly and Secure, and it expires within its 600s Max-Age. Any replay re-runs the same validation and re-fails with the same 400, and the code is only appended to the redirect after validation passes, so a surviving cookie cannot leak it to the untrusted URI. A retry mints a fresh handle and cookie and orphans the stale one, and for the loopback native-client flows this targets the branch never fires at all, since loopback validates identically at /authorize and /callback; it needs a same-origin UI redirect plus an origin shift between the two requests

Leaving it as intentional on that basis. If strict parity across every branch is wanted later, the safe way is to attach the delete-cookie to the raised exception's headers so the 400 contract and the handler both stay intact, rather than returning a response

@mateo-berri

Copy link
Copy Markdown
Contributor

@greptileai

@mateo-berri mateo-berri left a comment

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.

LGTM; thanks!

@tin-berri
tin-berri merged commit fc3c21e into litellm_internal_staging Jul 6, 2026
129 checks passed
@tin-berri
tin-berri deleted the litellm_mcp_short_oauth_state branch July 6, 2026 22:47
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.

2 participants