Merge ee7a7c59d0 into 2e427c6247

litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py

@@ -145,10 +145,16 @@ every session-token claim set. Four 256-character ASCII URIs seal to roughly 1.5
 the encoded client_id is checked against its own cap before registration succeeds.
 VS Code registers four callbacks for its web and desktop environments."""
 
-MAX_STATE_LENGTH: Final = 1024
-"""Bound on the client ``state`` sealed into the flow cookie and echoed on the auth-code
-redirect. An unbounded ``state`` can push the sealed cookie past the browser's ~4KB cap
-(silently dropped, breaking the flow); spec clients send a short opaque value."""
+MAX_STATE_LENGTH: Final = 2048
+"""Coarse bound on the client ``state`` sealed into the flow cookie and echoed on the
+auth-code redirect. RFC 6749 sets no limit and real clients send long opaque values
+(Gemini web sends ~1.1KB), so this only rejects clearly abusive input; the sealed cookie
+itself is checked against ``MAX_FLOW_COOKIE_BYTES``."""
+
+MAX_FLOW_COOKIE_BYTES: Final = 4000
+"""Bound on the flow cookie's name plus sealed value. Browsers drop cookies over ~4096
+bytes silently, which would break the flow at the connect page with an opaque "unknown or
+expired connect flow", so an oversized authorize request is rejected up front instead."""
 
 MIN_CODE_VERIFIER_LENGTH: Final = 43
 MAX_CODE_VERIFIER_LENGTH: Final = 128
@@ -758,8 +764,11 @@ def aggregate_authorize(
         upstream_scopes=upstream_scopes,
     )
     connect_url: Final = _append_query_params(f"{base_url}/ui/connect", (("connect_flow", handle),))
+    sealed_flow: Final = _seal(_UNPREFIXED, flow)
+    if _flow_cookie_too_large(handle, sealed_flow):
+        return _flow_cookie_too_large_error()
     response: Final = RedirectResponse(connect_url, status_code=303)
-    _set_flow_cookie(response, request, handle, flow)
+    _set_flow_cookie(response, request, handle, sealed_flow)
     return response
 
 
@@ -804,6 +813,9 @@ async def native_client_authorize(
         resource_server_ids=None,
         audience=PROXY_API_AUDIENCE,
     )
+    sealed_flow: Final = _seal(_UNPREFIXED, flow)
+    if _flow_cookie_too_large(handle, sealed_flow):
+        return _flow_cookie_too_large_error()
     page: Final = render_native_client_consent_page(
         client_origin=_origin_only(redirect_uri),
         user_id=session_user_id,
@@ -812,7 +824,7 @@ async def native_client_authorize(
         complete_url=f"{base_url}/authorize/complete",
     )
     response: Final = HTMLResponse(page, headers=_CONSENT_PAGE_HEADERS)
-    _set_flow_cookie(response, request, handle, flow)
+    _set_flow_cookie(response, request, handle, sealed_flow)
     return response
 
 
@@ -936,11 +948,19 @@ def _new_connect_flow(
     )
 
 
-def _set_flow_cookie(response: Response, request: Request, handle: str, flow: _ConnectFlow) -> None:
+def _flow_cookie_too_large(handle: str, sealed_flow: str) -> bool:
+    return len(_flow_cookie_name(handle)) + len(sealed_flow) > MAX_FLOW_COOKIE_BYTES
+
+
+def _flow_cookie_too_large_error() -> Response:
+    return _oauth_error(400, "invalid_request", "authorization request is too large (shorten state or client_id)")
+
+
+def _set_flow_cookie(response: Response, request: Request, handle: str, sealed_flow: str) -> None:
     path, secure = _cookie_path_and_secure(request)
     response.set_cookie(
         key=_flow_cookie_name(handle),
-        value=_seal(_UNPREFIXED, flow),
+        value=sealed_flow,
         max_age=CONNECT_FLOW_TTL_SECONDS,
         path=path,
         secure=secure,

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

@@ -23,6 +23,8 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import (
     JWT_BEARER_GRANT_TYPE,
     MANUAL_DELIVERY_AUTH_CODE_TTL_SECONDS,
     MAX_CLIENT_ID_LENGTH,
+    MAX_FLOW_COOKIE_BYTES,
+    MAX_STATE_LENGTH,
     SUBJECT_TOKEN_TYPES,
     TOKEN_EXCHANGE_GRANT_TYPE,
     AssertionIdentity,
@@ -631,7 +633,7 @@ async def test_authorize_rejects_over_long_state():
         request=_request(query=f"client_id={client_id}"),
         client_id=client_id,
         redirect_uri=REDIRECT_URI,
-        state="s" * 2000,
+        state="s" * (MAX_STATE_LENGTH + 1),
         code_challenge=CODE_CHALLENGE,
         code_challenge_method="S256",
         response_type="code",
@@ -641,6 +643,41 @@ async def test_authorize_rejects_over_long_state():
     assert json.loads(response.body)["error"] == "invalid_request"
 
 
+def _authorize_with_state(client_id, redirect_uri, state):
+    return aggregate_authorize(
+        request=_request(query=f"client_id={client_id}"),
+        client_id=client_id,
+        redirect_uri=redirect_uri,
+        state=state,
+        code_challenge=CODE_CHALLENGE,
+        code_challenge_method="S256",
+        response_type="code",
+        session_user_id="u1",
+    )
+
+
+@pytest.mark.asyncio
+async def test_authorize_accepts_long_opaque_state_from_gemini():
+    """Gemini web sends a ~1.1KB opaque state (#46414); RFC 6749 sets no limit on it."""
+    client_id = (await _register([REDIRECT_URI]))["client_id"]
+    response = _authorize_with_state(client_id, REDIRECT_URI, "s" * 1112)
+    assert response.status_code == 303
+    flow_cookie = next(header.decode() for name, header in response.raw_headers if name == b"set-cookie")
+    name_and_value = flow_cookie.split(";", 1)[0]
+    assert len(name_and_value) - 1 <= MAX_FLOW_COOKIE_BYTES
+
+
+@pytest.mark.asyncio
+async def test_authorize_rejects_request_whose_flow_cookie_would_exceed_browser_limit():
+    """A within-cap state plus a maximal registration must not set a cookie the browser
+    would silently drop; the request is rejected up front instead."""
+    client_id = (await _register(list(MAX_LENGTH_REDIRECT_URIS)))["client_id"]
+    response = _authorize_with_state(client_id, MAX_LENGTH_REDIRECT_URIS[0], "s" * MAX_STATE_LENGTH)
+    assert response.status_code == 400
+    assert json.loads(response.body)["error"] == "invalid_request"
+    assert not any(name == b"set-cookie" for name, _ in response.raw_headers)
+
+
 @pytest.mark.asyncio
 async def test_non_ascii_code_challenge_fails_grant_not_500():
     """A non-ASCII code_challenge (unvalidated from the client) must yield a clean