From dd404e3bdcf9b0d50e639f1578b39d697b5cfd9c Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Thu, 9 Apr 2026 04:47:31 +0530 Subject: [PATCH] feat(mcp): expose custom x-* headers in logging callbacks (#25338) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(mcp): pass custom x-* headers to logging callbacks Custom headers (x-* prefixed) sent in MCP tool call requests are now extracted and stored in StandardLoggingMCPToolCall.custom_headers, and also propagated to requester_custom_headers in StandardLoggingMetadata so logging callbacks can access them via kwargs. Sensitive headers (x-api-key, x-litellm-*, x-mcp-server-auth-*) are filtered out before logging. * Fix mcp tool call logging * fix(proxy): wire MCP REST /tools/call to litellm logging callbacks Use logging_obj from common_processing_pre_call_logic (not the JSON body) and mirror call_mcp_tool post_call, async_post_mcp_tool_call_hook, and async_success_handler so custom loggers fire for mcp-rest tool execution. Extract _execute_mcp_tool_rest_with_post_hooks to satisfy PLR0915. Made-with: Cursor * fix(mcp): Greptile review — custom header logging diagnostics and tests - Add _apply_requester_custom_headers_for_mcp_logging with debug logs when litellm_params or metadata is not a dict. - Expand _extract_custom_headers deny-list for common auth-style x-* names; document non-exhaustive filtering in the docstring. - Extend tests for deny-list and metadata wiring paths. Made-with: Cursor * fix test * Fix greptile concern * Fix greptile concern --- litellm/litellm_core_utils/litellm_logging.py | 26 +++- .../mcp_server/rest_endpoints.py | 9 +- .../proxy/_experimental/mcp_server/server.py | 118 ++++++++++++-- litellm/types/utils.py | 6 + litellm/utils.py | 5 +- .../test_mcp_custom_headers_logging.py | 144 ++++++++++++++++++ 6 files changed, 285 insertions(+), 23 deletions(-) create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_custom_headers_logging.py diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 7395b65626f..308d3f20d0c 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -287,6 +287,28 @@ def _get_cached_prometheus_logger(): return _PrometheusLogger +def _normalize_mcp_tool_response_for_post_hook(response_obj: Any) -> List[Any]: + """ + Coerce MCP tool execution results into content blocks for MCPPostCallResponseObject. + + ``CallToolResult`` exposes ``content`` as a list of MCP content blocks; callers or + tests may pass a plain dict or other JSON-serializable value instead. + """ + if response_obj is None: + return [] + if hasattr(response_obj, "content"): + content = getattr(response_obj, "content") + if isinstance(content, list): + return list(content) + if isinstance(response_obj, list): + return list(response_obj) + try: + text = json.dumps(response_obj, default=str) + except Exception: + text = str(response_obj) + return [{"type": "text", "text": text}] + + class Logging(LiteLLMLoggingBaseClass): global supabaseClient, promptLayerLogger, weightsBiasesLogger, logfireLogger, capture_exception, add_breadcrumb, lunaryLogger, logfireLogger, prometheusLogger, slack_app custom_pricing: bool = False @@ -1312,9 +1334,11 @@ class Logging(LiteLLMLoggingBaseClass): dynamic_success_callbacks=self.dynamic_success_callbacks, global_callbacks=litellm.success_callback, ) + normalized_mcp_content = _normalize_mcp_tool_response_for_post_hook(response_obj) post_mcp_tool_call_response_obj: MCPPostCallResponseObject = ( MCPPostCallResponseObject( - mcp_tool_call_response=response_obj, hidden_params=HiddenParams() + mcp_tool_call_response=normalized_mcp_content, + hidden_params=HiddenParams(), ) ) for callback in callbacks: diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 32560a2211d..241f765d3e8 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -798,20 +798,19 @@ if MCP_AVAILABLE: target_server, user_api_key_dict ) - # Call execute_mcp_tool directly (permission checks already done) - result = await execute_mcp_tool( + mcp_tool_start_time = datetime.now() + return await execute_mcp_tool( name=tool_name, arguments=tool_arguments, allowed_mcp_servers=allowed_mcp_servers, - start_time=datetime.now(), + start_time=mcp_tool_start_time, user_api_key_auth=data.get("user_api_key_auth"), mcp_auth_header=data.get("mcp_auth_header"), mcp_server_auth_headers=data.get("mcp_server_auth_headers"), oauth2_headers=user_oauth_extra_headers or data.get("oauth2_headers"), raw_headers=data.get("raw_headers"), - litellm_logging_obj=data.get("litellm_logging_obj"), + litellm_logging_obj=logging_obj, ) - return result except BlockedPiiEntityError as e: verbose_logger.error(f"BlockedPiiEntityError in MCP tool call: {str(e)}") raise HTTPException( diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 7fc28b68e9c..6b4b98b893f 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -1842,6 +1842,10 @@ if MCP_AVAILABLE: Returns: CallToolResult: Tool execution result + + With ``litellm_logging_obj``, runs MCP post-call and success logging once + before returning (the ``@client`` wrapper skips duplicate async success + logging for ``call_mcp_tool``). """ # Track resolved MCP server for both permission checks and dispatch mcp_server: Optional[MCPServer] = None @@ -1877,6 +1881,12 @@ if MCP_AVAILABLE: server_name=server_name, ) ) + + # Extract custom headers for logging callbacks + custom_headers = _extract_custom_headers(raw_headers) + if custom_headers: + standard_logging_mcp_tool_call["custom_headers"] = custom_headers + litellm_logging_obj: Optional[LiteLLMLoggingObj] = kwargs.get( "litellm_logging_obj", None ) @@ -1885,6 +1895,11 @@ if MCP_AVAILABLE: standard_logging_mcp_tool_call ) litellm_logging_obj.model = f"MCP: {name}" + # Populate requester_custom_headers in metadata so it flows into + # StandardLoggingMetadata via _STANDARD_LOGGING_METADATA_KEYS + _apply_requester_custom_headers_for_mcp_logging( + litellm_logging_obj, custom_headers + ) # Resolve the MCP server early so BYOK checks and credential injection # apply to ALL dispatch paths (local tool registry AND managed MCP server). if mcp_server is None: @@ -1977,6 +1992,22 @@ if MCP_AVAILABLE: local_content = await _handle_local_mcp_tool(original_tool_name, arguments) response = CallToolResult(content=cast(Any, local_content), isError=False) + if litellm_logging_obj: + litellm_logging_obj.post_call(original_response=response) + end_time = datetime.now() + await litellm_logging_obj.async_post_mcp_tool_call_hook( + kwargs=litellm_logging_obj.model_call_details, + response_obj=response, + start_time=start_time, + end_time=end_time, + ) + litellm_logging_obj.call_type = CallTypes.call_mcp_tool.value + await litellm_logging_obj.async_success_handler( + result=response, + start_time=start_time, + end_time=end_time, + ) + return response @client @@ -1995,9 +2026,6 @@ if MCP_AVAILABLE: Call a specific tool with the provided arguments (handles prefixed tool names). """ start_time = datetime.now() - litellm_logging_obj: Optional[LiteLLMLoggingObj] = kwargs.get( - "litellm_logging_obj", None - ) try: if arguments is None: @@ -2057,19 +2085,6 @@ if MCP_AVAILABLE: ) raise - if litellm_logging_obj: - litellm_logging_obj.post_call(original_response=response) - end_time = datetime.now() - await litellm_logging_obj.async_post_mcp_tool_call_hook( - kwargs=litellm_logging_obj.model_call_details, - response_obj=response, - start_time=start_time, - end_time=end_time, - ) - litellm_logging_obj.call_type = CallTypes.call_mcp_tool.value - await litellm_logging_obj.async_success_handler( - result=response, start_time=start_time, end_time=end_time - ) return response async def mcp_get_prompt( @@ -2172,6 +2187,77 @@ if MCP_AVAILABLE: raw_headers=raw_headers, ) + def _apply_requester_custom_headers_for_mcp_logging( + litellm_logging_obj: LiteLLMLoggingObj, + custom_headers: Optional[Dict[str, str]], + ) -> None: + """Copy MCP custom headers into litellm_params.metadata for standard logging.""" + if not custom_headers: + return + _lp = litellm_logging_obj.model_call_details.get("litellm_params") + if isinstance(_lp, dict): + _meta = _lp.get("metadata") + if isinstance(_meta, dict): + _meta["requester_custom_headers"] = custom_headers + else: + verbose_logger.debug( + "execute_mcp_tool: skipping requester_custom_headers " + "— metadata is not a dict (type=%s)", + type(_meta).__name__, + ) + else: + verbose_logger.debug( + "execute_mcp_tool: skipping requester_custom_headers " + "— litellm_params is not a dict (type=%s)", + type(_lp).__name__, + ) + + _SENSITIVE_CUSTOM_HEADER_PREFIXES = frozenset( + { + "x-mcp-server-auth-", + "x-litellm-", + } + ) + + _SENSITIVE_CUSTOM_HEADERS = frozenset( + { + "x-api-key", + "x-auth-token", + "x-access-token", + "x-goog-api-key", + "x-forwarded-authorization", + } + ) + + def _extract_custom_headers( + raw_headers: Optional[Dict[str, str]], + ) -> Optional[Dict[str, str]]: + """Extract ``x-*`` headers for MCP logging (non-secrets only). + + Only names starting with ``x-`` are included. Exclusions are a + **non-exhaustive** deny-list: fixed sensitive names, plus ``x-litellm-`` + and ``x-mcp-server-auth-`` prefixes. Other ``x-*`` headers may still + carry secrets; operators should audit what clients send and extend + ``_SENSITIVE_CUSTOM_HEADERS`` / prefixes if needed. + """ + if not raw_headers: + return None + custom: Dict[str, str] = {} + for k, v in raw_headers.items(): + key_lower = k.lower() + if not key_lower.startswith("x-"): + continue + if key_lower in _SENSITIVE_CUSTOM_HEADERS: + continue + if any( + key_lower.startswith(prefix) + for prefix in _SENSITIVE_CUSTOM_HEADER_PREFIXES + ): + continue + if v is not None and isinstance(v, str): + custom[k] = v + return custom if custom else None + def _get_standard_logging_mcp_tool_call( name: str, arguments: Dict[str, Any], diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 3f6e6e5aa5a..60f040c489f 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2556,6 +2556,12 @@ class StandardLoggingMCPToolCall(TypedDict, total=False): Cost per query for the MCP server tool call """ + custom_headers: Optional[Dict[str, str]] + """ + Custom (x-*) headers sent by the client in the MCP tool call request. + Filtered to include only x-* prefixed headers, excluding sensitive auth headers. + """ + class StandardLoggingVectorStoreRequest(TypedDict, total=False): """ diff --git a/litellm/utils.py b/litellm/utils.py index f902644e760..73a086080b8 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1945,7 +1945,10 @@ def client(original_function): # noqa: PLR0915 # parent's logging obj; skip async logging here so only the outer call bills once. # NOTE: streaming requests return early (before this point) via # CustomStreamWrapper, so this block is non-streaming only. - if not _is_litellm_internal_call: + # call_mcp_tool: execute_mcp_tool already runs post_call, + # async_post_mcp_tool_call_hook, and async_success_handler; skip the + # wrapper worker to avoid duplicate spend logs. + if not _is_litellm_internal_call and call_type != CallTypes.call_mcp_tool.value: if getattr(logging_obj, "_defer_async_logging", False): def _enqueue_deferred_logging() -> None: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_custom_headers_logging.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_custom_headers_logging.py new file mode 100644 index 00000000000..f58883e4199 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_custom_headers_logging.py @@ -0,0 +1,144 @@ +""" +Tests for custom MCP headers flowing into logging callbacks. + +Covers: +- ``_extract_custom_headers`` deny-list and inclusion rules for ``x-*`` headers. +- ``_apply_requester_custom_headers_for_mcp_logging`` wiring into + ``litellm_params.metadata`` (the path used by ``execute_mcp_tool``). +""" + +import logging + +from litellm.proxy._experimental.mcp_server.server import ( + _apply_requester_custom_headers_for_mcp_logging, + _extract_custom_headers, +) + + +class _FakeLiteLLMLogging: + __slots__ = ("model_call_details",) + + def __init__(self, model_call_details: dict): + self.model_call_details = model_call_details + + +class TestExtractCustomHeaders: + """Unit tests for _extract_custom_headers helper.""" + + def test_returns_none_for_none_input(self): + assert _extract_custom_headers(None) is None + + def test_returns_none_for_empty_dict(self): + assert _extract_custom_headers({}) is None + + def test_extracts_x_prefixed_headers(self): + raw = { + "x-custom-header-foo": "bar", + "x-request-id": "abc123", + "content-type": "application/json", + "authorization": "Bearer token", + } + result = _extract_custom_headers(raw) + assert result == { + "x-custom-header-foo": "bar", + "x-request-id": "abc123", + } + + def test_excludes_x_api_key(self): + raw = { + "x-api-key": "secret", + "x-custom-foo": "bar", + } + result = _extract_custom_headers(raw) + assert result == {"x-custom-foo": "bar"} + + def test_excludes_common_auth_style_x_headers(self): + raw = { + "x-auth-token": "secret", + "x-access-token": "secret", + "x-goog-api-key": "secret", + "x-forwarded-authorization": "secret", + "x-safe-correlation": "ok", + } + result = _extract_custom_headers(raw) + assert result == {"x-safe-correlation": "ok"} + + def test_excludes_x_litellm_prefixed(self): + raw = { + "x-litellm-api-key": "secret", + "x-litellm-mcp-debug": "true", + "x-custom-foo": "bar", + } + result = _extract_custom_headers(raw) + assert result == {"x-custom-foo": "bar"} + + def test_excludes_x_mcp_server_auth_prefixed(self): + raw = { + "x-mcp-server-auth-token": "secret", + "x-custom-foo": "bar", + } + result = _extract_custom_headers(raw) + assert result == {"x-custom-foo": "bar"} + + def test_preserves_original_key_casing(self): + raw = {"X-Custom-Header": "value"} + result = _extract_custom_headers(raw) + assert result == {"X-Custom-Header": "value"} + + def test_returns_none_when_all_headers_filtered(self): + raw = { + "content-type": "application/json", + "authorization": "Bearer token", + "x-api-key": "secret", + } + assert _extract_custom_headers(raw) is None + + def test_excludes_non_string_values(self): + raw = { + "x-good": "value", + "x-bad": None, # type: ignore + } + result = _extract_custom_headers(raw) + assert result == {"x-good": "value"} + + def test_case_insensitive_prefix_matching(self): + """Header keys with mixed case should still be filtered correctly.""" + raw = { + "X-API-Key": "secret", + "X-Litellm-Something": "hidden", + "X-Custom-Foo": "visible", + } + result = _extract_custom_headers(raw) + assert result == {"X-Custom-Foo": "visible"} + + +class TestApplyRequesterCustomHeadersForMcpLogging: + """Unit tests for metadata wiring used by execute_mcp_tool.""" + + def test_sets_requester_custom_headers_when_metadata_is_dict(self): + meta: dict = {} + fake = _FakeLiteLLMLogging({"litellm_params": {"metadata": meta}}) + headers = {"x-correlation-id": "abc"} + _apply_requester_custom_headers_for_mcp_logging(fake, headers) + assert meta["requester_custom_headers"] == headers + + def test_noop_when_custom_headers_empty(self): + meta = {"before": True} + fake = _FakeLiteLLMLogging({"litellm_params": {"metadata": meta}}) + _apply_requester_custom_headers_for_mcp_logging(fake, None) + _apply_requester_custom_headers_for_mcp_logging(fake, {}) + assert "requester_custom_headers" not in meta + + def test_skips_when_litellm_params_not_dict(self, caplog): + fake = _FakeLiteLLMLogging({"litellm_params": None}) + with caplog.at_level(logging.DEBUG, logger="LiteLLM"): + _apply_requester_custom_headers_for_mcp_logging(fake, {"x": "y"}) + assert "skipping requester_custom_headers" in caplog.text + assert "litellm_params is not a dict" in caplog.text + + def test_skips_when_metadata_not_dict(self, caplog): + fake = _FakeLiteLLMLogging({"litellm_params": {"metadata": object()}}) + with caplog.at_level(logging.DEBUG, logger="LiteLLM"): + _apply_requester_custom_headers_for_mcp_logging(fake, {"x": "y"}) + assert "skipping requester_custom_headers" in caplog.text + assert "metadata is not a dict" in caplog.text