diff --git a/litellm/proxy/_experimental/mcp_server/operations.py b/litellm/proxy/_experimental/mcp_server/operations.py index 2bd6d186f24..26bf68d9932 100644 --- 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, diff --git a/tests/integration/mcp/test_mcp_oauth_flows.py b/tests/integration/mcp/test_mcp_oauth_flows.py index bc83ca7ea50..4efb78dacc2 100644 --- 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) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_operations.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_operations.py index abb925ddc77..81877c38389 100644 --- 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()