Skip to content

feat(mcp): aggregate DCR register, authorize, complete, and token flow for the gateway front door - #33189

Merged
tin-berri merged 2 commits into
litellm_lit3637_session_tokenfrom
litellm_lit3637_aggregate_flow
Jul 17, 2026
Merged

feat(mcp): aggregate DCR register, authorize, complete, and token flow for the gateway front door#33189
tin-berri merged 2 commits into
litellm_lit3637_session_tokenfrom
litellm_lit3637_aggregate_flow

Conversation

@tin-berri

Copy link
Copy Markdown
Contributor

Relevant issues

Stacked on #33188 (UI session cookie exp) -> #33182 (session tokens) -> #33174 (aggregate discovery). The base of this PR is litellm_lit3637_ui_session_exp, not staging; review and merge after those

Linear ticket

Part of LIT-3637 (PR 3 of the stack: the aggregate DCR register/authorize/complete/token flow that turns the discovery front door into a working sign-in)

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 received a Greptile Confidence Score of at least 4/5 before requesting a maintainer review (Greptile reviews automatically once the PR is opened; only comment @greptileai to re-request a review after pushing changes)

Delays in PR merge?

If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).

Screenshots / Proof of Fix

Proven live end to end at commit head on a real proxy (:4138, Postgres) with mcp_gateway_dcr: true, driving the exact request sequence a DCR client (Claude Desktop, MCP Inspector) makes. The upstream-server authorization step is the existing per-user vault flow and is not re-exercised here; this PR is the identity front door

1. UI login mints an exp-bearing session cookie
   $ curl -s -X POST :4138/login -d 'username=admin&password=sk-1234' -c jar.txt
   cookie exp present: True | user_id: default_user_id | login_method: username_password | ttl_hrs: 24.0

2. DCR register (stateless public client)
   $ curl -s -X POST :4138/register -d '{"redirect_uris":["http://localhost:9666/callback"],"client_name":"curl-dcr"}'
   client_id: llm_dcrc_vGkDD-R3OEQ...  token_endpoint_auth_method: none  (no client_secret)

3. signed-in GET /authorize -> 303 to the connect grid, per-flow HttpOnly cookie set
   location: /ui/chat/integrations?connect_flow=AMdstIy3kMNB-H0x7l0DjLlFHPiafES6&connect_client=http://localhost:9666

4. POST /authorize/complete -> 303 back to the client with a sealed code, state echoed
   location: http://localhost:9666/callback?code=llm_gcode_wV2vWaDMQlTYpf...&state=rig-state-1

5. POST /token (authorization_code + PKCE verifier) -> identity-only session pair
   access_token: llm_session_eyJhbGci...  refresh_token: llm_srefresh_eyJhbGciO...  expires_in: 3600

6. the access token carries ONLY the litellm user, never an upstream secret
   iss: litellm-mcp-gateway | kind: session | user_id: default_user_id
   claims: [client_id, exp, iat, iss, jti, kind, user_id]   no upstream token in claims: True

7. replay the used authorization code -> invalid_grant   (single-use guard)
8. refresh_token grant -> new access minted, refresh rotated
9. wrong PKCE verifier on a fresh code -> invalid_grant
10. anonymous GET /authorize -> 303 /sso/key/generate?return_to=/authorize%3Fclient_id%3Dllm_dcrc_...  (same-origin round-trip)

Type

🆕 New Feature

Changes

PR 1 gave the aggregate /mcp endpoint OAuth discovery; this PR is the flow that discovery points at, so an OAuth-only DCR client can actually sign in. Everything is gated behind mcp_gateway_dcr and behind the llm_dcrc_ client-id prefix, so with the flag off, or for any client that did not register through this flow, the existing per-server authorize/token behavior is byte-identical. The whole flow lives in one new module, gateway_dcr_flow.py, wired into the existing root /register, /authorize, and /token handlers plus one new POST /authorize/complete

Registration is stateless. The client_id is the registration: the client's redirect URIs are sealed into it with the repo's authenticated symmetric helper, so nothing is persisted, open registration cannot fill storage, and a forged or tampered client_id simply fails to open. Clients are always registered as public (token_endpoint_auth_method "none") because the gateway issues no client secrets; mandatory S256 PKCE is what protects the code. Redirect URIs must be https, or http on a loopback host for local dev clients (RFC 8252), fragment-free, and bounded in count and length so the sealed id stays small

Authorize validates the client and redirect URI, requires S256 PKCE, and interposes LiteLLM sign-in. A validation failure answers 400 directly and never redirects, per RFC 6749 4.1.2.1: an unvalidated redirect URI must not receive an error redirect. Without a session cookie the browser is sent through /sso/key/generate with a strictly-relative return_to (a new same-origin arm in the SSO callback, so login can only ever bounce the browser back to this same authorize request, never off-origin). With a session, the flow parameters and the SSO user are sealed into a per-flow HttpOnly cookie (the same handle-plus-cookie pattern the upstream OAuth relay already uses) and the browser lands on the connect grid, where the user authorizes individual servers before finishing

