scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

fix(mcp): isolate gateway attribution headers

Browse source
This commit is contained in:
MokshaVora 2026-08-10 16:30:11 +05:30
parent 06a0e2422b
commit c8ca0d6228
4 changed files with 77 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
README.md
View file

@ -255,17 +255,13 @@ curl -X POST 'http://0.0.0.0:4000/v1/chat/completions' \
"LiteLLM": {
"url": "http://localhost:4000/mcp/",
"headers": {
"x-litellm-api-key": "Bearer sk-1234",
"x-litellm-end-user-id": "user-123",
"x-litellm-tags": "application:orders,service:checkout"
"x-litellm-api-key": "Bearer sk-1234"
}
}
}
}
```
`x-litellm-end-user-id` associates the request with a canonical end-user identity. `x-litellm-tags` uses comma-separated tags for additional spend attribution dimensions. LiteLLM does not forward either header to the upstream MCP server.
[**Docs: MCP Gateway**](https://docs.litellm.ai/docs/mcp)
</details>

8
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
View file

@ -773,6 +773,10 @@ def _without_authorization(
return filtered or None
def _is_gateway_attribution_header(header_name: str) -> bool:
return header_name.lower() in {"x-litellm-tags", "x-litellm-end-user-id"}
def _format_byok_openapi_auth_header(mcp_server: MCPServer, mcp_auth_header: str) -> str:
"""Format a raw BYOK credential for OpenAPI tool ``Authorization`` injection."""
if mcp_server.auth_type == MCPAuth.api_key:
@ -799,6 +803,8 @@ def _openapi_forwarded_extra_headers(
for header_name in mcp_server.extra_headers:
if not isinstance(header_name, str):
continue
if _is_gateway_attribution_header(header_name):
continue
if skip_caller_authorization and header_name.lower() == "authorization":
continue
value = normalized_raw.get(header_name.lower())
@ -4841,6 +4847,8 @@ class MCPServerManager:
for header in mcp_server.extra_headers:
if not isinstance(header, str):
continue
if _is_gateway_attribution_header(header):
continue
if header.lower() == "authorization" and strip_caller_authorization:
continue
header_value = normalized_raw_headers.get(header.lower())

3
litellm/proxy/_experimental/mcp_server/server.py
View file

@ -428,6 +428,7 @@ if MCP_AVAILABLE:
MCPServerManager,
_caller_authorization_fans_out,
_client_forwarded_authorization_headers,
_is_gateway_attribution_header,
_should_strip_caller_authorization,
_without_authorization,
global_mcp_server_manager,
@ -1765,7 +1766,7 @@ if MCP_AVAILABLE:
for header in server.extra_headers:
if not isinstance(header, str):
continue
if header.lower() in {"x-litellm-tags", "x-litellm-end-user-id"}:
if _is_gateway_attribution_header(header):
continue
if header.lower() == "authorization" and (
strip_caller_authorization or withhold_forwarded_authorization

67
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py
View file

@ -37,6 +37,7 @@ from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
MCPServerManager,
_deserialize_json_dict,
_flow_endpoints_missing,
_openapi_forwarded_extra_headers,
_oauth_endpoints_unresolved,
_deserialize_json_list,
_normalize_mcp_server_cost_info,
@ -912,7 +913,7 @@ class TestMCPServerManager:
async def capture_create_mcp_client(
server, mcp_auth_header, extra_headers, stdio_env, **kwargs
): # pragma: no cover - helper
):
nonlocal captured_extra_headers
captured_extra_headers = extra_headers
return mock_client
@ -934,6 +935,70 @@ class TestMCPServerManager:
assert captured_extra_headers == {"Authorization": "Bearer token"}
assert isinstance(result, CallToolResult)
@pytest.mark.asyncio
async def test_call_regular_mcp_tool_does_not_forward_gateway_attribution_headers(self):
manager = MCPServerManager()
server = MCPServer(
server_id="server-attribution-call",
name="attribution-call-server",
url="https://example.com",
transport=MCPTransport.http,
auth_type=MCPAuth.authorization,
extra_headers=["x-request-id", "X-LiteLLM-Tags", "X-LiteLLM-End-User-Id"],
)
mock_client = AsyncMock()
mock_client.call_tool = AsyncMock(return_value=CallToolResult(content=[], isError=False))
captured_extra_headers = None
async def capture_create_mcp_client(
server, mcp_auth_header, extra_headers, stdio_env, **kwargs
): # pragma: no cover - helper
nonlocal captured_extra_headers
captured_extra_headers = extra_headers
return mock_client
manager._create_mcp_client = AsyncMock(side_effect=capture_create_mcp_client)
await manager._call_regular_mcp_tool(
mcp_server=server,
original_tool_name="tool",
arguments={},
tasks=[],
mcp_auth_header=None,
mcp_server_auth_headers=None,
oauth2_headers=None,
raw_headers={
"x-request-id": "request-123",
"x-litellm-tags": "application:orders,service:checkout",
"x-litellm-end-user-id": "user-123",
},
proxy_logging_obj=None,
)
assert captured_extra_headers == {"x-request-id": "request-123"}
def test_openapi_forwarded_extra_headers_excludes_gateway_attribution_headers(self):
server = MCPServer(
server_id="server-attribution-openapi",
name="attribution-openapi-server",
url="https://example.com",
transport=MCPTransport.http,
auth_type=MCPAuth.authorization,
extra_headers=["x-request-id", "X-LiteLLM-Tags", "X-LiteLLM-End-User-Id"],
)
forwarded_headers = _openapi_forwarded_extra_headers(
mcp_server=server,
raw_headers={
"x-request-id": "request-123",
"x-litellm-tags": "application:orders,service:checkout",
"x-litellm-end-user-id": "user-123",
},
user_api_key_auth=None,
)
assert forwarded_headers == {"x-request-id": "request-123"}
async def _capture_list_subject_token(self, server, oauth2_headers, raw_headers=None):
"""Run _get_tools_from_server and return the subject_token it threaded to _create_mcp_client."""
manager = MCPServerManager()