diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 554b6ea952e..95a3806e8ad 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -1190,7 +1190,7 @@ class MCPRequestHandler: DEPRECATED: This method is deprecated in favor of server-specific auth headers using the format x-mcp-{{server_alias}}-{{header_name}} instead. """ - mcp_client_side_auth_header_name: Final[str] = MCPRequestHandler._get_mcp_client_side_auth_header_name() + mcp_client_side_auth_header_name: Final[str] = MCPRequestHandler.get_mcp_client_side_auth_header_name() auth_header: Final = headers.get(mcp_client_side_auth_header_name) if auth_header: verbose_logger.warning( @@ -1265,7 +1265,7 @@ class MCPRequestHandler: return oauth2_headers @staticmethod - def _get_mcp_client_side_auth_header_name() -> str: + def get_mcp_client_side_auth_header_name() -> str: """ Get the header name used to pass the MCP auth header to the MCP server diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index c8ff6e262d2..a1adda2bc95 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -118,6 +118,7 @@ from litellm.proxy._experimental.mcp_server.utils import ( is_short_mcp_tool_prefix_enabled, iter_known_server_prefixes, iter_known_tool_name_spellings, + logging_safe_mcp_headers, match_known_server_prefix, match_known_tool_name, merge_mcp_headers, @@ -4603,6 +4604,7 @@ class MCPServerManager: ), "user_api_key_hash": (getattr(user_api_key_auth, "api_key_hash", None) if user_api_key_auth else None), "incoming_bearer_token": incoming_bearer_token, + "headers": logging_safe_mcp_headers(raw_headers), } # Create MCP request object for processing diff --git a/litellm/proxy/_experimental/mcp_server/sampling_handler.py b/litellm/proxy/_experimental/mcp_server/sampling_handler.py index a9a3367cd93..125dc3d773d 100644 --- a/litellm/proxy/_experimental/mcp_server/sampling_handler.py +++ b/litellm/proxy/_experimental/mcp_server/sampling_handler.py @@ -1042,100 +1042,15 @@ def _build_sampling_request( raw_headers: dict[str, str] | None = None, client_ip: str | None = None, ) -> "Request": - """Build a synthetic FastAPI Request for sampling sub-calls. + """The synthetic FastAPI Request for sampling sub-calls, carrying the original + MCP connection's headers and client IP.""" + from litellm.proxy._experimental.mcp_server.utils import build_synthetic_mcp_request - Converts the original MCP connection's HTTP headers into ASGI - scope format so that ``add_litellm_data_to_request`` can apply - header-dependent guardrails, tag-based routing, trace correlation, - and ``forward_llm_provider_auth_headers``. - - Key fields populated: - - **headers**: All original HTTP headers are forwarded (except - hop-by-hop: content-length, transfer-encoding). This ensures - ``traceparent``, ``authorization``, ``user-agent``, and - ``x-litellm-api-key`` are visible to pre-call utils. - - **client**: The ASGI ``(host, port)`` tuple so that - ``request.client.host`` returns the real client IP for - IP-based routing and guardrails. - - **server**: Derived from the running proxy's ``server_host`` - / ``server_port`` when available, avoiding the misleading - ``127.0.0.1:0`` placeholder. - - **x-forwarded-for**: Injected from ``client_ip`` if the - original headers don't already carry it, as a fallback for - IP attribution. - """ - from fastapi import Request - - # --- Build ASGI headers --- - _scope_headers: Final[list[tuple[bytes, bytes]]] = [(b"content-type", b"application/json")] - # Hop-by-hop headers that must NOT be forwarded into the - # synthetic request (they describe the original HTTP framing, - # not the logical request). - _HOP_BY_HOP: Final = frozenset( - { - "content-length", - "transfer-encoding", - "connection", - "keep-alive", - "upgrade", - "te", - "trailer", - } + return build_synthetic_mcp_request( + path="/mcp/sampling/createMessage", + raw_headers=raw_headers, + client_ip=client_ip, ) - if raw_headers: - for hdr_name, hdr_value in raw_headers.items(): - _key = hdr_name.lower() - # Skip content-type (already set), x-forwarded-for (use resolved - # client_ip instead to prevent spoofing), and hop-by-hop headers - if _key in {"content-type", "x-forwarded-for"} or _key in _HOP_BY_HOP: - continue - _scope_headers.append( - ( - _key.encode("latin-1", errors="replace"), - hdr_value.encode("utf-8"), - ) - ) - - # Inject x-forwarded-for from captured client_ip if the - # original headers don't already carry it - if client_ip and not any(h[0] == b"x-forwarded-for" for h in _scope_headers): - _scope_headers.append((b"x-forwarded-for", client_ip.encode("utf-8"))) - - # --- Derive server (host, port) from the running proxy --- - _server_host = "127.0.0.1" - _server_port = 4000 # LiteLLM default - try: - from litellm.proxy import proxy_server - - _proxy_host: Final[str | None] = getattr(proxy_server, "server_host", None) - _proxy_port: Final[str | int | None] = getattr(proxy_server, "server_port", None) - - if _proxy_host: - _server_host = str(_proxy_host) - if _proxy_port: - _server_port = int(_proxy_port) - except (ImportError, AttributeError, TypeError, ValueError): - pass - - # --- Build ASGI client tuple for request.client.host --- - _client_tuple = None - if client_ip: - _client_tuple = (client_ip, 0) - - scope: Final[dict[str, object]] = { - "type": "http", - "method": "POST", - "path": "/mcp/sampling/createMessage", - "scheme": "http", - "server": (_server_host, _server_port), - "query_string": b"", - "root_path": "", - "headers": _scope_headers, - } - if _client_tuple is not None: - scope["client"] = _client_tuple - - return Request(scope=scope) async def _build_completion_kwargs( diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 49a1f1314f0..f237529b319 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -58,9 +58,11 @@ from litellm.proxy._experimental.mcp_server.utils import ( LITELLM_MCP_SERVER_VERSION, MCPMissingUserEnvVarsError, add_server_prefix_to_name, + build_synthetic_mcp_request, extract_mcp_tool_result_error_message, get_server_prefix, iter_known_server_prefixes, + logging_safe_mcp_headers, match_known_tool_name, ) from litellm.proxy._types import ( @@ -860,11 +862,11 @@ if MCP_AVAILABLE: name: str, arguments: dict[str, object], user_api_key_auth: UserAPIKeyAuth, + raw_headers: Mapping[str, str] | None = None, + client_ip: str | None = None, ) -> LiteLLMLoggingObj | None: """Run the pre-call pipeline (guardrails + logging setup) for a virtual mcp_tool_call so the SSE path spend-logs like the REST path.""" - from fastapi import Request - from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, ) @@ -874,13 +876,10 @@ if MCP_AVAILABLE: proxy_logging_obj, ) - request: Final = Request( - scope={ - "type": "http", - "method": "POST", - "path": "/mcp/tools/call", - "headers": [(b"content-type", b"application/json")], - } + request: Final = build_synthetic_mcp_request( + path="/mcp/tools/call", + raw_headers=raw_headers, + client_ip=client_ip, ) _, virtual_logging_obj = await ProxyBaseLLMRequestProcessing( data={"name": name, "arguments": arguments} @@ -952,7 +951,11 @@ if MCP_AVAILABLE: assert user_api_key_auth is not None # guaranteed by the flag check above virtual_logging_obj: Final = await _build_virtual_call_logging_obj( - name=name, arguments=args, user_api_key_auth=user_api_key_auth + name=name, + arguments=args, + user_api_key_auth=user_api_key_auth, + raw_headers=raw_headers, + client_ip=client_ip, ) return await handle_mcp_tool_call( tool_name=args.get("tool_name", ""), @@ -979,7 +982,6 @@ if MCP_AVAILABLE: Raises: HTTPException: If tool not found or arguments missing """ - from fastapi import Request from mcp.server.lowlevel.server import request_ctx from mcp.types import CallToolResult @@ -1041,13 +1043,10 @@ if MCP_AVAILABLE: body_data["litellm_trace_id"] = chain_id body_data["litellm_session_id"] = chain_id - request: Final = Request( - scope={ - "type": "http", - "method": "POST", - "path": "/mcp/tools/call", - "headers": [(b"content-type", b"application/json")], - } + request: Final = build_synthetic_mcp_request( + path="/mcp/tools/call", + raw_headers=raw_headers, + client_ip=_client_ip, ) if user_api_key_auth is not None: data = await add_litellm_data_to_request( @@ -1905,6 +1904,7 @@ if MCP_AVAILABLE: "litellm_trace_id": effective_litellm_trace_id, "metadata": { "spend_logs_metadata": spend_logs_metadata, + "headers": logging_safe_mcp_headers(raw_headers), **({"tags": request_tags} if request_tags else {}), }, # Provide a small input payload for standard logging diff --git a/litellm/proxy/_experimental/mcp_server/utils.py b/litellm/proxy/_experimental/mcp_server/utils.py index 5534b8b5a3c..4cf84dd0725 100644 --- a/litellm/proxy/_experimental/mcp_server/utils.py +++ b/litellm/proxy/_experimental/mcp_server/utils.py @@ -7,6 +7,7 @@ import importlib import json import os import re +import typing from collections.abc import Iterable, Iterator, Mapping, MutableMapping, MutableSequence from collections.abc import Set as AbstractSet from typing import Any, Final, Protocol @@ -14,6 +15,9 @@ from urllib.parse import quote from litellm.types.mcp_server.mcp_server_manager import MCPServer +if typing.TYPE_CHECKING: + from fastapi import Request + class _McpServerLike(Protocol): @property @@ -862,3 +866,146 @@ def set_mcp_tool_result_structured_content(result: object, value: object) -> boo return True except (AttributeError, TypeError, ValueError): return False + + +_HOP_BY_HOP_HEADERS: Final = frozenset( + { + "content-length", + "transfer-encoding", + "connection", + "keep-alive", + "upgrade", + "te", + "trailer", + } +) + +_SYNTHETIC_REQUEST_EXCLUDED_HEADERS: Final = _HOP_BY_HOP_HEADERS | frozenset({"content-type", "x-forwarded-for"}) + +_SYNTHETIC_REQUEST_SERVER: Final = ("127.0.0.1", 4000) + +_MCP_SERVER_AUTH_HEADER_PREFIX: Final = "x-mcp-" + + +def _custom_litellm_key_header_name() -> str | None: + """``general_settings.litellm_key_header_name``, the deployment's custom header name for + the proxy virtual key, so it is stripped from observability copies like the standard ones.""" + try: + from litellm.proxy.proxy_server import general_settings + except ImportError: + return None + return general_settings.get("litellm_key_header_name") if general_settings else None + + +def _mcp_client_side_auth_header_name() -> str: + """The header name the client passes the upstream MCP credential in, falling back to the + default when ``general_settings`` is unavailable (the SDK, outside a running proxy).""" + from .auth.user_api_key_auth_mcp import MCPRequestHandler + + try: + return MCPRequestHandler.get_mcp_client_side_auth_header_name() + except ImportError: + return MCPRequestHandler.LITELLM_MCP_AUTH_HEADER_NAME + + +def _upstream_credential_headers(header_names: Iterable[str]) -> frozenset[str]: + """Lowercased names of the headers in ``header_names`` that carry an upstream MCP + credential rather than request context: the configured client side auth header and + the per-server ``x-mcp-{alias}-{header}`` family. ``clean_headers`` only knows the + credential headers of the chat completions path, so these are dropped on top of it. + """ + from .auth.user_api_key_auth_mcp import MCPRequestHandler + + non_credential: Final = frozenset( + { + MCPRequestHandler.LITELLM_MCP_SERVERS_HEADER_NAME.lower(), + MCPRequestHandler.LITELLM_MCP_ACCESS_GROUPS_HEADER_NAME.lower(), + } + ) + client_side_auth: Final = _mcp_client_side_auth_header_name().lower() + return frozenset( + name + for name in (raw_name.lower() for raw_name in header_names) + if name == client_side_auth or (name.startswith(_MCP_SERVER_AUTH_HEADER_PREFIX) and name not in non_credential) + ) + + +def build_synthetic_mcp_request( + *, + path: str, + raw_headers: Mapping[str, str] | None = None, + client_ip: str | None = None, +) -> "Request": + """A synthetic FastAPI ``Request`` carrying the MCP connection's HTTP headers. + + The MCP protocol transports do not hand a per-call ``Request`` to the tool + handlers, so one is reconstructed from the connection's ``raw_headers``. That + lets ``add_litellm_data_to_request`` derive ``metadata.headers``, + ``proxy_server_request``, header-based tags, guardrails and trace correlation + exactly as on the chat completions path. Hop-by-hop headers describe the + original HTTP framing rather than the logical request, so they are dropped, and + ``x-forwarded-for`` comes from the resolved ``client_ip`` to avoid spoofing. Upstream + MCP credentials and the deployment's proxy key header, including a custom + ``litellm_key_header_name``, are dropped so they cannot reach a callback or a guardrail + through the derived metadata even when a caller omits ``general_settings``. + """ + from fastapi import Request + + custom_key_header: Final = _custom_litellm_key_header_name() + excluded: Final = ( + _SYNTHETIC_REQUEST_EXCLUDED_HEADERS + | _upstream_credential_headers(raw_headers.keys() if raw_headers else ()) + | (frozenset({custom_key_header.lower()}) if custom_key_header else frozenset()) + ) + forwarded: Final = tuple( + ( + name.lower().encode("latin-1", errors="replace"), + value.encode("utf-8", errors="replace"), + ) + for name, value in (raw_headers.items() if raw_headers else ()) + if name.lower() not in excluded + ) + xff: Final = ((b"x-forwarded-for", client_ip.encode("utf-8")),) if client_ip else () + return Request( + scope={ + "type": "http", + "method": "POST", + "path": path, + "scheme": "http", + "server": _SYNTHETIC_REQUEST_SERVER, + "query_string": b"", + "root_path": "", + "headers": ((b"content-type", b"application/json"), *forwarded, *xff), + **({"client": (client_ip, 0)} if client_ip else {}), + } + ) + + +def logging_safe_mcp_headers(raw_headers: Mapping[str, str] | None) -> Mapping[str, str]: + """The MCP request's client headers, sanitized the way the chat completions path + sanitizes them before they reach a logging callback or a guardrail: proxy key + headers stripped, including the custom key header name the deployment configured, + upstream MCP credentials dropped, and credential-bearing values masked. + + Client-controlled behaviour flags (``litellm-disable-message-redaction``) are dropped + too: these headers are read back out of the metadata to change proxy behaviour, so + leaving one in place would let any MCP client turn off the redaction an admin + configured. This path carries no key or team object to authorize an opt-out with, so + it always strips them.""" + from starlette.datastructures import Headers + + from litellm.proxy.litellm_pre_call_utils import ( + UNTRUSTED_REQUEST_HEADER_CONTROL_FIELDS, + clean_headers, + redact_credential_headers, + ) + + excluded: Final = ( + _upstream_credential_headers(raw_headers.keys() if raw_headers else ()) + | UNTRUSTED_REQUEST_HEADER_CONTROL_FIELDS + ) + cleaned: Final = clean_headers( + Headers(raw_headers), + litellm_key_header_name=_custom_litellm_key_header_name(), + ) + return redact_credential_headers({name: value for name, value in cleaned.items() if name.lower() not in excluded}) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 251ed1feb10..0a5626ba0a7 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -274,7 +274,7 @@ _UNTRUSTED_METADATA_CONTROL_FIELDS: Final = ( PRE_CALL_EXECUTED_GUARDRAILS_KEY, ) -_UNTRUSTED_REQUEST_HEADER_CONTROL_FIELDS: Final = frozenset( +UNTRUSTED_REQUEST_HEADER_CONTROL_FIELDS: Final = frozenset( { "litellm-disable-message-redaction", } @@ -355,7 +355,7 @@ def _strip_untrusted_request_header_controls( return for header_name in list(headers.keys()): - if isinstance(header_name, str) and header_name.lower() in _UNTRUSTED_REQUEST_HEADER_CONTROL_FIELDS: + if isinstance(header_name, str) and header_name.lower() in UNTRUSTED_REQUEST_HEADER_CONTROL_FIELDS: if allow_client_message_redaction_opt_out: continue headers.pop(header_name, None) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index ef0376ade84..d3ca2fa64ed 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -679,6 +679,7 @@ class ProxyLogging: # (e.g. MCPJWTSigner) to independently verify the caller's identity # before re-signing an outbound token (FR-5 verify+re-sign). "incoming_bearer_token": kwargs.get("incoming_bearer_token"), + "metadata": {"headers": kwargs.get("headers") or {}}, } return synthetic_data diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index c6e17502e5d..56818717c09 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -11,6 +11,7 @@ from litellm._logging import verbose_logger from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._experimental.mcp_server.utils import ( + logging_safe_mcp_headers, split_server_prefix_from_name, strip_known_server_prefix, ) @@ -653,6 +654,7 @@ class LiteLLM_Proxy_MCP_Handler: tool_results: Final[list[MCPToolResult]] = [] tool_call_id: str | None = None rules_obj: Final = Rules() + logging_safe_headers: Final = logging_safe_mcp_headers(raw_headers) for tool_call in tool_calls: logging_request_data: dict[str, object] = {} tool_name: str | None = None @@ -697,6 +699,7 @@ class LiteLLM_Proxy_MCP_Handler: "tool_call_id": tool_call_id, "tool_name": sanitized_tool_name, "server_name": server_name, + "headers": logging_safe_headers, } logging_request_data = { "model": f"MCP: {tool_name}", @@ -708,7 +711,7 @@ class LiteLLM_Proxy_MCP_Handler: "proxy_server_request": { "url": "/mcp/tools/call", "method": "POST", - "headers": {}, + "headers": logging_safe_headers, "body": { "name": sanitized_tool_name, "arguments": parsed_arguments, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 0b95a497882..52dc91ce24d 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -2874,7 +2874,7 @@ class TestMCPCustomHeaderName: mock_general_settings.get.return_value = general_setting # Call the method - result = MCPRequestHandler._get_mcp_client_side_auth_header_name() + result = MCPRequestHandler.get_mcp_client_side_auth_header_name() # Assert the result assert result == expected_header_name @@ -2938,7 +2938,7 @@ class TestMCPCustomHeaderName: # Mock the header name method with patch.object( MCPRequestHandler, - "_get_mcp_client_side_auth_header_name", + "get_mcp_client_side_auth_header_name", return_value=custom_header_name, ): # Create headers from the test data @@ -2963,7 +2963,7 @@ class TestMCPCustomHeaderName: # Mock the custom header name with patch.object( MCPRequestHandler, - "_get_mcp_client_side_auth_header_name", + "get_mcp_client_side_auth_header_name", return_value="custom-auth-header", ): # Create ASGI scope with custom header diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py index b56a12db5b1..4081681daef 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py @@ -1196,3 +1196,45 @@ class TestOpenApiResolvedUpstreamAuth: ) assert resolved is None lookup.assert_not_awaited() + + +class TestPreCallToolCheckExposesClientHeaders: + """The pre_mcp_call guardrail payload must carry the caller's sanitized HTTP headers.""" + + @pytest.mark.asyncio + async def test_sanitized_client_headers_reach_the_guardrail_payload(self): + manager = MCPServerManager() + server = MCPServer( + server_id="test-id", + name="test_server", + server_name="test_server", + url="https://example.com", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + ) + + captured: Dict[str, Any] = {} + + def capture(request_obj, kwargs): + captured.update(kwargs) + return {"model": "fake"} + + proxy_logging = MagicMock(spec=ProxyLogging) + proxy_logging._create_mcp_request_object_from_kwargs = MagicMock(return_value=MagicMock()) + proxy_logging._convert_mcp_to_llm_format = MagicMock(side_effect=capture) + proxy_logging.pre_call_hook = AsyncMock(return_value=None) + + with patch.object(manager, "check_allowed_or_banned_tools", return_value=True): + with patch.object(manager, "check_tool_permission_for_key_team", new_callable=AsyncMock): + with patch.object(manager, "validate_allowed_params"): + await manager.pre_call_tool_check( + name="test_tool", + arguments={"key": "val"}, + server_name="test_server", + user_api_key_auth=None, + proxy_logging_obj=proxy_logging, + server=server, + raw_headers={"x-nuid": "nuid-1", "x-litellm-api-key": "sk-proxy"}, + ) + + assert captured["headers"] == {"x-nuid": "nuid-1"} diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 850d01c6e34..7df83065865 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -77,7 +77,7 @@ async def test_mcp_server_tool_call_body_contains_request_data(): # Mock the add_litellm_data_to_request function to capture the data captured_data = {} - async def mock_add_litellm_data_to_request(data, request, user_api_key_dict, proxy_config): + async def mock_add_litellm_data_to_request(data, request, user_api_key_dict, proxy_config, **kwargs): captured_data.update(data) # Simulate the proxy_server_request creation captured_data["proxy_server_request"] = { @@ -116,6 +116,107 @@ async def test_mcp_server_tool_call_body_contains_request_data(): assert body["arguments"] == tool_arguments +@pytest.mark.asyncio +async def test_mcp_server_tool_call_forwards_client_headers_to_logging(): + """The MCP protocol path must hand the connection's client headers to the pre-call + pipeline, so logging callbacks and guardrails see them the way the REST path does.""" + try: + from litellm.proxy._experimental.mcp_server.server import ( + mcp_server_tool_call, + set_auth_context, + ) + except ImportError: + pytest.skip("MCP server not available") + + set_auth_context( + UserAPIKeyAuth(api_key="test_key", user_id="test_user"), + raw_headers={ + "x-nuid": "nuid-1", + "x-app-id": "app-1", + "content-length": "42", + "x-forwarded-for": "9.9.9.9", + }, + client_ip="1.2.3.4", + ) + + captured_headers = {} + + async def mock_add_litellm_data_to_request(data, request, user_api_key_dict, proxy_config, **kwargs): + captured_headers.update(request.headers) + return data + + async def mock_call_mcp_tool(*args, **kwargs): + return [{"type": "text", "text": "mocked response"}] + + with patch( + "litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request", + mock_add_litellm_data_to_request, + ): + with patch( + "litellm.proxy._experimental.mcp_server.server.call_mcp_tool", + mock_call_mcp_tool, + ): + with patch("litellm.proxy.proxy_server.proxy_config", MagicMock()): + await mcp_server_tool_call("test_tool", {"param": "value"}) + + assert captured_headers.get("x-nuid") == "nuid-1" + assert captured_headers.get("x-app-id") == "app-1" + assert "content-length" not in captured_headers + assert captured_headers.get("x-forwarded-for") == "1.2.3.4" + + +@pytest.mark.asyncio +async def test_mcp_server_tool_call_strips_custom_litellm_key_header(): + """The deployment can rename the proxy key header via general_settings.litellm_key_header_name. + The pre-call pipeline only knows that name if it is passed in, so without it the virtual key + reaches metadata.headers and proxy_server_request.headers in plaintext.""" + try: + from litellm.proxy._experimental.mcp_server.server import ( + mcp_server_tool_call, + set_auth_context, + ) + from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request + except ImportError: + pytest.skip("MCP server not available") + + set_auth_context( + UserAPIKeyAuth(api_key="test_key", user_id="test_user"), + raw_headers={"x-company-key": "sk-proxy-secret", "x-nuid": "nuid-1"}, + client_ip="1.2.3.4", + ) + + captured_data = {} + + async def capturing_add_litellm_data_to_request(**kwargs): + data = await add_litellm_data_to_request(**kwargs) + captured_data.update(data) + return data + + async def mock_call_mcp_tool(*args, **kwargs): + return [{"type": "text", "text": "mocked response"}] + + with patch( + "litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request", + capturing_add_litellm_data_to_request, + ): + with patch( + "litellm.proxy._experimental.mcp_server.server.call_mcp_tool", + mock_call_mcp_tool, + ): + with patch("litellm.proxy.proxy_server.proxy_config", MagicMock()): + with patch.dict( + "litellm.proxy.proxy_server.general_settings", + {"litellm_key_header_name": "x-company-key"}, + clear=False, + ): + await mcp_server_tool_call("test_tool", {"param": "value"}) + + metadata_headers = captured_data["metadata"]["headers"] + assert metadata_headers.get("x-nuid") == "nuid-1" + assert "x-company-key" not in metadata_headers + assert "x-company-key" not in captured_data["proxy_server_request"]["headers"] + + @pytest.mark.asyncio async def test_mcp_server_tool_call_relays_upstream_auth_error_as_iserror(): """The MCP session manager serializes handler exceptions as JSON-RPC errors, so a mid-session @@ -133,7 +234,7 @@ async def test_mcp_server_tool_call_relays_upstream_auth_error_as_iserror(): set_auth_context(UserAPIKeyAuth(api_key="test_key", user_id="test_user")) - async def mock_add_litellm_data_to_request(data, request, user_api_key_dict, proxy_config): + async def mock_add_litellm_data_to_request(data, request, user_api_key_dict, proxy_config, **kwargs): return data async def mock_call_mcp_tool(*args, **kwargs): @@ -1245,7 +1346,7 @@ async def test_mcp_server_tool_call_body_with_none_arguments(): # Mock the add_litellm_data_to_request function to capture the data captured_data = {} - async def mock_add_litellm_data_to_request(data, request, user_api_key_dict, proxy_config): + async def mock_add_litellm_data_to_request(data, request, user_api_key_dict, proxy_config, **kwargs): captured_data.update(data) captured_data["proxy_server_request"] = { "url": str(request.url), diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_utils.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_utils.py index 73fdee9cde3..00ed4e91efa 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_utils.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_utils.py @@ -1,7 +1,11 @@ +from unittest.mock import patch + import pytest from fastapi import HTTPException from litellm.proxy._experimental.mcp_server.utils import ( + build_synthetic_mcp_request, + logging_safe_mcp_headers, validate_and_normalize_mcp_server_payload, validate_tool_display_names, ) @@ -47,3 +51,99 @@ class TestValidateAndNormalizeMcpServerPayload: tool_name_to_display_name={"read_wiki_structure": "browse_repo_docs"}, ) validate_and_normalize_mcp_server_payload(payload) + + +class TestLoggingSafeMcpHeaders: + def test_returns_empty_for_missing_headers(self): + assert logging_safe_mcp_headers(None) == {} + assert logging_safe_mcp_headers({}) == {} + + def test_exposes_custom_headers_and_masks_credentials(self): + safe = logging_safe_mcp_headers( + { + "x-nuid": "nuid-1", + "x-app-id": "app-1", + "x-litellm-api-key": "sk-proxy", + "cookie": "session=secret", + } + ) + assert safe == { + "x-nuid": "nuid-1", + "x-app-id": "app-1", + "cookie": "***REDACTED***", + } + + def test_strips_custom_litellm_key_header(self): + """general_settings.litellm_key_header_name carries the proxy virtual key, so it must + never reach a callback or a guardrail even though clean_headers cannot know its name.""" + with patch.dict( + "litellm.proxy.proxy_server.general_settings", + {"litellm_key_header_name": "x-company-key"}, + clear=False, + ): + safe = logging_safe_mcp_headers({"x-company-key": "sk-proxy", "x-nuid": "nuid-1"}) + + assert safe == {"x-nuid": "nuid-1"} + + def test_strips_client_controlled_redaction_opt_out(self): + """litellm-disable-message-redaction is read back out of the logged metadata to turn off + redaction, so leaving it in place lets any MCP client undo what an admin configured.""" + safe = logging_safe_mcp_headers({"litellm-disable-message-redaction": "true", "x-nuid": "nuid-1"}) + + assert safe == {"x-nuid": "nuid-1"} + + def test_strips_upstream_mcp_credentials(self): + safe = logging_safe_mcp_headers( + { + "x-mcp-auth": "Bearer upstream", + "x-mcp-github-authorization": "Bearer gh_token", + "x-mcp-zapier-x-api-key": "zapier-key", + "x-nuid": "nuid-1", + } + ) + + assert safe == {"x-nuid": "nuid-1"} + + def test_strips_custom_mcp_client_side_auth_header(self): + with patch.dict( + "litellm.proxy.proxy_server.general_settings", + {"mcp_client_side_auth_header_name": "x-upstream-token"}, + clear=False, + ): + safe = logging_safe_mcp_headers({"x-upstream-token": "Bearer upstream", "x-nuid": "nuid-1"}) + + assert safe == {"x-nuid": "nuid-1"} + + +class TestBuildSyntheticMcpRequest: + def test_forwards_client_headers_without_upstream_credentials(self): + """The synthetic request feeds add_litellm_data_to_request, which derives + metadata.headers, so upstream MCP credentials must not ride along.""" + request = build_synthetic_mcp_request( + path="/mcp/tools/call", + raw_headers={ + "x-nuid": "nuid-1", + "x-mcp-auth": "Bearer upstream", + "x-mcp-github-authorization": "Bearer gh_token", + }, + ) + + assert request.headers.get("x-nuid") == "nuid-1" + assert "x-mcp-auth" not in request.headers + assert "x-mcp-github-authorization" not in request.headers + + def test_drops_custom_litellm_key_header(self): + """Callers such as the sampling flow build metadata off this request, so the + deployment's custom proxy key header must never be forwarded on it.""" + with patch.dict( + "litellm.proxy.proxy_server.general_settings", + {"litellm_key_header_name": "x-company-key"}, + clear=False, + ): + request = build_synthetic_mcp_request( + path="/mcp/sampling/createMessage", + raw_headers={"x-company-key": "sk-proxy-secret", "x-nuid": "nuid-1"}, + ) + + assert request.headers.get("x-nuid") == "nuid-1" + assert "x-company-key" not in request.headers diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_mcp_bridging.py b/tests/test_litellm/proxy/utils/proxy_logging/test_mcp_bridging.py index 9defb309863..56057dce7e0 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_mcp_bridging.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_mcp_bridging.py @@ -76,6 +76,23 @@ def test_convert_mcp_to_llm_format_defaults_model(proxy_logging, make_mcp_reques } +def test_convert_mcp_to_llm_format_exposes_headers_on_metadata(proxy_logging, make_mcp_request_obj): + """Guardrails read the caller's HTTP headers off ``metadata.headers`` on the chat + completions path, so the MCP bridge has to put them in the same place.""" + req = make_mcp_request_obj() + out = proxy_logging._convert_mcp_to_llm_format( + request_obj=req, + kwargs={"headers": {"x-nuid": "nuid-1"}}, + ) + assert out["metadata"]["headers"] == {"x-nuid": "nuid-1"} + + +def test_convert_mcp_to_llm_format_defaults_headers_to_empty(proxy_logging, make_mcp_request_obj): + req = make_mcp_request_obj() + out = proxy_logging._convert_mcp_to_llm_format(request_obj=req, kwargs={}) + assert out["metadata"]["headers"] == {} + + def test_convert_mcp_to_llm_format_missing_request_obj_raises(proxy_logging): with pytest.raises(AttributeError): proxy_logging._convert_mcp_to_llm_format(request_obj=None, kwargs={}) diff --git a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py index d60fff66c44..4981caa10c3 100644 --- a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py +++ b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py @@ -536,6 +536,37 @@ async def test_get_mcp_tools_from_manager_forwards_request_tags(monkeypatch): assert mock_get_tools.await_args.kwargs["request_tags"] == ["team-a"] +@pytest.mark.asyncio +async def test_execute_tool_calls_exposes_sanitized_client_headers_to_logging(monkeypatch): + """The Responses API MCP bridge used to log an empty header dict, hiding the caller's + headers from logging callbacks and hooks.""" + _setup_proxy_logging(monkeypatch) + _setup_mcp_call_environment(monkeypatch) + + captured = {} + + def fake_function_setup(*_args, **kwargs): + captured.update(kwargs) + return None, None + + handler_module = importlib.import_module( + "litellm.responses.mcp.litellm_proxy_mcp_handler" + ) + monkeypatch.setattr(handler_module, "function_setup", fake_function_setup) + + tool_name = "deepwiki-read_wiki_structure" + await LiteLLM_Proxy_MCP_Handler._execute_tool_calls( + tool_server_map={tool_name: "deepwiki"}, + tool_calls=[{"id": "call-1", "function": {"name": tool_name, "arguments": "{}"}}], + user_api_key_auth=None, + raw_headers={"x-nuid": "nuid-1", "x-litellm-api-key": "sk-proxy", "cookie": "s=1"}, + ) + + expected = {"x-nuid": "nuid-1", "cookie": "***REDACTED***"} + assert captured["metadata"]["headers"] == expected + assert captured["proxy_server_request"]["headers"] == expected + + @pytest.mark.asyncio async def test_execute_tool_calls_propagates_request_tags_to_function_setup(monkeypatch): _setup_proxy_logging(monkeypatch) diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 0104733bcf3..894d99c92e0 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22943 + "limit": 22941 }, "LIT002": { - "limit": 27141 + "limit": 27139 }, "LIT003": { "limit": 269 @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16722 + "limit": 16716 }, "LIT011": { "limit": 5596