From ccdd7b59ac2aaab9dba09e73130f8bc5a3c6f789 Mon Sep 17 00:00:00 2001
From: jesus <jesus@berri.ai>
Date: Thu, 17 Sep 2026 02:12:12 +0000
Subject: [PATCH] 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()