From c8ca0d6228096e00fc98a130a6b4575e2f5ddbf1 Mon Sep 17 00:00:00 2001 From: MokshaVora Date: Mon, 10 Aug 2026 16:30:11 +0530 Subject: [PATCH] fix(mcp): isolate gateway attribution headers --- README.md | 6 +- .../mcp_server/mcp_server_manager.py | 8 +++ .../proxy/_experimental/mcp_server/server.py | 3 +- .../mcp_server/test_mcp_server_manager.py | 67 ++++++++++++++++++- 4 files changed, 77 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index 1a7e84da01f..32b0160dbaa 100644 --- a/README.md +++ b/README.md @@ -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) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index c8ff6e262d2..8c60caf9692 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -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()) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 60fa45ace52..fa863536c44 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -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 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 54fd5242d5f..99b12dfdbe4 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -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()