Complete is a deliberate, separate step reached by POST, bound to the HttpOnly flow cookie and to an exact match between the signed-in user and the user sealed into the flow. A GET could be triggered cross-site; a POST with a SameSite=Lax cookie and the identity check cannot mint a code for a victim's session. It seals a short-lived (120s), single-use, PKCE- and client-bound authorization code and redirects back to the registered redirect URI

Token exchanges the code for the identity-only session tokens from the module below it in the stack. The authorization_code grant verifies the code has not expired, is bound to this client and redirect URI, passes PKCE, and has not already been used (a best-effort single-use guard over the shared cache, in-memory or Redis), then re-validates that the litellm user is still active before minting. The refresh_token grant re-validates the same way and rotates the pair, with the refresh token bound to its issuing client (RFC 6749 section 6). Every error is an RFC 6749 5.2 body carrying no token, code, or URL material

The one change outside this module and its wiring is the SSO callback's same-origin return_to arm, needed so the login round-trip returns the browser to the authorize request. Sealed values open totally (bad input maps to an OAuth error, never a raise); no upstream server credential appears anywhere in this flow, they are vaulted per user by the existing /v1/mcp authorize endpoints and resolved at egress by user id (unchanged here)

QA runbook

Run a proxy with mcp_gateway_dcr: true, a master key, and a database, then drive the numbered sequence above with curl (the full script is in the proof). Confirm: the login cookie carries an exp; register returns an llm_dcrc_ public client; a signed-in authorize 303s to /ui/chat/integrations?connect_flow=... and sets the per-flow cookie; complete 303s back with an llm_gcode_ code and the original state; token returns llm_session_/llm_srefresh_ with expires_in: 3600; the used code replays as invalid_grant; refresh rotates; a wrong PKCE verifier is invalid_grant; and an anonymous authorize 303s to /sso/key/generate with a relative return_to. With the flag off, confirm the aggregate well-known routes 404 and /register returns its former dummy body

Final Attestation

  • The tests check the right things, including the edge cases, and regressions in the respective real-world customer use-cases are not possible after this PR

@codecov

codecov Bot commented Jul 14, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 94.29658% with 15 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
...proxy/_experimental/mcp_server/gateway_dcr_flow.py 95.76% 10 Missing ⚠️
litellm/proxy/management_endpoints/ui_sso.py 50.00% 5 Missing ⚠️

📢 Thoughts on this report? Let us know!

@greptile-apps

greptile-apps Bot commented Jul 14, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR implements the full aggregate DCR OAuth front door for the MCP gateway (mcp_gateway_dcr: true): stateless client registration, an authorize \u2192 login-round-trip \u2192 connect-page flow, a POST /authorize/complete step that mints a short-lived PKCE-bound code, and an identity-only token/refresh exchange. It also adds the same-origin return_to arm in the SSO callback needed so a browser sent through login arrives back at the authorize request.

  • gateway_dcr_flow.py is a new self-contained module. The stateless client_id design (redirect URIs sealed with authenticated encryption) avoids any storage, and the flow cookie\u2019s HttpOnly + SameSite=Lax + POST requirement at /authorize/complete correctly prevents cross-site code minting.
  • discoverable_endpoints.py wires the new flow into existing root endpoints behind is_gateway_dcr_client_id + is_mcp_gateway_dcr_enabled guards, leaving all non-DCR paths byte-identical.
  • ui_sso.py adds _is_same_origin_return_path and uses it as an additional acceptor for return_to in the SSO login/callback path; tests cover absolute URLs, protocol-relative, backslash tricks, and javascript: schemes.

Confidence Score: 3/5

The stateless DCR design and cookie-bound complete step are sound, but a transient DB outage during the token exchange permanently invalidates the authorization code, making the gateway's own 503 non-retryable.

The mark_used call preceding reload_user in _authorization_code_grant causes valid authorization codes to be permanently invalidated on transient DB outages — a present behavioral defect on the new token exchange path.

gateway_dcr_flow.py: the ordering of mark_used and reload_user in _authorization_code_grant (lines 495-500) and the empty-prefix _seal call for the flow cookie (line 321).

Important Files Changed

Filename Overview
litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py New module implementing the complete aggregate DCR OAuth flow. Core logic is clean; TOCTOU race between already_used/mark_used allows concurrent code replay, and marking the code as used before reload_user permanently invalidates valid codes on transient DB failures.
litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py Wires new DCR flow into existing root endpoints with routing guards ensuring zero behavioral change for non-DCR clients.
litellm/proxy/management_endpoints/ui_sso.py Adds _is_same_origin_return_path helper; check correctly rejects protocol-relative and backslash variants and is well-tested.
tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py Comprehensive tests covering full register→authorize→complete→token walk plus security edges. No real network calls.

