Merge pull request #34848 from BerriAI/litellm_lit4863_headless_oauth

feat(mcp): manual authorization-code delivery for headless MCP clients

litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py

@@ -1778,10 +1778,12 @@ async def token_endpoint(
 
 
 @router.post("/authorize/complete")
-async def authorize_complete(request: Request, flow: str = Form(...)):
+async def authorize_complete(request: Request, flow: str = Form(...), delivery: str | None = Form(None)):
     """Finish an aggregate connect flow: mint the gateway authorization code for the
-    signed-in user and redirect back to the DCR client. POST plus the per-flow HttpOnly
-    cookie set at /authorize; an anonymous or bad-flow request just 400s."""
+    signed-in user and hand it back to the DCR client, by 303 redirect (default) or, for
+    a loopback client on a different machine, as a copyable callback URL
+    (``delivery=manual``). POST plus the per-flow HttpOnly cookie set at /authorize; an
+    anonymous or bad-flow request just 400s."""
     from litellm.proxy.proxy_server import user_api_key_cache  # noqa: PLC0415  # circular import at module load
 
     return await complete_connect_flow(
@@ -1789,6 +1791,7 @@ async def authorize_complete(request: Request, flow: str = Form(...)):
         flow_handle=flow,
         session_user_id=_session_cookie_user_id(request),
         cache=user_api_key_cache,
+        delivery=delivery,
     )
 
 

litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py

@@ -39,6 +39,7 @@ from __future__ import annotations
 
 import hashlib
 import hmac
+import html
 import secrets
 from base64 import urlsafe_b64encode
 from collections.abc import Mapping
@@ -47,7 +48,7 @@ from typing import Awaitable, Callable, Literal, TypeVar
 from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse
 
 from fastapi import HTTPException, Request
-from fastapi.responses import JSONResponse, RedirectResponse, Response
+from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse, Response
 from pydantic import BaseModel, ConfigDict, Field, ValidationError
 from typing_extensions import assert_never
 
@@ -94,6 +95,13 @@ server-side session store, and the sealed value never appears in a URL)."""
 
 CONNECT_FLOW_TTL_SECONDS = 600
 GATEWAY_AUTH_CODE_TTL_SECONDS = 120
+MANUAL_DELIVERY_AUTH_CODE_TTL_SECONDS = 300
+"""Lifetime of a code the user delivers by hand (headless/remote client, LIT-4863 class):
+copy-pasting a callback URL from a laptop browser to an SSH session is slower than a
+browser redirect, so manual-delivery codes get 5 minutes instead of 2, still well under
+the 10-minute ceiling RFC 6749 section 4.1.2 recommends. Single-use and PKCE binding are
+unchanged, so the longer window only extends how long the legitimate holder has to paste
+it, not what an observer could do with it."""
 _CLAIM_TTL_BUFFER_SECONDS = 60
 _USED_CODE_CACHE_PREFIX = "mcp_gateway_dcr_code_used:"
 _USED_FLOW_CACHE_PREFIX = "mcp_gateway_dcr_flow_used:"
@@ -390,6 +398,7 @@ async def complete_connect_flow(
     flow_handle: str,
     session_user_id: str | None,
     cache: DualCache,
+    delivery: str | None = None,
 ) -> Response:
     """The deliberate finish step of the connect flow: mint the gateway authorization
     code and send the browser back to the client.
@@ -399,7 +408,24 @@ async def complete_connect_flow(
     into the flow: a link crafted by another party dies here with ``access_denied``
     instead of minting a code for the victim's identity. The flow is single-use (an atomic
     claim on its ``jti``), so a double-submit cannot mint two codes from one sign-in.
+
+    ``delivery`` chooses how the code reaches the client. Default (absent or
+    ``"redirect"``) is the 303 to the client's registered redirect URI. ``"manual"``
+    renders the callback URL on a page instead, for a client whose redirect URI is a
+    loopback host but which runs on a DIFFERENT machine than the browser (EC2/SSH box,
+    container): the 303 would dereference the browser machine's loopback and the code
+    would never arrive, so the user carries it over by pasting the URL into the client or
+    fetching it from the client machine's terminal. Manual delivery is honored only for
+    loopback redirect URIs; a routable redirect URI works from any browser by
+    construction, so those flows always redirect. The user who sees the page is exactly
+    the user the 303 would have carried the code to, and the same user already sees the
+    code today in the dead redirect's address bar, so the page exposes the code to no new
+    party. Unknown ``delivery`` values are rejected rather than defaulted: a client that
+    asked for manual delivery and got a dead redirect instead would silently lose its
+    code.
     """
+    if delivery not in (None, "redirect", "manual"):
+        return _oauth_error(400, "invalid_request", "delivery must be 'redirect' or 'manual'")
     sealed_flow = request.cookies.get(_flow_cookie_name(flow_handle))
     if sealed_flow is None:
         return _oauth_error(400, "invalid_request", "unknown or expired connect flow")
@@ -417,6 +443,8 @@ async def complete_connect_flow(
         f"{_USED_FLOW_CACHE_PREFIX}{flow.jti}", CONNECT_FLOW_TTL_SECONDS + _CLAIM_TTL_BUFFER_SECONDS
     ):
         return _oauth_error(400, "invalid_request", "this connect flow was already completed; restart the connection")
+    manual_delivery = delivery == "manual" and is_loopback_redirect_host(urlparse(flow.redirect_uri))
+    code_ttl = MANUAL_DELIVERY_AUTH_CODE_TTL_SECONDS if manual_delivery else GATEWAY_AUTH_CODE_TTL_SECONDS
     code = _seal(
         GATEWAY_AUTH_CODE_PREFIX,
         _GatewayAuthCode(
@@ -426,16 +454,46 @@ async def complete_connect_flow(
             code_challenge=flow.code_challenge,
             jti=secrets.token_urlsafe(24),
             iat=int(now.timestamp()),
-            exp=int(now.timestamp()) + GATEWAY_AUTH_CODE_TTL_SECONDS,
+            exp=int(now.timestamp()) + code_ttl,
         ),
     )
     params = {"code": code, **({"state": flow.state} if flow.state else {})}
-    response = RedirectResponse(_append_query_params(flow.redirect_uri, params), status_code=303)
+    callback_url = _append_query_params(flow.redirect_uri, params)
+    response: Response = (
+        _manual_delivery_response(callback_url) if manual_delivery else RedirectResponse(callback_url, status_code=303)
+    )
     path, secure = _cookie_path_and_secure(request)
     response.delete_cookie(key=_flow_cookie_name(flow_handle), path=path, secure=secure, httponly=True, samesite="lax")
     return response
 
 
+def _manual_delivery_response(callback_url: str) -> Response:
+    """The manual code-delivery page: the callback URL the 303 would have followed,
+    rendered for the user to carry to the machine the client actually runs on (paste into
+    the client's prompt, or fetch with curl from that machine's terminal). Served
+    no-store because the body holds a live single-use code, and the URL is HTML-escaped
+    because it is client-influenced. The page renders the URL as data only, never as a
+    ready-to-paste shell command: no single quoting of an attacker-influenced string is
+    correct across POSIX shells, cmd.exe, and PowerShell (cmd.exe ignores single quotes
+    and percent-expands inside double quotes), so any command string this page suggested
+    would be wrong for some shell the user might paste it into."""
+    safe_url = html.escape(callback_url, quote=True)
+    minutes = MANUAL_DELIVERY_AUTH_CODE_TTL_SECONDS // 60
+    body = (
+        "<html><head><title>Finish connecting</title></head><body>"
+        "<h2>Almost done</h2>"
+        "<p>Your MCP client runs on a different machine, so this browser cannot deliver the"
+        " authorization code to it. On the machine where the client runs, paste this URL into"
+        " the client's prompt (Claude Code accepts the pasted callback URL), or pass it as the"
+        " quoted argument of a curl command from that machine's terminal:</p>"
+        f'<p><input type="text" value="{safe_url}" readonly size="100" onclick="this.select()"></p>'
+        f"<p>The code is single-use and expires in {minutes} minutes. You can close this window"
+        " once the client confirms it is connected.</p>"
+        "</body></html>"
+    )
+    return HTMLResponse(body, headers=TOKEN_NO_CACHE_HEADERS)
+
+
 def _pkce_verifier_matches(code_verifier: str, code_challenge: str) -> bool:
     """RFC 7636 S256 verification, total over hostile input. The comparison is over bytes
     so a non-ASCII ``code_challenge`` (which reaches here unvalidated from the client's
@@ -601,9 +659,11 @@ async def _authorization_code_grant(
     if failure is not None:
         return _reload_failure_response(failure)
     # Atomic single-use claim is the gate: on a concurrent double-redeem exactly one caller
-    # wins, and a claim that cannot be recorded fails closed.
+    # wins, and a claim that cannot be recorded fails closed. The marker's TTL derives from
+    # the code's own remaining lifetime so it outlives whichever lifetime the code was minted with.
     if not await guard.claim(
-        f"{_USED_CODE_CACHE_PREFIX}{parsed.jti}", GATEWAY_AUTH_CODE_TTL_SECONDS + _CLAIM_TTL_BUFFER_SECONDS
+        f"{_USED_CODE_CACHE_PREFIX}{parsed.jti}",
+        parsed.exp - int(now.timestamp()) + _CLAIM_TTL_BUFFER_SECONDS,
     ):
         return _oauth_error(400, "invalid_grant", "the authorization code was already used")
     return _session_token_pair(SessionPrincipal(user_id=parsed.user_id, client_id=client_id), keys, now)

tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py

@@ -1,6 +1,7 @@
 """Tests for the aggregate gateway DCR flow (register, authorize, complete, token)."""
 
 import hashlib
+import html
 import json
 from base64 import urlsafe_b64encode
 from datetime import datetime, timedelta, timezone
@@ -16,7 +17,10 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import (
     GATEWAY_AUTH_CODE_PREFIX,
     GATEWAY_AUTH_CODE_TTL_SECONDS,
     GATEWAY_DCR_CLIENT_ID_PREFIX,
+    MANUAL_DELIVERY_AUTH_CODE_TTL_SECONDS,
+    _AUTH_CODE_DEBUG_KEY,
     _GatewayAuthCode,
+    _open_sealed,
     _seal,
     aggregate_authorize,
     aggregate_token,
@@ -588,3 +592,206 @@ async def test_single_use_guard_fails_closed_when_redis_errors():
 
     guard = _SingleUseGuard(cache)
     assert await guard.claim("jti-fault", 60) is False  # fail closed, not a fallback count of 1
+
+
+LOOPBACK_REDIRECT_URI = "http://localhost:3118/callback"
+
+
+async def _complete(redirect_uri: str, delivery, cookies=None, handle=None, session_user_id="u1"):
+    client_id = (await _register([redirect_uri]))["client_id"]
+    if cookies is None:
+        handle, cookies = _flow_cookie_from(_authorize(client_id, session_user_id="u1", redirect_uri=redirect_uri))
+    response = await complete_connect_flow(
+        request=_request("/authorize/complete", cookies=cookies, method="POST"),
+        flow_handle=handle,
+        session_user_id=session_user_id,
+        cache=DualCache(),
+        delivery=delivery,
+    )
+    return client_id, response
+
+
+def _callback_url_from_page(response) -> str:
+    import html as html_lib
+    import re
+
+    match = re.search(r'value="([^"]+)"', response.body.decode())
+    assert match is not None
+    return html_lib.unescape(match.group(1))
+
+
+@pytest.mark.asyncio
+async def test_manual_delivery_renders_pasteable_callback_url_for_loopback_client():
+    """The LIT-4863 headless path: a loopback client on another machine gets the callback
+    URL on a page instead of a dead 303, and the code on that page is a full-fidelity
+    authorization code (PKCE-bound, single-use, redeemable at /token)."""
+    client_id, response = await _complete(LOOPBACK_REDIRECT_URI, delivery="manual")
+    assert response.status_code == 200
+    assert response.headers["content-type"].startswith("text/html")
+    assert response.headers["cache-control"] == "no-store"
+    assert f"{CONNECT_FLOW_COOKIE_PREFIX}" in response.headers["set-cookie"]
+
+    callback_url = _callback_url_from_page(response)
+    parsed = urlparse(callback_url)
+    assert f"{parsed.scheme}://{parsed.netloc}{parsed.path}" == LOOPBACK_REDIRECT_URI
+    params = parse_qs(parsed.query)
+    assert params["state"] == ["client-state-123"]
+    code = params["code"][0]
+    assert code.startswith(GATEWAY_AUTH_CODE_PREFIX)
+
+    cache = DualCache()
+    token_response = await aggregate_token(
+        request=_request("/token", method="POST"),
+        grant_type="authorization_code",
+        code=code,
+        redirect_uri=LOOPBACK_REDIRECT_URI,
+        client_id=client_id,
+        code_verifier=CODE_VERIFIER,
+        refresh_token=None,
+        master_key=MASTER_KEY,
+        reload_user=_reload_user_active,
+        cache=cache,
+    )
+    assert token_response.status_code == 200
+
+    replay = await aggregate_token(
+        request=_request("/token", method="POST"),
+        grant_type="authorization_code",
+        code=code,
+        redirect_uri=LOOPBACK_REDIRECT_URI,
+        client_id=client_id,
+        code_verifier=CODE_VERIFIER,
+        refresh_token=None,
+        master_key=MASTER_KEY,
+        reload_user=_reload_user_active,
+        cache=cache,
+    )
+    assert json.loads(replay.body)["error"] == "invalid_grant"
+
+
+@pytest.mark.asyncio
+async def test_manual_delivery_code_gets_the_longer_ttl_and_redirect_code_does_not():
+    _, manual = await _complete(LOOPBACK_REDIRECT_URI, delivery="manual")
+    manual_code = parse_qs(urlparse(_callback_url_from_page(manual)).query)["code"][0]
+    opened_manual = _open_sealed(manual_code, GATEWAY_AUTH_CODE_PREFIX, _GatewayAuthCode, _AUTH_CODE_DEBUG_KEY)
+    assert opened_manual is not None
+    assert opened_manual.exp - opened_manual.iat == MANUAL_DELIVERY_AUTH_CODE_TTL_SECONDS
+
+    _, redirected = await _complete(LOOPBACK_REDIRECT_URI, delivery=None)
+    redirect_code = parse_qs(urlparse(redirected.headers["location"]).query)["code"][0]
+    opened_redirect = _open_sealed(redirect_code, GATEWAY_AUTH_CODE_PREFIX, _GatewayAuthCode, _AUTH_CODE_DEBUG_KEY)
+    assert opened_redirect is not None
+    assert opened_redirect.exp - opened_redirect.iat == GATEWAY_AUTH_CODE_TTL_SECONDS
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("delivery", [None, "redirect"])
+async def test_loopback_client_still_redirects_when_manual_not_requested(delivery):
+    _, response = await _complete(LOOPBACK_REDIRECT_URI, delivery=delivery)
+    assert response.status_code == 303
+    assert response.headers["location"].startswith(LOOPBACK_REDIRECT_URI)
+
+
+@pytest.mark.asyncio
+async def test_manual_delivery_is_ignored_for_routable_redirect_uri():
+    """A routable redirect URI works from any browser by construction, so manual is a
+    no-op there and the flow keeps its normal shape."""
+    _, response = await _complete(REDIRECT_URI, delivery="manual")
+    assert response.status_code == 303
+    assert response.headers["location"].startswith(REDIRECT_URI)
+
+
+@pytest.mark.asyncio
+async def test_unknown_delivery_value_is_rejected_before_the_flow_is_consumed():
+    """A typo'd delivery must not burn the single-use flow: the user fixes the form and
+    finishes normally."""
+    client_id = (await _register([LOOPBACK_REDIRECT_URI]))["client_id"]
+    handle, cookies = _flow_cookie_from(_authorize(client_id, session_user_id="u1", redirect_uri=LOOPBACK_REDIRECT_URI))
+
+    rejected = await complete_connect_flow(
+        request=_request("/authorize/complete", cookies=cookies, method="POST"),
+        flow_handle=handle,
+        session_user_id="u1",
+        cache=DualCache(),
+        delivery="carrier-pigeon",
+    )
+    assert rejected.status_code == 400
+    assert json.loads(rejected.body)["error"] == "invalid_request"
+
+    retried = await complete_connect_flow(
+        request=_request("/authorize/complete", cookies=cookies, method="POST"),
+        flow_handle=handle,
+        session_user_id="u1",
+        cache=DualCache(),
+        delivery="manual",
+    )
+    assert retried.status_code == 200
+
+
+@pytest.mark.asyncio
+async def test_manual_delivery_page_escapes_client_influenced_values():
+    """redirect_uri (and everything else on the page) is client-registered input; a quote
+    or tag in its path must render inert."""
+    hostile_uri = 'http://127.0.0.1:9/cb"><script>alert(1)</script>'
+    _, response = await _complete(hostile_uri, delivery="manual")
+    assert response.status_code == 200
+    body = response.body.decode()
+    assert "<script>alert(1)</script>" not in body
+    assert "&lt;script&gt;" in body
+
+
+class _TtlRecordingCache(DualCache):
+    """Captures the TTL of every single-use claim recorded through the in-memory arm."""
+
+    def __init__(self):
+        super().__init__()
+        self.claim_ttls: dict = {}
+
+    async def async_increment_cache(self, key, value, ttl=None, **kwargs):
+        self.claim_ttls[key] = ttl
+        return await super().async_increment_cache(key, value, ttl=ttl, **kwargs)
+
+
+@pytest.mark.asyncio
+async def test_used_code_marker_outlives_the_manually_delivered_code():
+    """Veria review finding on the LIT-4863 change: a manual code lives 300s, but the
+    used-code marker was retained for the 120s redirect lifetime plus buffer, so a client
+    could redeem, wait out the marker, and redeem the still-valid code again. The marker's
+    TTL must cover the code's own remaining lifetime plus the claim buffer."""
+    client_id, response = await _complete(LOOPBACK_REDIRECT_URI, delivery="manual")
+    code = parse_qs(urlparse(_callback_url_from_page(response)).query)["code"][0]
+
+    cache = _TtlRecordingCache()
+    token_response = await aggregate_token(
+        request=_request("/token", method="POST"),
+        grant_type="authorization_code",
+        code=code,
+        redirect_uri=LOOPBACK_REDIRECT_URI,
+        client_id=client_id,
+        code_verifier=CODE_VERIFIER,
+        refresh_token=None,
+        master_key=MASTER_KEY,
+        reload_user=_reload_user_active,
+        cache=cache,
+    )
+    assert token_response.status_code == 200
+
+    marker_ttls = [ttl for key, ttl in cache.claim_ttls.items() if key.startswith("mcp_gateway_dcr_code_used:")]
+    assert len(marker_ttls) == 1
+    assert marker_ttls[0] >= MANUAL_DELIVERY_AUTH_CODE_TTL_SECONDS
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("redirect_uri", [LOOPBACK_REDIRECT_URI, "http://127.0.0.1:9/cb$(whoami)&calc& rem x"])
+async def test_manual_delivery_page_renders_the_url_as_data_never_as_a_shell_command(redirect_uri):
+    """Two review rounds proved no single command string is safe across POSIX shells,
+    cmd.exe, and PowerShell (single quotes are not quoting in cmd.exe; percent expands
+    there even inside double quotes), so the page must render the callback URL as data
+    only and never as a ready-to-paste command."""
+    _, response = await _complete(redirect_uri, delivery="manual")
+    assert response.status_code == 200
+    body = response.body.decode()
+    assert "<code>" not in body
+    assert 'curl "' not in body
+    assert "curl '" not in body
+    assert 'value="' in body

ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.test.tsx

@@ -1,6 +1,6 @@
 import { afterEach, describe, expect, it, vi } from "vitest";
 import { render, screen } from "@testing-library/react";
-import ConnectFlowBanner from "./ConnectFlowBanner";
+import ConnectFlowBanner, { isLoopbackOrigin } from "./ConnectFlowBanner";
 
 vi.mock("@/components/networking", () => ({
   getProxyBaseUrl: () => "https://gateway.example.com",
@@ -36,6 +36,39 @@ describe("ConnectFlowBanner", () => {
     expect(screen.getAllByText(/the application/).length).toBeGreaterThan(0);
   });
 
+  it("offers manual delivery for a loopback client, posted only when checked", () => {
+    const { container } = render(
+      <ConnectFlowBanner flowHandle="flow-handle-123" clientOrigin="http://localhost:3118" />,
+    );
+
+    const checkbox = container.querySelector('input[type="checkbox"][name="delivery"]') as HTMLInputElement;
+    expect(checkbox).not.toBeNull();
+    expect(checkbox.value).toBe("manual");
+    expect(checkbox.checked).toBe(false);
+    expect(screen.getByText(/remote or SSH machine/i)).toBeInTheDocument();
+  });
+
+  it("does not offer manual delivery for a routable client origin or an unknown one", () => {
+    const routable = render(<ConnectFlowBanner flowHandle="h" clientOrigin="https://claude.ai" />);
+    expect(routable.container.querySelector('input[name="delivery"]')).toBeNull();
+
+    const unknown = render(<ConnectFlowBanner flowHandle="h" clientOrigin={null} />);
+    expect(unknown.container.querySelector('input[name="delivery"]')).toBeNull();
+  });
+
+  it("classifies loopback origins like the server does", () => {
+    expect(isLoopbackOrigin("http://localhost:3118")).toBe(true);
+    expect(isLoopbackOrigin("http://127.0.0.1:8080")).toBe(true);
+    expect(isLoopbackOrigin("http://127.5.4.3:1")).toBe(true);
+    expect(isLoopbackOrigin("http://[::1]:9000")).toBe(true);
+    expect(isLoopbackOrigin("http://[0:0:0:0:0:0:0:1]:9000")).toBe(true);
+    expect(isLoopbackOrigin("https://claude.ai")).toBe(false);
+    expect(isLoopbackOrigin("http://127.evil.com")).toBe(false);
+    expect(isLoopbackOrigin("http://localhost.evil.com")).toBe(false);
+    expect(isLoopbackOrigin(null)).toBe(false);
+    expect(isLoopbackOrigin("not a url")).toBe(false);
+  });
+
   it("does NOT complete the flow on pagehide (completion requires the explicit button)", () => {
     // Security regression: an attacker could lure a signed-in victim to their own client's
     // authorize URL; the victim merely closing the tab must NOT deliver a victim-bound code.

ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.tsx

@@ -26,9 +26,20 @@ interface Props {
  * (no click). Merely visiting the authorize URL is attacker-inducible, so completion has to be a
  * deliberate user action, not a side effect of leaving the page.
  */
+export function isLoopbackOrigin(origin: string | null): boolean {
+  if (!origin) return false;
+  try {
+    const hostname = new URL(origin).hostname.replace(/^\[|\]$/g, "");
+    return hostname === "localhost" || hostname === "::1" || /^127(\.\d{1,3}){3}$/.test(hostname);
+  } catch {
+    return false;
+  }
+}
+
 const ConnectFlowBanner: React.FC<Props> = ({ flowHandle, clientOrigin }) => {
   const action = `${getProxyBaseUrl()}/authorize/complete`;
   const clientLabel = clientOrigin ?? "the application";
+  const loopbackClient = isLoopbackOrigin(clientOrigin);
 
   return (
     <div className="mb-6 rounded-lg border border-primary/30 bg-primary/5 px-5 py-4">
@@ -50,6 +61,12 @@ const ConnectFlowBanner: React.FC<Props> = ({ flowHandle, clientOrigin }) => {
           >
             Finish connecting
           </button>
+          {loopbackClient && (
+            <label className="mt-2 flex items-center gap-2 text-[13px] text-muted-foreground">
+              <input type="checkbox" name="delivery" value="manual" />
+              My client is on a remote or SSH machine
+            </label>
+          )}
         </form>
       </div>
     </div>