From f1b5c0e4d5658d975bde2d8805093a040e676674 Mon Sep 17 00:00:00 2001 From: jesus Date: Thu, 17 Sep 2026 01:54:54 +0000 Subject: [PATCH 1/4] fix(mcp): surface missing OAuth state cookie on /callback instead of silent 200 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/discoverable_endpoints.py | 26 ++++++- .../mcp_server/test_discoverable_endpoints.py | 78 ++++++++++++++++++- 2 files changed, 100 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index ffb27d5f92e..e8f864a136c 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -1,6 +1,7 @@ import asyncio import html as _html import json +import os import secrets import time from collections.abc import Callable, Mapping @@ -437,6 +438,10 @@ def _resolve_encoded_oauth_state(request: Request, state: str) -> str: return cookie_value if cookie_value else state +def _oauth_state_cookie_present(request: Request, state: str) -> bool: + return _oauth_state_cookie_name(state) in request.cookies + + def _clear_oauth_state_cookie(response: Response, request: Request, state: str) -> None: cookie_name: Final = _oauth_state_cookie_name(state) if cookie_name not in request.cookies: @@ -2235,9 +2240,28 @@ async def callback( ) # 3. Successful authorization response. + encoded_state = _resolve_encoded_oauth_state(request, state) try: - encoded_state = _resolve_encoded_oauth_state(request, state) state_data = decode_state_hash(encoded_state) + except Exception: + cookie_present: Final = _oauth_state_cookie_present(request, state) + verbose_logger.warning( + "MCP /callback could not decode OAuth state (state_cookie_present=%s, request_base_url=%s, " + "PROXY_BASE_URL_set=%s). If the cookie is absent, /authorize and /callback were served from " + "different origins; set PROXY_BASE_URL to the public origin or configure mcp_trusted_proxy_ranges.", + cookie_present, + get_request_base_url(request), + bool(os.environ.get("PROXY_BASE_URL", "").strip()), + ) + description: Final = ( + "The OAuth session cookie set when authorization started did not arrive at the callback. " + "This usually means the authorize and callback requests used different origins. " + "Ask the gateway operator to set PROXY_BASE_URL to the public URL of this gateway, then retry." + if not cookie_present + else "The OAuth session could not be decoded. Start the authorization again." + ) + return _render_oauth_error_html("invalid_request", description) + try: original_state = state_data["original_state"] # Re-validate the client redirect URI at the sink. /authorize diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index aa45b2f6793..c09afa58cd0 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -4275,9 +4275,81 @@ async def test_oauth_callback_handles_invalid_state(): state="invalid_encrypted_state", ) - # Should return HTML error page - assert response.status_code == 200 - assert "Authentication incomplete" in response.body.decode() + # Should return a 400 error page that points at the missing state cookie + assert response.status_code == 400 + body = response.body.decode() + assert "did not arrive" in body + assert "PROXY_BASE_URL" in body + + +@pytest.mark.asyncio +async def test_oauth_callback_missing_state_cookie_logs_origin_hint(monkeypatch): + """A /callback whose per-flow state cookie never arrived must log a warning + carrying the resolved base URL and whether PROXY_BASE_URL is set, so the + origin-mismatch behind an ingress is diagnosable from logs.""" + try: + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + callback, + ) + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.decode_state_hash" + ) as mock_decode, patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.verbose_logger" + ) as mock_logger: + mock_decode.side_effect = ValueError("bad state") + + response = await callback( + request=_mock_callback_request(), + code="test_code", + state="relay-handle-without-cookie", + ) + + assert response.status_code == 400 + mock_logger.warning.assert_called_once() + warning_args = mock_logger.warning.call_args + assert "state_cookie_present=%s" in warning_args.args[0] + assert "PROXY_BASE_URL_set=%s" in warning_args.args[0] + assert warning_args.args[1] is False + assert warning_args.args[2] == "http://localhost:3000" + assert warning_args.args[3] is False + + +@pytest.mark.asyncio +async def test_oauth_callback_undecodable_state_with_cookie_present(): + """When the state cookie DID arrive but still does not decode, the error + page must not claim the cookie was lost in transit.""" + try: + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _oauth_state_cookie_name, + callback, + ) + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + state = "relay-handle" + request = _mock_callback_request() + request.cookies = {_oauth_state_cookie_name(state): "garbage"} + + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.decode_state_hash" + ) as mock_decode: + mock_decode.side_effect = ValueError("bad state") + + response = await callback( + request=request, + code="test_code", + state=state, + ) + + assert response.status_code == 400 + body = response.body.decode() + assert "could not be decoded" in body + assert "did not arrive" not in body @pytest.mark.asyncio From 4022928ac768df4fc8583eda0aa4bb21934db7a9 Mon Sep 17 00:00:00 2001 From: jesus Date: Thu, 17 Sep 2026 02:06:58 +0000 Subject: [PATCH 2/4] fix(mcp): suppress BLE001 on deliberate broad state-decode catch in /callback Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/_experimental/mcp_server/discoverable_endpoints.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index e8f864a136c..5bdcfe5be20 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -2243,7 +2243,7 @@ async def callback( encoded_state = _resolve_encoded_oauth_state(request, state) try: state_data = decode_state_hash(encoded_state) - except Exception: + except Exception: # noqa: BLE001 # any decode failure means the session is unusable; surface it, never crash the callback cookie_present: Final = _oauth_state_cookie_present(request, state) verbose_logger.warning( "MCP /callback could not decode OAuth state (state_cookie_present=%s, request_base_url=%s, " From ccdd7b59ac2aaab9dba09e73130f8bc5a3c6f789 Mon Sep 17 00:00:00 2001 From: jesus Date: Thu, 17 Sep 2026 02:12:12 +0000 Subject: [PATCH 3/4] test(mcp): exercise real state decode failure instead of patching internals Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/test_discoverable_endpoints.py | 66 ++++++++----------- 1 file changed, 26 insertions(+), 40 deletions(-) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index c09afa58cd0..cdc7373e763 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -4264,29 +4264,28 @@ async def test_oauth_callback_handles_invalid_state(): except ImportError: pytest.skip("MCP discoverable endpoints not available") - # Mock state decoding to raise an exception - with patch("litellm.proxy._experimental.mcp_server.discoverable_endpoints.decode_state_hash") as mock_decode: - mock_decode.side_effect = Exception("Failed to decrypt state") + # A state blob that cannot be decrypted takes the same path as the + # LIT-4197 relay handle echoed back without its per-flow cookie + response = await callback( + request=_mock_callback_request(), + code="test_code", + state="invalid_encrypted_state", + ) - # Call callback endpoint with invalid state - response = await callback( - request=_mock_callback_request(), - code="test_code", - state="invalid_encrypted_state", - ) - - # Should return a 400 error page that points at the missing state cookie - assert response.status_code == 400 - body = response.body.decode() - assert "did not arrive" in body - assert "PROXY_BASE_URL" in body + # Should return a 400 error page that points at the missing state cookie + assert response.status_code == 400 + body = response.body.decode() + assert "did not arrive" in body + assert "PROXY_BASE_URL" in body @pytest.mark.asyncio -async def test_oauth_callback_missing_state_cookie_logs_origin_hint(monkeypatch): +async def test_oauth_callback_missing_state_cookie_logs_origin_hint(monkeypatch, caplog): """A /callback whose per-flow state cookie never arrived must log a warning carrying the resolved base URL and whether PROXY_BASE_URL is set, so the origin-mismatch behind an ingress is diagnosable from logs.""" + import logging + try: from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( callback, @@ -4296,13 +4295,7 @@ async def test_oauth_callback_missing_state_cookie_logs_origin_hint(monkeypatch) monkeypatch.delenv("PROXY_BASE_URL", raising=False) - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.decode_state_hash" - ) as mock_decode, patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.verbose_logger" - ) as mock_logger: - mock_decode.side_effect = ValueError("bad state") - + with caplog.at_level(logging.WARNING, logger="LiteLLM"): response = await callback( request=_mock_callback_request(), code="test_code", @@ -4310,13 +4303,11 @@ async def test_oauth_callback_missing_state_cookie_logs_origin_hint(monkeypatch) ) assert response.status_code == 400 - mock_logger.warning.assert_called_once() - warning_args = mock_logger.warning.call_args - assert "state_cookie_present=%s" in warning_args.args[0] - assert "PROXY_BASE_URL_set=%s" in warning_args.args[0] - assert warning_args.args[1] is False - assert warning_args.args[2] == "http://localhost:3000" - assert warning_args.args[3] is False + matching = [r.getMessage() for r in caplog.records if "could not decode OAuth state" in r.getMessage()] + assert len(matching) == 1 + assert "state_cookie_present=False" in matching[0] + assert "request_base_url=http://localhost:3000" in matching[0] + assert "PROXY_BASE_URL_set=False" in matching[0] @pytest.mark.asyncio @@ -4335,16 +4326,11 @@ async def test_oauth_callback_undecodable_state_with_cookie_present(): request = _mock_callback_request() request.cookies = {_oauth_state_cookie_name(state): "garbage"} - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.decode_state_hash" - ) as mock_decode: - mock_decode.side_effect = ValueError("bad state") - - response = await callback( - request=request, - code="test_code", - state=state, - ) + response = await callback( + request=request, + code="test_code", + state=state, + ) assert response.status_code == 400 body = response.body.decode() From fdcdea0d372de3ee536ea3461136b186a4159d02 Mon Sep 17 00:00:00 2001 From: jesus Date: Thu, 17 Sep 2026 02:25:38 +0000 Subject: [PATCH 4/4] fix(mcp): clear stale state cookie and mention trusted-proxy config in callback error Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../_experimental/mcp_server/discoverable_endpoints.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 5bdcfe5be20..5b18b861d64 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -2256,11 +2256,14 @@ async def callback( description: Final = ( "The OAuth session cookie set when authorization started did not arrive at the callback. " "This usually means the authorize and callback requests used different origins. " - "Ask the gateway operator to set PROXY_BASE_URL to the public URL of this gateway, then retry." + "Ask the gateway operator to set PROXY_BASE_URL to the public URL of this gateway " + "(or configure mcp_trusted_proxy_ranges), then retry." if not cookie_present else "The OAuth session could not be decoded. Start the authorization again." ) - return _render_oauth_error_html("invalid_request", description) + response = _render_oauth_error_html("invalid_request", description) + _clear_oauth_state_cookie(response, request, state) + return response try: original_state = state_data["original_state"]