fix(mcp): return 401 challenge for REST token-exchange tool calls without a subject token (#42782)

* fix(mcp): return 401 challenge for REST token-exchange tool calls without a subject token Co-Authored-By: bot_apk <apk@cognition.ai> * fix(mcp): keep tool_server_mismatch when server_id disagrees with the tool prefix Co-Authored-By: bot_apk <apk@cognition.ai> * test(mcp): type the token-exchange challenge test helpers Co-Authored-By: bot_apk <apk@cognition.ai> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: bot_apk <apk@cognition.ai>
2026-09-26 01:12:21 +00:00 · 2026-09-23 14:13:16 -07:00 · 2026-09-23 14:13:16 -07:00 · dc14e76147
commit dc14e76147
parent 2b3a7f7f8e
3 changed files with 216 additions and 4 deletions
--- a/litellm/proxy/_experimental/mcp_server/operations.py
+++ b/litellm/proxy/_experimental/mcp_server/operations.py
@ -1721,6 +1721,39 @@ async def _check_byok_credential(
        )


+def _challenge_missing_token_exchange_subject(
+    server: MCPServer | None,
+    requested_server: MCPServer | None,
+    allowed_mcp_servers: list[MCPServer],
+    user_api_key_auth: UserAPIKeyAuth | None,
+    oauth2_headers: dict[str, str] | None,
+    raw_headers: dict[str, str] | None,
+) -> None:
+    """Raise the RFC 9728 challenge when a token-exchange server is called without a subject token.
+
+    The listing that fills a cold catalog absorbs the upstream 401 by design, so without this
+    check a missing subject surfaces as an unknown-tool error instead of the challenge the
+    warm path already raises. Gated to servers the key may reach so an unauthorized caller
+    learns nothing about the catalog.
+    """
+    if server is None or server.auth_type != MCPAuth.oauth2_token_exchange:
+        return
+    if requested_server is not None and requested_server.server_id != server.server_id:
+        return
+    if all(allowed.server_id != server.server_id for allowed in allowed_mcp_servers):
+        return
+    if global_mcp_server_manager._extract_subject_token(oauth2_headers, raw_headers, user_api_key_auth) is not None:
+        return
+    from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import (  # noqa: PLC0415  # lazy: adapter pulls MCP subgraph
+        raise_token_exchange_challenge,
+    )
+    from litellm.proxy.middleware.per_request_root_path_middleware import (  # noqa: PLC0415  # lazy: middleware imports proxy utils
+        get_request_root_path,
+    )
+
+    raise_token_exchange_challenge(server, root_path=get_request_root_path())
+
+
 async def _list_tools_before_first_call(
    server: MCPServer | None,
    tool_name: str,
@ -1864,6 +1897,14 @@ async def _execute_mcp_tool(
        if first_call_target is None or (requested_server is not None and not name_is_prefixed)
        else strip_known_server_prefix(name, first_call_target)
    )
+    _challenge_missing_token_exchange_subject(
+        server=first_call_target,
+        requested_server=requested_server,
+        allowed_mcp_servers=allowed_mcp_servers,
+        user_api_key_auth=user_api_key_auth,
+        oauth2_headers=oauth2_headers,
+        raw_headers=raw_headers,
+    )
    await _list_tools_before_first_call(
        server=first_call_target,
        tool_name=first_call_tool_name,
--- a/tests/integration/mcp/test_mcp_oauth_flows.py
+++ b/tests/integration/mcp/test_mcp_oauth_flows.py
@ -171,12 +171,37 @@ def test_token_exchange_without_a_subject_token_is_rejected_before_any_upstream_
        key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
        peer.drain()
        auth.drain()
-        response: Final = call_tool(gateway, key, identity, f"{alias}-add", ADD)
+        cold: Final = call_tool(gateway, key, identity, f"{alias}-add", ADD)
        assert tool_calls(peer.drain()) == ()
        assert auth.token_requests() == ()
-        if response.status_code == 500:
-            pytest.skip("BUG: /mcp-rest/tools/call without a subject token on a token-exchange server returns 500")
-        assert response.status_code == 401, response.text
+        _assert_subject_token_challenge(cold, alias)
+        warmed: Final = gateway.client.post(
+            "/mcp-rest/tools/call",
+            headers={"x-litellm-api-key": key, "Authorization": "Bearer subject-" + uuid.uuid4().hex},
+            json={"name": f"{alias}-add", "arguments": ADD, "server_id": identity},
+        )
+        assert warmed.status_code == 200, warmed.text
+        assert len(tool_calls(peer.drain())) == 1 and len(auth.token_requests()) == 1
+        auth.drain()
+        warm: Final = call_tool(gateway, key, identity, f"{alias}-add", ADD)
+        assert tool_calls(peer.drain()) == ()
+        assert auth.token_requests() == ()
+        _assert_subject_token_challenge(warm, alias)
+        as_subject: Final = gateway.client.post(
+            "/mcp-rest/tools/call",
+            headers={"x-litellm-api-key": key, "Authorization": f"Bearer {key}"},
+            json={"name": f"{alias}-add", "arguments": ADD, "server_id": identity},
+        )
+        assert tool_calls(peer.drain()) == ()
+        assert auth.token_requests() == ()
+        _assert_subject_token_challenge(as_subject, alias)
+
+
+def _assert_subject_token_challenge(response: httpx.Response, alias: str) -> None:
+    assert response.status_code == 401, response.text
+    challenge: Final = response.headers["www-authenticate"]
+    assert challenge.startswith("Bearer ") and 'error="invalid_token"' in challenge, challenge
+    assert f'resource_metadata="/.well-known/oauth-protected-resource/mcp/{alias}"' in challenge, challenge


@pytest.mark.parametrize("entry", ENTRY_POINTS)
--- a/tests/test_litellm/proxy/_experimental/mcp_server/test_operations.py
+++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_operations.py
@ -6,6 +6,8 @@ from mcp.types import GetPromptRequest, GetPromptRequestParams, GetPromptResult

 from litellm.proxy._experimental.mcp_server.operations import GatewayOperations, prepare_context
 from litellm.proxy._types import UserAPIKeyAuth
+from litellm.types.mcp import MCPAuth, MCPTransport
+from litellm.types.mcp_server.mcp_server_manager import MCPServer


@pytest.mark.asyncio
@ -363,3 +365,147 @@ async def test_explicit_proxy_context_lists_builtin_tools_and_blocks_direct_tool
    assert denied.is_error is True
    assert "unavailable on /mcp/proxy" in denied.content[0].text
    allowed.assert_not_awaited()
+
+
+def _server(server_id: str, auth_type: MCPAuth) -> MCPServer:
+    return MCPServer(
+        server_id=server_id,
+        name=f"{server_id}-server",
+        url="https://up.example.com/mcp",
+        transport=MCPTransport.http,
+        auth_type=auth_type,
+        token_exchange_endpoint="https://idp.example.com/token",
+        client_id="cid",
+        client_secret="csec",
+    )
+
+
+class TestChallengeMissingTokenExchangeSubject:
+    """The REST cold-catalog path must answer a missing OBO subject with the RFC 9728 401 challenge
+    before the best-effort listing swallows the upstream 401 and tool resolution turns it into a 500."""
+
+    @staticmethod
+    def _challenge(
+        server: MCPServer | None,
+        allowed: list[MCPServer],
+        *,
+        user: UserAPIKeyAuth | None = None,
+        oauth2_headers: dict[str, str] | None = None,
+        raw_headers: dict[str, str] | None = None,
+        requested_server: MCPServer | None = None,
+    ) -> None:
+        from litellm.proxy._experimental.mcp_server.operations import _challenge_missing_token_exchange_subject
+
+        return _challenge_missing_token_exchange_subject(
+            server=server,
+            requested_server=requested_server,
+            allowed_mcp_servers=allowed,
+            user_api_key_auth=user,
+            oauth2_headers=oauth2_headers,
+            raw_headers=raw_headers,
+        )
+
+    def test_missing_subject_raises_401_challenge(self):
+        from fastapi import HTTPException
+
+        server = _server("te-cold", MCPAuth.oauth2_token_exchange)
+        with pytest.raises(HTTPException) as exc_info:
+            self._challenge(
+                server,
+                [server],
+                user=UserAPIKeyAuth(api_key="sk-admission"),
+                raw_headers={"x-litellm-api-key": "sk-admission"},
+            )
+        assert exc_info.value.status_code == 401
+        challenge = (exc_info.value.headers or {}).get("WWW-Authenticate", "")
+        assert challenge.startswith("Bearer ") and 'error="invalid_token"' in challenge, challenge
+        assert "resource_metadata" in challenge, challenge
+
+    @pytest.mark.parametrize(
+        "authorization",
+        ["Bearer sk-admission", "Bearer sk-some-other-virtual-key"],
+        ids=["repeated-admission-key", "another-virtual-key"],
+    )
+    def test_litellm_key_in_authorization_is_not_a_subject(self, authorization: str):
+        from fastapi import HTTPException
+
+        server = _server("te-vk", MCPAuth.oauth2_token_exchange)
+        with pytest.raises(HTTPException) as exc_info:
+            self._challenge(
+                server,
+                [server],
+                user=UserAPIKeyAuth(api_key="sk-admission"),
+                oauth2_headers={"Authorization": authorization},
+                raw_headers={"x-litellm-api-key": "sk-admission", "authorization": authorization},
+            )
+        assert exc_info.value.status_code == 401
+
+    def test_subject_present_does_not_challenge(self):
+
+        server = _server("te-ok", MCPAuth.oauth2_token_exchange)
+        assert (
+            self._challenge(
+                server,
+                [server],
+                user=UserAPIKeyAuth(api_key="sk-admission"),
+                oauth2_headers={"Authorization": "Bearer idp-subject"},
+                raw_headers={"x-litellm-api-key": "sk-admission", "authorization": "Bearer idp-subject"},
+            )
+            is None
+        )
+
+    def test_server_outside_allowlist_is_not_challenged(self):
+
+        server = _server("te-hidden", MCPAuth.oauth2_token_exchange)
+        other = _server("te-visible", MCPAuth.oauth2_token_exchange)
+        assert self._challenge(server, [other], user=UserAPIKeyAuth(api_key="sk-admission")) is None
+        assert self._challenge(None, [other], user=UserAPIKeyAuth(api_key="sk-admission")) is None
+
+    def test_prefix_owner_differing_from_server_id_is_not_challenged(self):
+        """An explicit server_id that disagrees with the tool prefix keeps the existing mismatch answer."""
+        from fastapi import HTTPException
+
+        prefix_owner = _server("te-prefix", MCPAuth.oauth2_token_exchange)
+        requested = _server("te-requested", MCPAuth.oauth2_token_exchange)
+        user = UserAPIKeyAuth(api_key="sk-admission")
+        allowed = [prefix_owner, requested]
+        assert self._challenge(prefix_owner, allowed, user=user, requested_server=requested) is None
+        with pytest.raises(HTTPException):
+            self._challenge(prefix_owner, allowed, user=user, requested_server=prefix_owner)
+
+    @pytest.mark.parametrize(
+        "auth_type",
+        ["oauth2", "oauth_delegate", "oauth2_id_jag", "bearer_token", "api_key", "none"],
+    )
+    def test_other_auth_types_are_untouched(self, auth_type: str):
+
+        server = _server("na", MCPAuth(auth_type))
+        assert self._challenge(server, [server], user=UserAPIKeyAuth(api_key="sk-admission")) is None
+
+
+@pytest.mark.asyncio
+async def test_execute_mcp_tool_challenges_missing_subject_before_cold_listing():
+    """On a cold catalog the challenge fires before any listing or tool resolution is attempted."""
+    from fastapi import HTTPException
+    from datetime import datetime, timezone
+    from litellm.proxy._experimental.mcp_server import operations
+
+    server = _server("te-exec", MCPAuth.oauth2_token_exchange)
+    listing = AsyncMock()
+    with (
+        patch.object(operations.global_mcp_server_manager, "get_mcp_server_by_id", return_value=server),
+        patch.object(operations.global_mcp_server_manager, "server_exposes_tool", return_value=False),
+        patch.object(operations, "_get_tools_from_mcp_servers", listing),
+        pytest.raises(HTTPException) as exc_info,
+    ):
+        await operations.execute_mcp_tool(
+            name="add",
+            arguments={"a": 2, "b": 3},
+            allowed_mcp_servers=[server],
+            start_time=datetime.now(timezone.utc),
+            user_api_key_auth=UserAPIKeyAuth(api_key="sk-admission"),
+            raw_headers={"x-litellm-api-key": "sk-admission"},
+            requested_server_id=server.server_id,
+        )
+    assert exc_info.value.status_code == 401
+    listing.assert_not_awaited()