diff --git a/litellm/proxy/mcp_server/mcp_server_manager.py b/litellm/proxy/mcp_server/mcp_server_manager.py index 5b4ebe41948..d31e9246a38 100644 --- a/litellm/proxy/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/mcp_server/mcp_server_manager.py @@ -373,6 +373,12 @@ _UpstreamSessionKey: TypeAlias = tuple[ ] +@dataclass(frozen=True, slots=True) +class ResourceListingFailed: + server_name: str + error: Exception + + @dataclass(frozen=True, slots=True) class _OAuthDiscoveryResolved: server: MCPServer @@ -1621,13 +1627,14 @@ def _redacted_registry_dump(servers: dict[str, MCPServer]) -> dict[str, dict[str def caller_authorization_fans_out( server: MCPServer, - scope_servers: list[MCPServer] | None, + scope_servers: Sequence[MCPServer] | None, ) -> bool: """True when forwarding the caller's request-wide ``Authorization`` to ``server`` inside a listing fan-out would replay one credential against multiple upstreams: another server in the scope also consumes it (RFC 9700 cross-resource replay). ``scope_servers`` is None for - explicitly-addressed operations (tool call, get_prompt, read_resource, single-server routes), - where the client named the one target and the gateway is not choosing recipients.""" + explicitly-addressed operations (tool call, get_prompt, single-server routes), where the client + named the one target and the gateway is not choosing recipients. An aggregate read_resource passes + its scope because the gateway picks the owner from listings.""" if scope_servers is None: return False return any( @@ -5099,8 +5106,33 @@ class MCPServerManager: raw_headers: dict[str, str] | None = None, client_ip: str | None = None, ) -> list[Resource]: - if self._skip_blocked_stdio_listing(server, "resource"): + listing: Final = await self.list_resources_from_server( + server=server, + user_api_key_auth=user_api_key_auth, + mcp_auth_header=mcp_auth_header, + extra_headers=extra_headers, + add_prefix=add_prefix, + raw_headers=raw_headers, + client_ip=client_ip, + ) + if isinstance(listing, ResourceListingFailed): + verbose_logger.warning("Failed to get resources from server %s: %s", server.name, listing.error) return [] + return list(listing) + + async def list_resources_from_server( + self, + server: MCPServer, + user_api_key_auth: UserAPIKeyAuth | None, + mcp_auth_header: str | Mapping[str, str] | None = None, + extra_headers: Mapping[str, str] | None = None, + add_prefix: bool = True, + raw_headers: Mapping[str, str] | None = None, + client_ip: str | None = None, + ) -> tuple[Resource, ...] | ResourceListingFailed: + if self._skip_blocked_stdio_listing(server, "resource"): + return () + auth_header: Final = dict(mcp_auth_header) if isinstance(mcp_auth_header, Mapping) else mcp_auth_header try: headers: Final = ( dict( @@ -5115,7 +5147,7 @@ class MCPServerManager: subject_token: Final = self._obo_subject_token(server, raw_headers, user_api_key_auth) client: Final = await self.create_mcp_client( server=server, - mcp_auth_header=mcp_auth_header, + mcp_auth_header=auth_header, extra_headers=headers, stdio_env=stdio_env, subject_token=subject_token, @@ -5127,7 +5159,7 @@ class MCPServerManager: key: Final = self._discovery_key( server, user_api_key_auth, - mcp_auth_header, + auth_header, headers, stdio_env, subject_token, @@ -5139,10 +5171,9 @@ class MCPServerManager: return await client.list_resources_result(raise_on_error=True) items: Final = await self._resource_discovery_cache.get(key, fetch) - return self._create_prefixed_resources(items.resources, server, add_prefix=add_prefix) + return tuple(self._create_prefixed_resources(items.resources, server, add_prefix=add_prefix)) except Exception as error: - verbose_logger.warning("Failed to get resources from server %s: %s", server.name, error) - return [] + return ResourceListingFailed(server_name=server.name, error=error) async def get_resource_templates_from_server( self, diff --git a/litellm/proxy/mcp_server/operations.py b/litellm/proxy/mcp_server/operations.py index 73a16291c61..6c9b8184e1a 100644 --- a/litellm/proxy/mcp_server/operations.py +++ b/litellm/proxy/mcp_server/operations.py @@ -112,6 +112,7 @@ from litellm.proxy.mcp_server.interactions import ( from litellm.proxy.mcp_server.legacy_callbacks import get_scope_response from litellm.proxy.mcp_server.mcp_server_manager import ( # noqa: F401 # legacy module exports MCPServerManager, + ResourceListingFailed, _caller_authorization_fans_out, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export _client_forwarded_authorization_headers, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export _resolve_openapi_tool_auth, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export @@ -142,6 +143,7 @@ from litellm.proxy.mcp_server.result_conversion import ( handler_outcome, to_call_tool_result, ) +from litellm.proxy.mcp_server.stdio_gate import is_mcp_stdio_blocked from litellm.proxy.mcp_server.tool_registry import ( global_mcp_tool_registry, ) @@ -262,6 +264,61 @@ class _MCPServerRateLimitAdmission: rejected_servers: tuple[tuple[MCPServer, ProxyRateLimitError], ...] +@dataclass(frozen=True, slots=True) +class _ResourceOwnerRateLimited: + error: ProxyRateLimitError + + +@dataclass(frozen=True, slots=True) +class _ResourceOwnerMissing: + pass + + +@dataclass(frozen=True, slots=True) +class _ResourceOwnerAmbiguous: + server_names: tuple[str, ...] + + +_ResourceOwnerFailure: TypeAlias = ( + _ResourceOwnerRateLimited | ResourceListingFailed | _ResourceOwnerMissing | _ResourceOwnerAmbiguous +) + + +def _withhold_fanned_out_caller_authorization( + server: MCPServer, extra_headers: dict[str, str] | None, scope_servers: Sequence[MCPServer] +) -> dict[str, str] | None: + """Drop the caller's request-wide Authorization from a gateway-selected read when another server in + scope also consumes it. ``_prepare_mcp_server_headers`` already withholds it from client-forwarded + servers, but legacy oauth2 delegates and OAuth pass-through servers still carry it, and an aggregate + read must not replay one caller token to servers the gateway picked.""" + if extra_headers and caller_authorization_fans_out(server, scope_servers): + return without_header(extra_headers, DEFAULT_CREDENTIAL_HEADER) + return extra_headers + + +def _raise_resource_owner_failure(failure: _ResourceOwnerFailure, url: AnyUrl) -> NoReturn: + match failure: + case _ResourceOwnerRateLimited(error=error): + raise error + case ResourceListingFailed(server_name=server_name): + raise HTTPException( + status_code=502, + detail=f"Unable to list resources from MCP server '{server_name}'; resource owner is unknown.", + ) + case _ResourceOwnerMissing(): + raise HTTPException(status_code=404, detail=f"Resource '{url}' not found on any allowed MCP server.") + case _ResourceOwnerAmbiguous(server_names=server_names): + raise HTTPException( + status_code=400, + detail=( + f"Resource URI '{url}' is ambiguous across MCP servers: {', '.join(server_names)}. " + "Use /{server_name}/mcp or x-mcp-servers to select one server." + ), + ) + case _: + assert_never(failure) + + _mcp_server_admission_memo: Final[ContextVar[dict[str, asyncio.Task[ProxyRateLimitError | None]] | None]] = ContextVar( "mcp_server_admission_memo", default=None ) @@ -1000,7 +1057,7 @@ def _prepare_mcp_server_headers( oauth2_headers: dict[str, str] | None, raw_headers: dict[str, str] | None, user_api_key_auth: UserAPIKeyAuth | None = None, - scope_servers: list[MCPServer] | None = None, + scope_servers: Sequence[MCPServer] | None = None, ) -> tuple[dict[str, str] | str | None, dict[str, str] | None]: """Build auth and extra headers for a server. @@ -3131,13 +3188,58 @@ async def mcp_read_resource( detail="User not allowed to read this resource.", ) - if len(allowed_mcp_servers) != 1: - raise HTTPException( - status_code=400, - detail=("Multiple MCP servers configured; read_resource currently supports exactly one allowed server."), - ) + candidates: Final = tuple(server for server in allowed_mcp_servers if not is_mcp_stdio_blocked(server.transport)) - server: Final = allowed_mcp_servers[0] + async def owner(candidate: MCPServer) -> MCPServer | ResourceListingFailed | None: + candidate_auth, candidate_headers = _prepare_mcp_server_headers( + server=candidate, + mcp_server_auth_headers=mcp_server_auth_headers, + mcp_auth_header=mcp_auth_header, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + user_api_key_auth=user_api_key_auth, + scope_servers=candidates, + ) + listing: Final = await global_mcp_server_manager.list_resources_from_server( + server=candidate, + user_api_key_auth=user_api_key_auth, + mcp_auth_header=candidate_auth, + extra_headers=_withhold_fanned_out_caller_authorization(candidate, candidate_headers, candidates), + add_prefix=False, + raw_headers=raw_headers, + client_ip=client_ip, + ) + if isinstance(listing, ResourceListingFailed): + return listing + return candidate if any(str(resource.uri) == str(url) for resource in listing) else None + + async def resolve_owner() -> MCPServer | _ResourceOwnerFailure: + admission: Final = await _admit_mcp_servers(candidates, user_api_key_auth) + if admission.rejected_servers: + return _ResourceOwnerRateLimited(error=admission.rejected_servers[0][1]) + lookups: Final = tuple(asyncio.create_task(owner(candidate)) for candidate in candidates) + try: + for lookup in asyncio.as_completed(lookups): + if isinstance(failure := await lookup, ResourceListingFailed): + return failure + finally: + for lookup in lookups: + lookup.cancel() + await asyncio.gather(*lookups, return_exceptions=True) + results: Final = tuple(lookup.result() for lookup in lookups) + owners: Final = tuple(result for result in results if isinstance(result, MCPServer)) + if len(owners) == 1: + return owners[0] + if not owners: + return _ResourceOwnerMissing() + return _ResourceOwnerAmbiguous(server_names=tuple(owner.name for owner in owners)) + + if len(allowed_mcp_servers) == 1: + await _enforce_mcp_server_rate_limit(user_api_key_auth, allowed_mcp_servers[0]) + resolved: Final = allowed_mcp_servers[0] if len(allowed_mcp_servers) == 1 else await resolve_owner() + if isinstance(resolved, _ResourceOwnerFailure): + _raise_resource_owner_failure(resolved, url) + server: Final = resolved server_auth_header, extra_headers = _prepare_mcp_server_headers( server=server, @@ -3146,15 +3248,15 @@ async def mcp_read_resource( oauth2_headers=oauth2_headers, raw_headers=raw_headers, user_api_key_auth=user_api_key_auth, + scope_servers=candidates, ) - await _enforce_mcp_server_rate_limit(user_api_key_auth, server) return await global_mcp_server_manager.read_resource_from_server( server=server, user_api_key_auth=user_api_key_auth, url=url, mcp_auth_header=server_auth_header, - extra_headers=extra_headers, + extra_headers=_withhold_fanned_out_caller_authorization(server, extra_headers, candidates), raw_headers=raw_headers, client_ip=client_ip, input_responses=input_responses, @@ -3690,21 +3792,27 @@ async def _execute_read_resource( _client_ip, ) = context.legacy_auth() - read_resource_result: Final = await mcp_read_resource( - url=params.uri, - user_api_key_auth=user_api_key_auth, - mcp_auth_header=mcp_auth_header, - mcp_servers=mcp_servers, - mcp_server_auth_headers=mcp_server_auth_headers, - oauth2_headers=oauth2_headers, - raw_headers=raw_headers, - client_ip=_client_ip, - input_responses=params.input_responses, - request_state=params.request_state, - allow_input_required=context.wire_compat is WireCompat.MODERN, - ) + from mcp.shared.exceptions import MCPError + from mcp.types import INVALID_REQUEST - return read_resource_result + try: + return await mcp_read_resource( + url=params.uri, + user_api_key_auth=user_api_key_auth, + mcp_auth_header=mcp_auth_header, + mcp_servers=mcp_servers, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + client_ip=_client_ip, + input_responses=params.input_responses, + request_state=params.request_state, + allow_input_required=context.wire_compat is WireCompat.MODERN, + ) + except MCPError: + raise + except HTTPException as error: + raise MCPError(code=INVALID_REQUEST, message=_http_detail_message(error.detail)) from error def _reject_mcp_proxy_operation() -> NoReturn: diff --git a/tests/unit/proxy/mcp_server/test_mcp_server_tool_calls_and_headers.py b/tests/unit/proxy/mcp_server/test_mcp_server_tool_calls_and_headers.py index e9478de2ece..5a4e039c56a 100644 --- a/tests/unit/proxy/mcp_server/test_mcp_server_tool_calls_and_headers.py +++ b/tests/unit/proxy/mcp_server/test_mcp_server_tool_calls_and_headers.py @@ -1241,6 +1241,7 @@ async def test_mcp_read_resource_success(): oauth2_headers=None, raw_headers=None, user_api_key_auth=user_api_key_auth, + scope_servers=(server,), ) mock_manager.read_resource_from_server.assert_awaited_once_with( server=server, @@ -1305,35 +1306,6 @@ async def test_read_resource_preserves_content_metadata(_mcp_request_ctx, kind, } -@pytest.mark.asyncio -async def test_mcp_read_resource_multiple_servers_error(): - try: - from litellm.proxy.mcp_server.server import mcp_read_resource - except ImportError: - pytest.skip("MCP server not available") - - user_api_key_auth = UserAPIKeyAuth(api_key="key", user_id="user") - - server_a = MagicMock() - server_b = MagicMock() - server_a.name = "server_a" - server_b.name = "server_b" - - with patch( - "litellm.proxy.mcp_server.operations._get_allowed_mcp_servers", - AsyncMock(return_value=[server_a, server_b]), - ) as mock_allowed: - with pytest.raises(HTTPException) as exc_info: - await mcp_read_resource( - url="https://example.com/resource", - user_api_key_auth=user_api_key_auth, - ) - - mock_allowed.assert_awaited_once() - assert exc_info.value.status_code == 400 - assert "Multiple MCP servers" in str(exc_info.value.detail) - - @pytest.mark.asyncio async def test_get_tools_from_mcp_servers_continues_when_one_server_fails(): """Test that _get_tools_from_mcp_servers continues when one server fails""" diff --git a/tests/unit/proxy/mcp_server/test_operations.py b/tests/unit/proxy/mcp_server/test_operations.py index 5394e3df97f..881344476df 100644 --- a/tests/unit/proxy/mcp_server/test_operations.py +++ b/tests/unit/proxy/mcp_server/test_operations.py @@ -4,10 +4,21 @@ from types import SimpleNamespace from typing import Final, Literal from unittest.mock import AsyncMock, patch +import httpx2 import pytest from mcp.server.context import ServerRequestContext -from mcp.types import GetPromptRequest, GetPromptRequestParams, GetPromptResult, ListToolsResult +from mcp.types import ( + GetPromptRequest, + GetPromptRequestParams, + GetPromptResult, + JSONRPCMessage, + JSONRPCRequest, + ListToolsResult, + ReadResourceResult, + TextResourceContents, +) from mcp.types import Tool as MCPTool +from pydantic import TypeAdapter import litellm from litellm.caching.dual_cache import DualCache @@ -24,6 +35,7 @@ from litellm.proxy.mcp_server.mcp_server_manager import ( ) from litellm.proxy.mcp_server.mcp_server_manager import ( ListedToolsCaller, + MCPServerManager, ) from litellm.proxy.mcp_server.operations import GatewayOperations, prepare_context from litellm.proxy.mcp_server.tool_registry import global_mcp_tool_registry @@ -46,6 +58,191 @@ async def _served_catalog_tool() -> str: return "ok" +@pytest.mark.asyncio +@pytest.mark.parametrize( + "case", + [ + "unique", + "ambiguous", + "missing", + "failure", + "slow_failure", + "blocked_stdio", + "scoped_bearer", + "legacy_delegates", + "rate_limited", + ], +) +async def test_resource_read_requires_unique_owner_from_complete_listings( + monkeypatch: pytest.MonkeyPatch, case: str +) -> None: + uri: Final = "ui://widget/view.html" + caller_bearer: Final = "Bearer caller-upstream-token" + reads: Final[asyncio.Queue[tuple[str, str | None]]] = asyncio.Queue() + listing_authorizations: Final[asyncio.Queue[str | None]] = asyncio.Queue() + slow_listing_cancelled: Final = asyncio.Event() + message: Final = TypeAdapter(JSONRPCMessage) + + async def respond(request: httpx2.Request) -> httpx2.Response: + if request.method == "DELETE": + return httpx2.Response(200) + payload: Final = message.validate_json(request.content) + if not isinstance(payload, JSONRPCRequest): + return httpx2.Response(202) + if payload.method == "initialize": + return httpx2.Response( + 200, + json={ + "jsonrpc": "2.0", + "id": payload.id, + "result": { + "protocolVersion": "2025-03-26", + "serverInfo": {"name": "upstream", "version": "1"}, + "capabilities": {"resources": {}}, + }, + }, + ) + if payload.method == "resources/read": + reads.put_nowait((request.url.host, request.headers.get("authorization"))) + return httpx2.Response( + 200, + json={ + "jsonrpc": "2.0", + "id": payload.id, + "result": { + "contents": [{"uri": uri, "text": request.url.host}], + }, + }, + ) + assert payload.method == "resources/list" + listing_authorizations.put_nowait(request.headers.get("authorization")) + if case == "failure" and request.url.host == "other.example": + return httpx2.Response(503, text="private upstream diagnostic") + if case == "slow_failure" and request.url.host == "owner.example": + return httpx2.Response(503, text="private upstream diagnostic") + if case == "slow_failure": + try: + await asyncio.Event().wait() + finally: + slow_listing_cancelled.set() + if not (payload.params or {}).get("cursor"): + return httpx2.Response( + 200, + json={ + "jsonrpc": "2.0", + "id": payload.id, + "result": { + "resources": [], + "nextCursor": "second-page", + }, + }, + ) + owns: Final = case not in ("missing", "rate_limited") and ( + request.url.host == "owner.example" or case == "ambiguous" + ) + return httpx2.Response( + 200, + json={ + "jsonrpc": "2.0", + "id": payload.id, + "result": { + "resources": [{"name": "widget", "uri": uri}] if owns else [], + }, + }, + ) + + manager: Final = MCPServerManager(upstream_transport=httpx2.MockTransport(respond)) + forwarding: Final = { + "owner": case in ("scoped_bearer", "blocked_stdio"), + "other": case == "scoped_bearer", + } + http_servers: Final = tuple( + MCPServer( + server_id=name, + name=name, + url=f"https://{name}.example/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2 + if case == "legacy_delegates" + else (MCPAuth.true_passthrough if forwarding[name] else None), + delegate_auth_to_upstream=case == "legacy_delegates", + allow_all_keys=True, + rpm=1, + ) + for name in ("owner", "other") + ) + stdio_server: Final = MCPServer( + server_id="local", + name="local", + transport=MCPTransport.stdio, + command="unused", + auth_type=MCPAuth.true_passthrough, + allow_all_keys=True, + rpm=1, + ) + servers: Final = (*http_servers, stdio_server) if case == "blocked_stdio" else http_servers + caller: Final = UserAPIKeyAuth(api_key=hash_token("sk-resource-read")) + proxy_logging: Final = _mcp_rate_limited_proxy_logging() + monkeypatch.setattr(operations, "global_mcp_server_manager", manager) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging) + monkeypatch.delenv("LITELLM_ENABLE_MCP_STDIO", raising=False) + for server in servers: + monkeypatch.setitem(manager.config_mcp_servers, server.server_id, server) + if case == "blocked_stdio": + await proxy_logging.enforce_mcp_server_rate_limits(caller, stdio_server) + + from mcp.shared.exceptions import MCPError + from mcp.types import INVALID_REQUEST, ReadResourceRequest, ReadResourceRequestParams + + async def read() -> object: + return await GatewayOperations().execute( + ReadResourceRequest(params=ReadResourceRequestParams(uri=uri)), + prepare_context( + caller, + oauth2_headers={"Authorization": caller_bearer}, + raw_headers={"authorization": caller_bearer, "x-litellm-api-key": "sk-resource-read"}, + ), + ) + + if case in ("unique", "blocked_stdio", "scoped_bearer", "legacy_delegates"): + result: Final = await read() + assert result == ReadResourceResult(contents=[TextResourceContents(uri=uri, text="owner.example")]) + assert reads.qsize() == 1 + host, authorization = reads.get_nowait() + assert host == "owner.example" + assert (authorization == caller_bearer) == (case == "blocked_stdio"), ( + "the caller bearer reaches the owner only when no other listable server also consumes it" + ) + listed_with: Final = tuple(listing_authorizations.get_nowait() for _ in range(listing_authorizations.qsize())) + assert listed_with + assert (caller_bearer in listed_with) == (case == "blocked_stdio"), ( + "owner lookups must not replay the caller bearer to servers that also consume it" + ) + return + if case == "rate_limited": + with pytest.raises(MCPError) as first: + await read() + assert first.value.error.code == INVALID_REQUEST + assert "not found" in first.value.error.message.lower() + with pytest.raises(MCPError) as limited: + await read() + assert limited.value.error.code == INVALID_REQUEST + assert "rate limit" in limited.value.error.message.lower() + assert reads.empty() + return + with pytest.raises(MCPError) as rejected: + await asyncio.wait_for(read(), timeout=30) + assert rejected.value.error.code == INVALID_REQUEST, "owner failures must reach MCP clients as MCPError, not 500" + assert {"ambiguous": "ambiguous", "missing": "not found", "failure": "list resources", "slow_failure": "owner"}[ + case + ] in rejected.value.error.message.lower() + assert reads.empty() + assert "private upstream diagnostic" not in rejected.value.error.message + if case == "slow_failure": + await asyncio.wait_for(slow_listing_cancelled.wait(), timeout=30) + + def _listing_client() -> SimpleNamespace: return SimpleNamespace(discovery_auth_fingerprint=AsyncMock(return_value="test-credential-hash")) @@ -932,7 +1129,10 @@ async def test_catalog_helpers_preserve_context_results_and_failure_policy(metho patch.object(operations.global_mcp_server_manager, manager_method, upstream), ): if collection is None and state != "success": - expected_error = RuntimeError if state == "upstream_failure" else HTTPException + from mcp.shared.exceptions import MCPError + + access_error = MCPError if method == "resources/read" else HTTPException + expected_error = RuntimeError if state == "upstream_failure" else access_error with pytest.raises(expected_error): await getattr(server, handler_name)(ctx, operation.params) else: