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

litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py

@@ -1,14 +1,15 @@
 import asyncio
 import html as _html
 import json
+import secrets
 import time
 from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Tuple
 from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse
 
 import httpx
 from fastapi import APIRouter, Form, HTTPException, Request
-from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse
+from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse, Response
 from pydantic import BaseModel, ValidationError
 
 from litellm._logging import verbose_logger
@@ -137,6 +138,72 @@ def decode_state_hash(encrypted_state: str) -> dict:
     return state_data
 
 
+# LIT-4197: some upstream authorization servers reject an over-long ``state``
+# (the encrypted OAuth session blob routinely exceeds their limit). The upstream
+# only needs an opaque value it echoes back on ``/callback``, so we forward a
+# short random handle and keep the encrypted session in a per-flow HttpOnly
+# cookie bound to that handle. The browser carries the cookie across the
+# upstream round trip, so the flow stays correct with no server-side session
+# store (works across proxy replicas, unlike an in-process map).
+_OAUTH_STATE_COOKIE_PREFIX = "mcp_oauth_state_"
+_OAUTH_STATE_COOKIE_TTL_SECONDS = 600
+_OAUTH_STATE_HANDLE_BYTES = 32
+
+
+def _oauth_state_cookie_name(relay_state: str) -> str:
+    return f"{_OAUTH_STATE_COOKIE_PREFIX}{relay_state}"
+
+
+def _oauth_state_cookie_path_and_secure(request: Request) -> tuple[str, bool]:
+    parsed = urlparse(get_request_base_url(request))
+    return parsed.path or "/", parsed.scheme == "https"
+
+
+def _set_oauth_state_cookie(
+    response: Response,
+    request: Request,
+    relay_state: str,
+    encoded_state: str,
+) -> None:
+    path, secure = _oauth_state_cookie_path_and_secure(request)
+    response.set_cookie(
+        key=_oauth_state_cookie_name(relay_state),
+        value=encoded_state,
+        max_age=_OAUTH_STATE_COOKIE_TTL_SECONDS,
+        path=path,
+        secure=secure,
+        httponly=True,
+        samesite="lax",
+    )
+
+
+def _resolve_encoded_oauth_state(request: Request, state: str) -> str:
+    """Return the encrypted OAuth session for a ``/callback`` request.
+
+    New flows carry it in a per-flow cookie keyed by the short handle we
+    forwarded upstream (the IdP echoes that handle back as ``state``). Flows
+    started before this change - or in flight across a deploy - carry the
+    encrypted blob directly in ``state``, so fall back to it when the cookie
+    is absent.
+    """
+    cookie_value = request.cookies.get(_oauth_state_cookie_name(state))
+    return cookie_value if cookie_value else state
+
+
+def _clear_oauth_state_cookie(response: Response, request: Request, state: str) -> None:
+    cookie_name = _oauth_state_cookie_name(state)
+    if cookie_name not in request.cookies:
+        return
+    path, secure = _oauth_state_cookie_path_and_secure(request)
+    response.delete_cookie(
+        key=cookie_name,
+        path=path,
+        secure=secure,
+        httponly=True,
+        samesite="lax",
+    )
+
+
 def _get_validated_client_redirect_uri(request: Request, state_data: Dict[str, Any]) -> str:
     """Return a trusted (same-origin, loopback, or ops-allowlisted)
     client redirect URI from OAuth state.
@@ -462,11 +529,12 @@ async def authorize_with_server(
         code_challenge_method=code_challenge_method,
         client_redirect_uri=redirect_uri,
     )
+    relay_state = secrets.token_urlsafe(_OAUTH_STATE_HANDLE_BYTES)
 
     params = {
         "client_id": mcp_server.client_id if mcp_server.client_id else client_id,
         "redirect_uri": f"{request_base_url}/callback",
-        "state": encoded_state,
+        "state": relay_state,
         "response_type": response_type or "code",
     }
     if scope:
@@ -483,7 +551,9 @@ async def authorize_with_server(
     existing_params = dict(parse_qsl(parsed_auth_url.query))
     existing_params.update(params)
     final_url = urlunparse(parsed_auth_url._replace(query=urlencode(existing_params)))
-    return RedirectResponse(final_url)
+    response = RedirectResponse(final_url)
+    _set_oauth_state_cookie(response, request, relay_state, encoded_state)
+    return response
 
 
 async def exchange_token_with_server(
@@ -1016,17 +1086,19 @@ async def callback(
             error_description,
         )
         if state:
+            encoded_state = _resolve_encoded_oauth_state(request, state)
             try:
-                state_data = decode_state_hash(state)
+                state_data = decode_state_hash(encoded_state)
                 original_state = state_data.get("original_state")
                 redirect_uri = _get_validated_client_redirect_uri(request, state_data)
-            except HTTPException:
-                # Untrusted/invalid client redirect_uri — surface inline rather
-                # than blindly forwarding the error to an attacker-controlled URL.
-                return _render_oauth_error_html(error, error_description)
             except Exception:
-                # State could not be decrypted (expired key, tampered, etc.).
-                return _render_oauth_error_html(error, error_description)
+                # Untrusted/invalid client redirect_uri (HTTPException), or an
+                # undecryptable state (expired key, tampered): surface the IdP
+                # error inline rather than forwarding it to an attacker-controlled
+                # URL, and drop the one-time cookie we can no longer consume.
+                response = _render_oauth_error_html(error, error_description)
+                _clear_oauth_state_cookie(response, request, state)
+                return response
 
             params: Dict[str, str] = {"error": error}
             if error_description:
@@ -1036,7 +1108,9 @@ async def callback(
             if original_state is not None:
                 params["state"] = original_state
             complete_returned_url = _append_query_params(redirect_uri, params)
-            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
 
         # No state — nothing to round-trip to. Show the user the error.
         return _render_oauth_error_html(error, error_description)
@@ -1052,7 +1126,8 @@ async def callback(
 
     # 3. Successful authorization response.
     try:
-        state_data = decode_state_hash(state)
+        encoded_state = _resolve_encoded_oauth_state(request, state)
+        state_data = decode_state_hash(encoded_state)
         original_state = state_data["original_state"]
 
         # Re-validate the client redirect URI at the sink. /authorize
@@ -1065,14 +1140,18 @@ async def callback(
 
         params = {"code": code, "state": original_state}
         complete_returned_url = _append_query_params(redirect_uri, params)
-        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
 
     except HTTPException:
         # Re-raise so a non-loopback base_url surfaces as 400 instead of
         # a generic "authentication incomplete" redirect.
         raise
     except Exception:
-        return HTMLResponse("<html><body>Authentication incomplete. You can close this window.</body></html>")
+        response = HTMLResponse("<html><body>Authentication incomplete. You can close this window.</body></html>")
+        _clear_oauth_state_cookie(response, request, state)
+        return response
 
 
 # ------------------------------

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

@@ -35,6 +35,7 @@ def _mock_callback_request(base_url: str = "http://localhost:3000/"):
     req = MagicMock()
     req.base_url = base_url
     req.headers = {}
+    req.cookies = {}
     return req
 
 
@@ -2630,6 +2631,162 @@ async def test_oauth_callback_accepts_same_origin_ui_redirect():
     assert "state=state-123" in response.headers["location"]
 
 
+@pytest.mark.asyncio
+async def test_authorize_forwards_short_state_and_round_trips_via_cookie(monkeypatch):
+    """LIT-4197: the ``state`` sent to the upstream authorization server must be
+    a short opaque handle, not the long encrypted OAuth session (some IdPs
+    reject an over-long state). The session must instead ride in a per-flow
+    HttpOnly cookie so ``/callback`` still recovers the client's original state
+    and redirects back to the client's redirect_uri."""
+    from http.cookies import SimpleCookie
+    from urllib.parse import parse_qs, urlparse
+
+    from fastapi import Request
+
+    from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
+        _oauth_state_cookie_name,
+        authorize_with_server,
+        callback,
+        decode_state_hash,
+    )
+    from litellm.proxy._types import MCPTransport
+    from litellm.types.mcp import MCPAuth
+    from litellm.types.mcp_server.mcp_server_manager import MCPServer
+
+    # Real encryption so the cookie value is a genuine encrypted session.
+    monkeypatch.setenv("LITELLM_SALT_KEY", "sk-test-salt-for-LIT-4197")
+
+    client_state = "ee230e3dfd4f19c7441941684f39c8a4e0e2c3c61a088e33403df5662b4047b8"
+    client_redirect_uri = "http://127.0.0.1:6274/oauth/callback/debug"
+
+    server = MCPServer(
+        server_id="leanix_server",
+        name="leanix",
+        server_name="leanix",
+        alias="leanix",
+        transport=MCPTransport.http,
+        auth_type=MCPAuth.oauth2,
+        client_id="upstream-client-id",
+        authorization_url="https://idp.example.com/oauth/authorize",
+        token_url="https://idp.example.com/oauth/token",
+    )
+
+    authorize_request = MagicMock(spec=Request)
+    authorize_request.base_url = "https://proxy.example.com/"
+    authorize_request.headers = {}
+
+    authorize_response = await authorize_with_server(
+        request=authorize_request,
+        mcp_server=server,
+        client_id="upstream-client-id",
+        redirect_uri=client_redirect_uri,
+        state=client_state,
+        code_challenge="challenge",
+        code_challenge_method="S256",
+    )
+
+    location = authorize_response.headers["location"]
+    upstream_state = parse_qs(urlparse(location).query)["state"][0]
+
+    # The upstream must receive a short handle, not the encrypted session blob.
+    assert len(upstream_state) <= 64
+    assert upstream_state != client_state
+
+    # The encrypted session rides in a per-flow HttpOnly cookie bound to it.
+    jar = SimpleCookie()
+    jar.load(authorize_response.headers["set-cookie"])
+    cookie_name = _oauth_state_cookie_name(upstream_state)
+    assert cookie_name in jar
+    morsel = jar[cookie_name]
+    assert morsel["httponly"]
+    assert morsel["samesite"].lower() == "lax"
+    assert len(morsel.value) > len(upstream_state)
+    session = decode_state_hash(morsel.value)
+    assert session["original_state"] == client_state
+    assert session["client_redirect_uri"] == client_redirect_uri
+
+    # /callback recovers the original state from the cookie (not the handle) and
+    # redirects back to the client with the client's own state.
+    callback_request = MagicMock(spec=Request)
+    callback_request.base_url = "https://proxy.example.com/"
+    callback_request.headers = {}
+    callback_request.cookies = {cookie_name: morsel.value}
+
+    callback_response = await callback(
+        request=callback_request,
+        code="upstream-auth-code",
+        state=upstream_state,
+    )
+
+    assert callback_response.status_code == 302
+    cb_query = parse_qs(urlparse(callback_response.headers["location"]).query)
+    assert callback_response.headers["location"].startswith(client_redirect_uri)
+    assert cb_query["code"] == ["upstream-auth-code"]
+    assert cb_query["state"] == [client_state]
+
+    # The one-time cookie is expired on the callback response so it cannot be replayed.
+    cleared = SimpleCookie()
+    cleared.load(callback_response.headers["set-cookie"])
+    assert cookie_name in cleared
+    assert cleared[cookie_name].value == ""
+    assert cleared[cookie_name]["max-age"] == "0"
+
+
+@pytest.mark.asyncio
+async def test_callback_error_path_reads_cookie_and_clears_it(monkeypatch):
+    """LIT-4197: an IdP error routed through /callback must recover the client's
+    original state from the cookie (not the short handle), propagate the error to
+    the client's redirect_uri, and expire the one-time cookie."""
+    from http.cookies import SimpleCookie
+    from urllib.parse import parse_qs, urlparse
+
+    from fastapi import Request
+
+    from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
+        _oauth_state_cookie_name,
+        callback,
+        encode_state_with_base_url,
+    )
+
+    monkeypatch.setenv("LITELLM_SALT_KEY", "sk-test-salt-for-LIT-4197")
+
+    client_state = "client-original-state-abc"
+    client_redirect_uri = "http://127.0.0.1:6274/oauth/callback/debug"
+    handle = "shortRelayHandle123"
+    encoded_state = encode_state_with_base_url(
+        base_url=client_redirect_uri,
+        original_state=client_state,
+        client_redirect_uri=client_redirect_uri,
+    )
+    cookie_name = _oauth_state_cookie_name(handle)
+
+    request = MagicMock(spec=Request)
+    request.base_url = "https://proxy.example.com/"
+    request.headers = {}
+    request.cookies = {cookie_name: encoded_state}
+
+    response = await callback(
+        request=request,
+        error="access_denied",
+        error_description="User declined access",
+        state=handle,
+    )
+
+    assert response.status_code == 302
+    location = response.headers["location"]
+    assert location.startswith(client_redirect_uri)
+    query = parse_qs(urlparse(location).query)
+    assert query["error"] == ["access_denied"]
+    # The client's own state is echoed back, recovered from the cookie.
+    assert query["state"] == [client_state]
+
+    cleared = SimpleCookie()
+    cleared.load(response.headers["set-cookie"])
+    assert cookie_name in cleared
+    assert cleared[cookie_name].value == ""
+    assert cleared[cookie_name]["max-age"] == "0"
+
+
 @pytest.mark.asyncio
 async def test_oauth_authorize_includes_scopes_from_server_config():
     """Test that authorize endpoint includes scopes from server configuration."""