scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-24 00:52:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 91

fix(mcp): preserve discovery attribution and sanitize logging headers (#42541)

Browse source 
Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com>
This commit is contained in:
joshua-berri 2026-09-22 21:26:48 +00:00 • committed by GitHub
parent 6764868861
commit b277be0867
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 109 additions and 19 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
litellm/proxy/_experimental/mcp_server/utils.py
View file

@ -986,7 +986,7 @@ def _forwarded_upstream_header_names() -> frozenset[str]:
)
def _upstream_credential_headers(header_names: Iterable[str]) -> frozenset[str]:
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, any
header name a configured server forwards upstream via ``extra_headers``, and the
@ -1038,7 +1038,7 @@ def build_synthetic_mcp_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 ())
| upstream_credential_headers(raw_headers.keys() if raw_headers else ())
| (frozenset({custom_key_header.lower()}) if custom_key_header else frozenset())
)
forwarded: Final = tuple(
@ -1086,7 +1086,7 @@ def logging_safe_mcp_headers(raw_headers: Mapping[str, str] | None) -> Mapping[s
)
excluded: Final = (
_upstream_credential_headers(raw_headers.keys() if raw_headers else ())
upstream_credential_headers(raw_headers.keys() if raw_headers else ())
| UNTRUSTED_REQUEST_HEADER_CONTROL_FIELDS
| frozenset({"host"})
)

9
litellm/proxy/litellm_pre_call_utils.py
View file

@ -2029,7 +2029,14 @@ async def add_litellm_data_to_request(
_headers,
allow_client_message_redaction_opt_out=_allow_client_message_redaction_opt_out,
)
_logging_safe_headers: Final = redact_credential_headers(_headers)
from litellm.proxy._experimental.mcp_server.utils import upstream_credential_headers
_mcp_credential_headers: Final = upstream_credential_headers(_headers)
_logging_safe_headers: Final = redact_credential_headers(
MappingProxyType(
{name: value for name, value in _headers.items() if name.lower() not in _mcp_credential_headers}
)
)
verbose_proxy_logger.debug("Request Headers: %s", _logging_safe_headers)
verbose_proxy_logger.debug("Raw Headers: %s", _raw_headers)

19
litellm/responses/main.py
View file

@ -209,16 +209,12 @@ async def aresponses_api_with_mcp(
user_api_key_auth = kwargs.get("user_api_key_auth") or kwargs.get("litellm_metadata", {}).get("user_api_key_auth")
# Extract MCP auth headers from request (for dynamic auth when fetching tools)
mcp_auth_header: str | None = None
mcp_server_auth_headers: dict[str, dict[str, str]] | None = None
secret_fields = kwargs.get("secret_fields")
if secret_fields and isinstance(secret_fields, dict):
(
mcp_auth_header,
mcp_server_auth_headers,
_,
_,
) = ResponsesAPIRequestUtils.extract_mcp_headers_from_request(secret_fields=secret_fields, tools=tools)
secret_fields: Final = kwargs.get("secret_fields")
mcp_auth_header, mcp_server_auth_headers, _, discovery_raw_headers = (
ResponsesAPIRequestUtils.extract_mcp_headers_from_request(secret_fields=secret_fields, tools=tools)
if isinstance(secret_fields, dict) and secret_fields
else (None, None, None, None)
)
# Get original MCP tools (for events) and OpenAI tools (for LLM) by reusing existing methods
(
@ -231,6 +227,7 @@ async def aresponses_api_with_mcp(
mcp_auth_header=mcp_auth_header,
mcp_server_auth_headers=mcp_server_auth_headers,
request_tags=LiteLLM_Proxy_MCP_Handler._get_parent_request_tags(kwargs),
raw_headers=discovery_raw_headers,
)
openai_tools: Final = LiteLLM_Proxy_MCP_Handler._transform_mcp_tools_to_openai(original_mcp_tools)
@ -330,7 +327,6 @@ async def aresponses_api_with_mcp(
user_api_key_auth = kwargs.get("litellm_metadata", {}).get("user_api_key_auth")
# Extract MCP auth headers from the request to pass to MCP server
secret_fields = kwargs.get("secret_fields")
(
mcp_auth_header,
mcp_server_auth_headers,
@ -416,6 +412,7 @@ async def aresponses_api_with_mcp(
mcp_auth_header=mcp_auth_header,
mcp_server_auth_headers=mcp_server_auth_headers,
request_tags=LiteLLM_Proxy_MCP_Handler._get_parent_request_tags(kwargs),
raw_headers=discovery_raw_headers,
)
final_response = LiteLLM_Proxy_MCP_Handler._add_mcp_output_elements_to_response(
response=final_response,

1
litellm/responses/mcp/chat_completions_handler.py
View file

@ -138,6 +138,7 @@ async def acompletion_with_mcp(
mcp_auth_header=mcp_auth_header,
mcp_server_auth_headers=mcp_server_auth_headers,
request_tags=request_tags,
raw_headers=raw_headers,
)
openai_tools: Final = LiteLLM_Proxy_MCP_Handler._transform_mcp_tools_to_openai(

4
litellm/responses/mcp/litellm_proxy_mcp_handler.py
View file

@ -236,6 +236,7 @@ class LiteLLM_Proxy_MCP_Handler:
mcp_auth_header: str | None = None,
mcp_server_auth_headers: dict[str, dict[str, str]] | None = None,
request_tags: list[str] | None = None,
raw_headers: dict[str, str] | None = None,
) -> tuple[list[MCPTool], list[str]]:
"""
Get available tools from the MCP server manager.
@ -326,6 +327,7 @@ class LiteLLM_Proxy_MCP_Handler:
list_tools_log_source="responses",
litellm_trace_id=litellm_trace_id,
request_tags=request_tags,
raw_headers=raw_headers,
)
tools: Final = listing.tools
@ -452,6 +454,7 @@ class LiteLLM_Proxy_MCP_Handler:
mcp_auth_header: str | None = None,
mcp_server_auth_headers: dict[str, dict[str, str]] | None = None,
request_tags: list[str] | None = None,
raw_headers: dict[str, str] | None = None,
) -> tuple[list[MCPTool], dict[str, str]]:
"""
Process MCP tools through filtering and deduplication pipeline without OpenAI transformation.
@ -482,6 +485,7 @@ class LiteLLM_Proxy_MCP_Handler:
mcp_auth_header=mcp_auth_header,
mcp_server_auth_headers=mcp_server_auth_headers,
request_tags=request_tags,
raw_headers=raw_headers,
)
# Step 2: Filter tools based on allowed_tools parameter

6
tests/test_litellm/proxy/_experimental/mcp_server/test_utils.py
View file

@ -4,7 +4,7 @@ import pytest
from fastapi import HTTPException
from litellm.proxy._experimental.mcp_server.utils import (
_upstream_credential_headers,
upstream_credential_headers,
build_synthetic_mcp_request,
logging_safe_mcp_headers,
validate_and_normalize_mcp_server_payload,
@ -189,8 +189,8 @@ class TestLoggingSafeMcpHeaders:
"""clean_headers already strips authorization, and claiming it here would change which
header authenticated_with_header resolves to on a config that lists it by design."""
with _configured_servers(_server_forwarding("Authorization", "X-GitHub-Token")):
assert "authorization" not in _upstream_credential_headers(["authorization", "x-github-token"])
assert "x-github-token" in _upstream_credential_headers(["authorization", "x-github-token"])
assert "authorization" not in upstream_credential_headers(["authorization", "x-github-token"])
assert "x-github-token" in upstream_credential_headers(["authorization", "x-github-token"])
def test_keeps_headers_when_no_server_forwards_them(self):
with _configured_servers(_server_forwarding("x-github-token")):

42
tests/test_litellm/proxy/test_litellm_pre_call_utils.py
View file

@ -8238,3 +8238,45 @@ def test_default_team_settings_bool_turn_off_message_logging_redacts():
)
is True
)
@pytest.mark.asyncio
@pytest.mark.parametrize("path", ["/mcp-rest/tools/call", "/v1/responses", "/v1/chat/completions"])
@pytest.mark.parametrize("custom_auth", ["x-mcp-auth", "x-private-mcp-token"])
async def test_mcp_credentials_only_removed_from_logging_copies(path: str, custom_auth: str):
from litellm.types.mcp_server.mcp_server_manager import MCPServer
metadata_name: Final = "litellm_metadata" if path == "/v1/responses" else "metadata"
secrets: Final = {
"X-MCP-Deepwiki-Authorization": "upstream-sentinel",
custom_auth: "client-auth-sentinel",
"x-service-token": "configured-secret-sentinel",
}
attribution: Final = {"x-app-id": "app-a", "x-nuid": "user-a", "x-user-id": "identity-a"}
request: Final = _make_request_mock(path, {"Content-Type": "application/json", **secrets, **attribution})
request.headers = Headers(request.headers)
settings: Final = {"mcp_client_side_auth_header_name": custom_auth, "user_header_name": "x-user-id"}
server: Final = MCPServer(
server_id="header-test", name="header-test", transport="http", url="https://example.com/mcp",
extra_headers=["x-service-token", "x-user-id"],
)
with (
patch("litellm.proxy.proxy_server.general_settings", settings),
patch.dict(
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager.config_mcp_servers",
{"header-test": server}, clear=True,
),
):
updated: Final = await add_litellm_data_to_request(
data={"model": "test-model", "messages": [{"role": "user", "content": "hello"}]},
request=request, user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
proxy_config=MagicMock(), general_settings=settings, version="test",
)
for header_dict in _all_header_dicts(updated, metadata_name):
assert not any(value in json.dumps(header_dict) for value in secrets.values())
assert updated[metadata_name]["headers"] == updated["proxy_server_request"]["headers"]
for name, value in attribution.items():
assert updated[metadata_name]["headers"][name] == value
for name, value in secrets.items():
assert updated["secret_fields"]["raw_headers"][name.lower()] == value
assert request.headers[name] == value

1
tests/test_litellm/responses/mcp/test_chat_completions_handler.py
View file

@ -159,6 +159,7 @@ async def test_acompletion_with_mcp_passes_mcp_server_auth_headers_to_process_to
secret_fields=secret_fields,
)
assert captured_process_kwargs["raw_headers"] == secret_fields["raw_headers"]
assert "mcp_server_auth_headers" in captured_process_kwargs
mcp_server_auth_headers = captured_process_kwargs["mcp_server_auth_headers"]
assert mcp_server_auth_headers is not None

40
tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py
View file

@ -3,7 +3,7 @@ import subprocess
import sys
import textwrap
import types
from typing import Any, cast
from typing import Any, Final, cast
from unittest.mock import AsyncMock, MagicMock
import pytest
@ -1226,6 +1226,7 @@ async def test_mcp_follow_up_call_is_stateless_when_store_is_false(
)
async def fake_process(**kwargs: Any) -> tuple[list[Any], dict[str, str]]:
assert kwargs["raw_headers"] == {"x-app-id": "follow-up-caller"}
return ([], {"foo": "litellm_proxy"})
async def fake_execute(**kwargs: Any) -> list[dict[str, Any]]:
@ -1245,6 +1246,7 @@ async def test_mcp_follow_up_call_is_stateless_when_store_is_false(
model="gpt-5",
tools=[{"type": "mcp", "server_url": "litellm_proxy", "require_approval": "never"}],
litellm_metadata={"guardrails": ["block-all"]},
secret_fields={"raw_headers": {"x-app-id": "follow-up-caller"}},
store=store,
previous_response_id=caller_previous_response_id,
)
@ -1257,3 +1259,39 @@ async def test_mcp_follow_up_call_is_stateless_when_store_is_false(
item for item in follow_up_call["input"] if isinstance(item, dict) and item.get("type") == "reasoning"
]
assert bool(reasoning_items) is (store is False)
@pytest.mark.asyncio
async def test_responses_discovery_logs_sanitized_caller_headers(monkeypatch: pytest.MonkeyPatch):
from litellm.proxy._experimental.mcp_server import operations
from litellm.proxy._experimental.mcp_server import mcp_server_manager
headers: Final = {
"x-app-id": "app-a", "x-nuid": "user-a", "x-user-id": "identity-a",
"x-mcp-deepwiki-authorization": "upstream-sentinel", "authorization": "proxy-sentinel",
}
manager: Final = types.SimpleNamespace(
get_registry=MagicMock(return_value={}),
get_allowed_mcp_servers=AsyncMock(return_value=[]),
get_mcp_servers_from_ids=MagicMock(return_value=[]),
)
logger: Final = MagicMock(model_call_details={})
logger.async_success_handler = AsyncMock()
setup: Final = MagicMock(return_value=(logger, None))
monkeypatch.setattr(mcp_server_manager, "global_mcp_server_manager", manager)
monkeypatch.setattr(operations, "_get_allowed_mcp_servers", AsyncMock(return_value=[]))
monkeypatch.setattr(operations, "function_setup", setup)
response: Final = ResponsesAPIResponse(
id="resp_test", created_at=1234567891, model="test-model", object="response",
status="completed", output=[], parallel_tool_calls=False, tool_choice="auto", tools=[],
)
monkeypatch.setattr(responses_main, "aresponses", AsyncMock(return_value=response))
result: Final = await responses_main.aresponses_api_with_mcp(
input="hi", model="test-model", tools=[{"type": "mcp", "server_url": "litellm_proxy"}],
secret_fields={"raw_headers": headers},
)
assert result is response
logger.async_success_handler.assert_awaited_once()
logged: Final = setup.call_args.kwargs["metadata"]["headers"]
assert logged == {"x-app-id": "app-a", "x-nuid": "user-a", "x-user-id": "identity-a"}
assert headers["x-mcp-deepwiki-authorization"] == "upstream-sentinel"