mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-23 00:41:40 +00:00
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.
This commit is contained in:
parent
33c3f13443
commit
ebf751f56e
3 changed files with 144 additions and 0 deletions
|
|
@ -1818,6 +1818,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
|
||||
)
|
||||
|
|
@ -1826,6 +1832,14 @@ if MCP_AVAILABLE:
|
|||
"mcp_tool_call_metadata"
|
||||
] = 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
|
||||
if custom_headers:
|
||||
_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
|
||||
# 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:
|
||||
|
|
@ -2113,6 +2127,41 @@ if MCP_AVAILABLE:
|
|||
raw_headers=raw_headers,
|
||||
)
|
||||
|
||||
_SENSITIVE_CUSTOM_HEADER_PREFIXES = frozenset(
|
||||
{
|
||||
"x-mcp-server-auth-",
|
||||
"x-litellm-",
|
||||
}
|
||||
)
|
||||
|
||||
_SENSITIVE_CUSTOM_HEADERS = frozenset(
|
||||
{
|
||||
"x-api-key",
|
||||
}
|
||||
)
|
||||
|
||||
def _extract_custom_headers(
|
||||
raw_headers: Optional[Dict[str, str]],
|
||||
) -> Optional[Dict[str, str]]:
|
||||
"""Extract x-* custom headers from raw HTTP headers, excluding sensitive ones."""
|
||||
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],
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -0,0 +1,89 @@
|
|||
"""
|
||||
Tests for custom MCP headers flowing into logging callbacks.
|
||||
|
||||
Validates that:
|
||||
1. _extract_custom_headers correctly filters raw HTTP headers
|
||||
2. Custom headers are stored in StandardLoggingMCPToolCall.custom_headers
|
||||
3. Custom headers populate requester_custom_headers in StandardLoggingMetadata
|
||||
"""
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.server import _extract_custom_headers
|
||||
|
||||
|
||||
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_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"}
|
||||
Loading…
Add table
Reference in a new issue