Reviews (1): Last reviewed commit: "feat(mcp): aggregate DCR register, autho..." | Re-trigger Greptile

Comment on lines +495 to +500
if await guard.already_used(parsed.jti):
return _oauth_error(400, "invalid_grant", "the authorization code was already used")
await guard.mark_used(parsed.jti)
failure = await reload_user(parsed.user_id)
if failure is not None:
return _reload_failure_response(failure)

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 Single-use guard has a TOCTOU race and marks the code used before the user check

Two concurrent token requests carrying the same code can both get False from already_used before either writes the mark, so both proceed to mint — a narrow but real asyncio-level race since the event loop can switch coroutines between the two awaits. More concretely, mark_used is called on line 497 before reload_user on line 498. When reload_user returns "unavailable" (transient DB outage), the code is already permanently marked; the client gets a 503 and reasonably retries, but then receives invalid_grant — forcing a complete re-authorization cycle through the browser even though the failure was entirely the gateway's own DB blip.

A safer ordering would finish validation (reload_user) before consuming the single-use slot, and ideally use an atomic set-if-not-exists primitive to close the TOCTOU window.

Comment on lines +165 to +166
def _seal(prefix: str, payload: BaseModel) -> str:
return prefix + encrypt_value_helper(payload.model_dump_json())

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.

P2 Empty-string prefix makes _open_sealed prefix check vacuous for flow cookies

_seal("", flow) seals with no prefix, so the prefix check in _open_sealed (value.startswith("")) is always true — the actual protection relies entirely on decrypt_value_helper's authenticated encryption plus Pydantic validation. This is functionally safe but means a sealed value of any other type (e.g. a GatewayDcrClient token) would reach Pydantic deserialization as a _ConnectFlow candidate before being rejected by field mismatch.

Using a distinct non-empty prefix for the flow cookie (e.g. "mcp_cf_") would make the routing intention explicit and add an early-exit path consistent with how GATEWAY_DCR_CLIENT_ID_PREFIX and GATEWAY_AUTH_CODE_PREFIX are used.

failure = await reload_user(opened.principal.user_id)
if failure is not None:
return _reload_failure_response(failure)
return _session_token_pair(opened.principal, keys, now)

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.

Medium: Refresh token replay

This mints a new token pair without invalidating the presented refresh token. An attacker who obtains that token can repeatedly mint fresh access and refresh tokens for up to its 14-day lifetime, including after the legitimate client has refreshed. Consume the refresh token's signed jti with an atomic shared-cache claim before minting, and reject subsequent uses.

@veria-ai

veria-ai Bot commented Jul 14, 2026

Copy link
Copy Markdown
Contributor

PR overview

This PR adds an aggregated MCP gateway front-door flow for dynamic client registration, authorization, completion, and token issuance. The touched code includes the gateway DCR flow handling token refresh behavior.

There is one open security issue in the token refresh path: refresh tokens can be reused because the presented token is not consumed or invalidated before minting a new token pair. If a refresh token is obtained by an attacker, it could be replayed during its validity window to maintain access under that client’s scope. No issues have been addressed yet, so the PR still needs a fix for refresh-token rotation/replay protection before the flow is in a safer state.

Open issues (1)

Fixed/addressed: 0 · PR risk: 6/10

@tin-berri
tin-berri force-pushed the litellm_lit3637_ui_session_exp branch from 9a81539 to 436673f Compare July 15, 2026 00:27
- atomic single-use guard (async_increment_cache) + reload-before-claim so a transient DB blip does not burn a valid code
- PKCE verify over bytes so a non-ASCII code_challenge fails invalid_grant instead of raising a 500; validate code_verifier length (RFC 7636)
- flag-off byte-identical for a server literally named mcp (AS well-known delegates to the named-server document)
- connect flow is single-use (atomic jti claim) so a double-submit cannot mint two codes
- extra=forbid on the sealed models; bound state length; drop unused request param and coarse dict on register
- _reload_failure_response exhaustive match+assert_never; dedupe ReloadUserFailure with _KeyResolutionFailure
- reject control/whitespace chars in the same-origin return_to
@tin-berri
tin-berri force-pushed the litellm_lit3637_aggregate_flow branch from f9b90ee to 90eb828 Compare July 15, 2026 00:38
Base automatically changed from litellm_lit3637_ui_session_exp to litellm_lit3637_session_token July 17, 2026 21:46
@tin-berri
tin-berri merged commit 186c862 into litellm_lit3637_session_token Jul 17, 2026
99 of 101 checks passed
@tin-berri
tin-berri deleted the litellm_lit3637_aggregate_flow branch July 17, 2026 21:52
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