diff --git a/litellm/constants.py b/litellm/constants.py index 8e0a5cfe50f..499a7fb1e1b 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -121,6 +121,25 @@ MCP_TOOL_LISTING_TIMEOUT = float(os.getenv("LITELLM_MCP_TOOL_LISTING_TIMEOUT", " MCP_METADATA_TIMEOUT = float(os.getenv("LITELLM_MCP_METADATA_TIMEOUT", "10.0")) MCP_HEALTH_CHECK_TIMEOUT = float(os.getenv("LITELLM_MCP_HEALTH_CHECK_TIMEOUT", "10.0")) +MCP_RESPONSE_HEADERS_META_KEY = "ai.litellm/responseHeaders" +MCP_RESPONSE_HEADER_DENYLIST = frozenset( + { + "authorization", + "proxy-authorization", + "www-authenticate", + "proxy-authenticate", + "set-cookie", + "cookie", + "mcp-session-id", + "connection", + "keep-alive", + "transfer-encoding", + "upgrade", + "te", + "trailer", + } +) + # Allowlist of commands permitted for MCP stdio transport. # Prevents arbitrary command execution via /mcp-rest/test/* endpoints or server creation. # Note: allowlisted runtimes can still execute code via args (e.g. python -c "..."). diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index da711463a44..b0374a014a7 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -4,6 +4,7 @@ LiteLLM Proxy uses this MCP Client to connnect to other MCP servers. import asyncio import base64 +import json import os from typing import ( Any, @@ -41,7 +42,12 @@ from mcp.types import ( from mcp.types import Tool as MCPTool from pydantic import AnyUrl from litellm._logging import verbose_logger -from litellm.constants import MCP_CLIENT_TIMEOUT, MCP_NPM_CACHE_DIR +from litellm.constants import ( + MCP_CLIENT_TIMEOUT, + MCP_NPM_CACHE_DIR, + MCP_RESPONSE_HEADER_DENYLIST, + MCP_RESPONSE_HEADERS_META_KEY, +) from litellm.llms.custom_httpx.http_handler import get_ssl_configuration from litellm.types.llms.custom_http import VerifyTypes from litellm.types.mcp import ( @@ -58,6 +64,35 @@ def to_basic_auth(auth_value: str) -> str: return base64.b64encode(auth_value.encode("utf-8")).decode() +def normalize_allowed_response_headers(names: list[str] | None) -> frozenset[str]: + """Lowercase the configured upstream response-header allowlist, dropping blanks and denylisted names. + + Normalizing in one place keeps a single notion of "blank" (None, "", whitespace) and guarantees a + denylisted header can never be surfaced to the caller even when an admin configures it explicitly. + """ + if not names: + return frozenset() + return frozenset( + stripped + for stripped in (name.strip().lower() for name in names if isinstance(name, str)) + if stripped and stripped not in MCP_RESPONSE_HEADER_DENYLIST + ) + + +def _is_tool_call_request(request: httpx.Request) -> bool: + """True when this HTTP request carries the JSON-RPC ``tools/call`` that produced the response. + + A streamable-HTTP session also POSTs ``initialize`` and notifications, and may DELETE the session + on close, so the tool call is identified by its JSON-RPC method rather than by being the last + response observed on the connection. + """ + try: + payload = json.loads(request.content) + except (httpx.RequestNotRead, ValueError, TypeError): + return False + return isinstance(payload, dict) and payload.get("method") == "tools/call" + + def _strip_header_whitespace(headers: Dict[str, str]) -> Dict[str, str]: return { (key.strip() if isinstance(key, str) else key): (value.strip() if isinstance(value, str) else value) @@ -217,6 +252,7 @@ class MCPClient: sampling_callback: Optional[Callable] = None, elicitation_callback: Optional[Callable] = None, logging_callback: Optional[Callable] = None, + allowed_response_headers: list[str] | None = None, ): self.server_url: str = server_url self.transport_type: MCPTransport = transport_type @@ -234,10 +270,35 @@ class MCPClient: self._sampling_callback: Optional[Callable] = sampling_callback self._elicitation_callback: Optional[Callable] = elicitation_callback self._logging_callback: Optional[Callable] = logging_callback + self._allowed_response_headers: frozenset[str] = normalize_allowed_response_headers(allowed_response_headers) + self._tool_call_response_headers: dict[str, str] | None = None # handle the basic auth value if provided if auth_value: self.update_auth_value(auth_value) + def _response_event_hooks( + self, + ) -> dict[str, list[Callable[[httpx.Response], Awaitable[None]]]] | None: + """An httpx response hook recording the allowlisted upstream ``tools/call`` headers, when configured. + + This is the only seam that sees the upstream HTTP response: the MCP SDK's transport keeps just + the session id and content-type and hands back a headerless ``CallToolResult``. Reading only + ``response.headers`` leaves the (possibly streamed) body untouched. + """ + if not self._allowed_response_headers: + return None + + async def _capture(response: httpx.Response) -> None: + if not _is_tool_call_request(response.request): + return + self._tool_call_response_headers = { + name.lower(): value + for name, value in response.headers.items() + if name.lower() in self._allowed_response_headers + } + + return {"response": [_capture]} + def _create_transport_context( self, ) -> Tuple[Any, Optional[httpx.AsyncClient]]: @@ -278,6 +339,7 @@ class MCPClient: http_client = httpx_client_factory( headers=headers, timeout=httpx.Timeout(self.timeout), + event_hooks=self._response_event_hooks(), ) transport_ctx = streamable_http_client( url=self.server_url, @@ -396,6 +458,7 @@ class MCPClient: http_client: Optional[httpx.AsyncClient] = None try: self._last_initialize_instructions = None + self._tool_call_response_headers = None transport_ctx, http_client = self._create_transport_context() return await self._execute_session_operation(transport_ctx, operation) except Exception: @@ -465,6 +528,7 @@ class MCPClient: headers: Optional[Dict[str, str]] = None, timeout: Optional[httpx.Timeout] = None, auth: Optional[httpx.Auth] = None, + event_hooks: dict[str, list[Callable[[httpx.Response], Awaitable[None]]]] | None = None, ) -> httpx.AsyncClient: """Create an httpx.AsyncClient with LiteLLM's SSL configuration.""" # Get unified SSL configuration using the same logic as http_handler.py @@ -481,6 +545,7 @@ class MCPClient: auth=effective_auth, verify=ssl_config, follow_redirects=True, + event_hooks=event_hooks, ) return factory @@ -537,6 +602,20 @@ class MCPClient: # Return empty list instead of raising to allow graceful degradation return [] + def _with_response_headers_meta(self, result: MCPCallToolResult) -> MCPCallToolResult: + """Surface the captured upstream response headers on the result's MCP ``_meta``. + + ``_meta`` is the protocol's reserved slot for out-of-band metadata and the only channel that + reaches the caller, since the gateway re-serializes the result over its own transport and + commits its HTTP headers before the tool runs. The key carries litellm's reverse-DNS prefix + because the spec reserves ``mcp``/``modelcontextprotocol`` prefixes for protocol use. + """ + headers = self._tool_call_response_headers + if not headers: + return result + merged_meta = {**(result.meta or {}), MCP_RESPONSE_HEADERS_META_KEY: headers} + return result.model_copy(update={"meta": merged_meta}) + @staticmethod def error_tool_result(exc: Exception) -> MCPCallToolResult: """The error result ``call_tool`` returns when it swallows a failure (no re-execution).""" @@ -586,7 +665,7 @@ class MCPClient: try: tool_result = await self.run_with_session(_call_tool_operation, quiet_on_error=raise_on_error) verbose_logger.info(f"MCP client tool call '{call_tool_request_params.name}' completed successfully") - return tool_result + return self._with_response_headers_meta(tool_result) except asyncio.CancelledError: verbose_logger.warning(f"MCP client tool call timed out after {self.timeout}s for {self.server_url}") raise diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 115ff2e492c..8d3c92c1f9d 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -1323,6 +1323,7 @@ class MCPServerManager: allowed_params=server_config.get("allowed_params", None), access_groups=server_config.get("access_groups", None), static_headers=server_config.get("static_headers", None), + allowed_response_headers=server_config.get("allowed_response_headers", None), env_vars=server_config.get("env_vars", None), allow_all_keys=bool(server_config.get("allow_all_keys", False)), available_on_public_internet=bool(server_config.get("available_on_public_internet", True)), @@ -2864,6 +2865,7 @@ class MCPServerManager: resolved_auth=resolved_auth, sampling_callback=sampling_cb, elicitation_callback=elicitation_cb, + allowed_response_headers=server.allowed_response_headers, ) # Create SigV4 auth if configured @@ -2889,6 +2891,7 @@ class MCPServerManager: aws_auth=aws_auth, sampling_callback=sampling_cb, elicitation_callback=elicitation_cb, + allowed_response_headers=server.allowed_response_headers, ) async def _get_tools_from_server( diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index d0d8cc4cb28..0aec6566f9c 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -58,6 +58,13 @@ class MCPServer(BaseModel): tool_name_to_description: Optional[Dict[str, str]] = None allowed_params: Optional[Dict[str, List[str]]] = None # map of tool names to allowed parameter lists static_headers: Optional[Dict[str, str]] = None # static headers to forward to the MCP server + allowed_response_headers: Optional[List[str]] = None + """Upstream ``tools/call`` response headers to surface to the caller on the MCP result's ``_meta``. + + Opt-in allowlist of header names; empty or unset forwards nothing. The gateway answers the caller + over a stream whose HTTP headers are already committed before the tool runs, so these are relayed + as protocol metadata rather than as HTTP response headers. Credential, cookie, upstream-session and + hop-by-hop headers are never forwarded (see ``MCP_RESPONSE_HEADER_DENYLIST``).""" # Admin-configured env vars. Each entry is {name, value, scope, description}. # scope=="global" values are interpolated into static_headers using ${NAME}. # scope=="user" values must be supplied per-user. diff --git a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py index c97856598f1..f5757e521fb 100644 --- a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py +++ b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py @@ -9,10 +9,15 @@ import pytest # Add the parent directory to the path so we can import litellm sys.path.insert(0, "../../../") +from mcp.types import CallToolRequestParams, CallToolResult, TextContent + import litellm.experimental_mcp_client.client as mcp_client_module +from litellm.constants import MCP_RESPONSE_HEADERS_META_KEY from litellm.experimental_mcp_client.client import ( MCPClient, _first_non_cancelled_cause, + _is_tool_call_request, + normalize_allowed_response_headers, ) from litellm.types.mcp import MCPAuth, MCPStdioConfig, MCPTransport @@ -695,3 +700,191 @@ async def test_run_with_session_quiet_on_error_demotes_warning_to_debug(): assert any("run_with_session failed" in m for m in warning_msgs), ( "the default path must keep the operator-visible warning" ) + + +def _jsonrpc_response(method: str, headers: dict) -> httpx.Response: + """An upstream HTTP response to the given JSON-RPC method, carrying ``headers``.""" + request = httpx.Request("POST", "http://upstream/mcp", json={"jsonrpc": "2.0", "id": 1, "method": method}) + return httpx.Response(200, headers=headers, request=request) + + +class TestNormalizeAllowedResponseHeaders: + """The allowlist normalizer: one notion of blank, case, and never-forward.""" + + def test_unset_forwards_nothing(self): + assert normalize_allowed_response_headers(None) == frozenset() + assert normalize_allowed_response_headers([]) == frozenset() + + def test_lowercases_and_strips(self): + assert normalize_allowed_response_headers([" X-Example-Header ", "X-Request-Id"]) == frozenset( + {"x-example-header", "x-request-id"} + ) + + def test_blank_entries_dropped(self): + assert normalize_allowed_response_headers(["", " ", "\t", "X-Ok"]) == frozenset({"x-ok"}) + + def test_denylisted_headers_never_forwarded_even_when_configured(self): + """An admin cannot opt a credential, cookie, upstream session or hop-by-hop header into the result.""" + configured = [ + "Authorization", + "Proxy-Authorization", + "Set-Cookie", + "Cookie", + "MCP-Session-Id", + "WWW-Authenticate", + "Connection", + "Transfer-Encoding", + "X-Keep", + ] + assert normalize_allowed_response_headers(configured) == frozenset({"x-keep"}) + + +class TestIsToolCallRequest: + """Only the tools/call response may be captured, never a sibling message on the same connection.""" + + def test_tool_call_matches(self): + assert _is_tool_call_request(_jsonrpc_response("tools/call", {}).request) is True + + @pytest.mark.parametrize("method", ["initialize", "tools/list", "notifications/initialized"]) + def test_other_methods_do_not_match(self, method): + assert _is_tool_call_request(_jsonrpc_response(method, {}).request) is False + + def test_non_json_body_does_not_match(self): + assert _is_tool_call_request(httpx.Request("DELETE", "http://upstream/mcp")) is False + + +class TestMCPResponseHeaderCapture: + """Capturing the upstream tools/call response headers off the httpx seam.""" + + def test_no_allowlist_installs_no_hook(self): + """Unconfigured servers keep the stock client: no hook, no capture, no behavior change.""" + client = MCPClient(server_url="http://upstream/mcp", transport_type=MCPTransport.http) + assert client._response_event_hooks() is None + + @pytest.mark.asyncio + async def test_captures_only_allowlisted_headers_from_tool_call(self): + client = MCPClient( + server_url="http://upstream/mcp", + transport_type=MCPTransport.http, + allowed_response_headers=["X-Example-Header"], + ) + hook = client._response_event_hooks()["response"][0] + + await hook(_jsonrpc_response("tools/call", {"X-Example-Header": "hello", "X-Other": "dropped"})) + + assert client._tool_call_response_headers == {"x-example-header": "hello"} + + @pytest.mark.asyncio + async def test_ignores_non_tool_call_responses(self): + """initialize and the session teardown share the connection; taking the last response would be wrong.""" + client = MCPClient( + server_url="http://upstream/mcp", + transport_type=MCPTransport.http, + allowed_response_headers=["X-Example-Header"], + ) + hook = client._response_event_hooks()["response"][0] + + await hook(_jsonrpc_response("initialize", {"X-Example-Header": "from-initialize"})) + + assert client._tool_call_response_headers is None + + @pytest.mark.asyncio + async def test_tool_call_headers_survive_a_later_sibling_response(self): + client = MCPClient( + server_url="http://upstream/mcp", + transport_type=MCPTransport.http, + allowed_response_headers=["X-Example-Header"], + ) + hook = client._response_event_hooks()["response"][0] + + await hook(_jsonrpc_response("tools/call", {"X-Example-Header": "hello"})) + await hook(_jsonrpc_response("tools/list", {"X-Example-Header": "later"})) + + assert client._tool_call_response_headers == {"x-example-header": "hello"} + + +class TestWithResponseHeadersMeta: + """Surfacing captured headers on the result's MCP ``_meta``.""" + + def _client(self): + return MCPClient( + server_url="http://upstream/mcp", + transport_type=MCPTransport.http, + allowed_response_headers=["X-Example-Header"], + ) + + def test_nothing_captured_leaves_result_untouched(self): + client = self._client() + result = CallToolResult(content=[TextContent(type="text", text="ok")]) + + assert client._with_response_headers_meta(result).meta is None + + def test_captured_headers_land_under_the_litellm_meta_key(self): + client = self._client() + client._tool_call_response_headers = {"x-example-header": "hello"} + result = CallToolResult(content=[TextContent(type="text", text="ok")]) + + assert client._with_response_headers_meta(result).meta == { + MCP_RESPONSE_HEADERS_META_KEY: {"x-example-header": "hello"} + } + + def test_preserves_meta_the_upstream_server_already_set(self): + """The upstream's own ``_meta`` is the caller's data; header surfacing must merge, not clobber.""" + client = self._client() + client._tool_call_response_headers = {"x-example-header": "hello"} + result = CallToolResult(content=[TextContent(type="text", text="ok")], _meta={"upstream/trace": "abc123"}) + + assert client._with_response_headers_meta(result).meta == { + "upstream/trace": "abc123", + MCP_RESPONSE_HEADERS_META_KEY: {"x-example-header": "hello"}, + } + + def test_does_not_mutate_the_original_result(self): + client = self._client() + client._tool_call_response_headers = {"x-example-header": "hello"} + result = CallToolResult(content=[TextContent(type="text", text="ok")]) + + client._with_response_headers_meta(result) + + assert result.meta is None + + @pytest.mark.asyncio + async def test_call_tool_surfaces_captured_headers_on_its_result(self): + """The wiring that makes the feature reachable: call_tool must surface what the session captured.""" + client = self._client() + + async def _session_that_captures(operation, *, quiet_on_error=False): + client._tool_call_response_headers = {"x-example-header": "hello"} + return CallToolResult(content=[TextContent(type="text", text="ok")]) + + with patch.object(client, "run_with_session", side_effect=_session_that_captures): + result = await client.call_tool(CallToolRequestParams(name="echo", arguments={})) + + assert result.meta == {MCP_RESPONSE_HEADERS_META_KEY: {"x-example-header": "hello"}} + + @pytest.mark.asyncio + async def test_call_tool_leaves_result_alone_when_nothing_captured(self): + client = self._client() + + async def _session_without_capture(operation, *, quiet_on_error=False): + return CallToolResult(content=[TextContent(type="text", text="ok")]) + + with patch.object(client, "run_with_session", side_effect=_session_without_capture): + result = await client.call_tool(CallToolRequestParams(name="echo", arguments={})) + + assert result.meta is None + + @pytest.mark.asyncio + async def test_opening_a_session_clears_a_previous_calls_headers(self): + """Without this reset a reused client would replay the prior call's headers onto a later result.""" + client = self._client() + client._tool_call_response_headers = {"x-example-header": "stale"} + + async def _op(session): + return "unused" + + with patch.object(client, "_create_transport_context", side_effect=RuntimeError("boom")): + with pytest.raises(RuntimeError): + await client.run_with_session(_op) + + assert client._tool_call_response_headers is None