From 6959d9de69af616b84d35605bdfbf1d1e14278dd Mon Sep 17 00:00:00 2001
From: Tin Chi Lo <tin@berri.ai>
Date: Mon, 20 Jul 2026 19:48:37 -0700
Subject: [PATCH] test(mcp): pin the token and register wall messages for
 url-less servers

---
 .../mcp_server/test_discoverable_endpoints.py | 80 +++++++++++++++++++
 1 file changed, 80 insertions(+)

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 5b1768ee563..636c7fbd3d5 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
@@ -8068,3 +8068,83 @@ async def test_authorize_wall_names_the_fix_for_urlless_servers():
     detail_text = str(exc_info.value.detail)
     assert "set Authorization URL and Token URL" in detail_text
     assert "Issuer" in detail_text
+
+
+@pytest.mark.asyncio
+async def test_token_wall_names_the_fix_for_urlless_servers():
+    """The /token wall is the second stop on the same misconfiguration (LIT-4629): after an admin
+    fills only the Authorization URL, the code exchange dies here; the detail must name the
+    remedies like the authorize wall does."""
+    from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
+        exchange_token_with_server,
+    )
+    from litellm.types.mcp import MCPAuth, MCPTransport
+    from litellm.types.mcp_server.mcp_server_manager import MCPServer
+
+    server = MCPServer(
+        server_id="urlless-token-wall",
+        name="sheets_token_wall",
+        server_name="sheets_token_wall",
+        url=None,
+        transport=MCPTransport.http,
+        auth_type=MCPAuth.oauth2,
+        spec_path="https://example.com/openapi.yaml",
+        authorization_url="https://accounts.google.com/o/oauth2/v2/auth",
+    )
+    mock_request = MagicMock()
+    mock_request.base_url = "https://litellm.example.com/"
+    mock_request.headers = {}
+
+    with pytest.raises(HTTPException) as exc_info:
+        await exchange_token_with_server(
+            request=mock_request,
+            mcp_server=server,
+            grant_type="authorization_code",
+            code="auth-code",
+            redirect_uri="http://localhost/callback",
+            client_id="client",
+            client_secret=None,
+            code_verifier="verifier",
+        )
+    assert exc_info.value.status_code == 400
+    detail_text = str(exc_info.value.detail)
+    assert "set Token URL manually" in detail_text
+    assert "Issuer" in detail_text
+
+
+@pytest.mark.asyncio
+async def test_register_wall_names_the_fix_for_urlless_servers():
+    """The /register wall serves the same missing-authorization-url 400 as authorize; its detail
+    must carry the same actionable remedies."""
+    from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
+        register_client_with_server,
+    )
+    from litellm.types.mcp import MCPAuth, MCPTransport
+    from litellm.types.mcp_server.mcp_server_manager import MCPServer
+
+    server = MCPServer(
+        server_id="urlless-register-wall",
+        name="sheets_register_wall",
+        server_name="sheets_register_wall",
+        url=None,
+        transport=MCPTransport.http,
+        auth_type=MCPAuth.oauth2,
+        spec_path="https://example.com/openapi.yaml",
+    )
+    mock_request = MagicMock()
+    mock_request.base_url = "https://litellm.example.com/"
+    mock_request.headers = {}
+
+    with pytest.raises(HTTPException) as exc_info:
+        await register_client_with_server(
+            request=mock_request,
+            mcp_server=server,
+            client_name="client",
+            grant_types=None,
+            response_types=None,
+            token_endpoint_auth_method=None,
+        )
+    assert exc_info.value.status_code == 400
+    detail_text = str(exc_info.value.detail)
+    assert "set Authorization URL and Token URL" in detail_text
+    assert "Issuer" in detail_text