diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 12fd1ef89a5..b3eca56a182 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -497,7 +497,8 @@ class MCPServerManager: # Add any static headers from server config. # # Note: `extra_headers` on MCPServer is a List[str] of header names to forward - # from the client request (not available in this OpenAPI tool generation step). + # from each client MCP request; values are applied at call time via + # `_request_extra_headers` in server.py (not baked in here). # `static_headers` is a dict of concrete headers to always send. headers = ( merge_mcp_headers( diff --git a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py index 718435cce6f..973f6ee59df 100644 --- a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py +++ b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py @@ -55,6 +55,13 @@ _request_auth_header: contextvars.ContextVar[Optional[str]] = contextvars.Contex "_request_auth_header", default=None ) +# Per-request extra headers forwarded from the client request. +# Populated from MCPServer.extra_headers names matched against raw request +# headers in server.py before dispatching to a local/OpenAPI tool handler. +_request_extra_headers: contextvars.ContextVar[Optional[Dict[str, str]]] = ( + contextvars.ContextVar("_request_extra_headers", default=None) +) + def _sanitize_path_parameter_value(param_value: Any, param_name: str) -> str: """Ensure path params cannot introduce directory traversal.""" @@ -297,6 +304,20 @@ def build_input_schema(operation: Dict[str, Any]) -> Dict[str, Any]: } +def _merge_openapi_tool_request_headers( + static_headers: Dict[str, str] +) -> Dict[str, str]: + """Merge static closure headers with per-request ContextVar overrides.""" + effective_headers = dict(static_headers) + request_extra = _request_extra_headers.get() + if request_extra: + effective_headers.update(request_extra) + override_auth = _request_auth_header.get() + if override_auth: + effective_headers["Authorization"] = override_auth + return effective_headers + + def create_tool_function( path: str, method: str, @@ -334,14 +355,7 @@ def create_tool_function( The function safely handles parameter names that aren't valid Python identifiers by using **kwargs instead of named parameters. """ - # Allow per-request auth override (e.g. BYOK credential set via ContextVar). - # The ContextVar holds the full Authorization header value, including the - # correct prefix (Bearer / ApiKey / Basic) formatted by the caller in - # server.py based on the server's configured auth_type. - effective_headers = dict(headers) - override_auth = _request_auth_header.get() - if override_auth: - effective_headers["Authorization"] = override_auth + effective_headers = _merge_openapi_tool_request_headers(headers) # Build URL from base_url and path url = base_url + path diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 54d9bbe6e28..b0dc3eff189 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -158,6 +158,7 @@ if MCP_AVAILABLE: ) from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( _request_auth_header, + _request_extra_headers, ) from litellm.proxy._experimental.mcp_server.sse_transport import SseServerTransport from litellm.proxy._experimental.mcp_server.tool_registry import ( @@ -2195,11 +2196,42 @@ if MCP_AVAILABLE: auth_header_value = f"Basic {mcp_auth_header}" else: auth_header_value = f"Bearer {mcp_auth_header}" + + # Forward named client headers to OpenAPI tool upstream requests. + # MCPServer.extra_headers lists header names to copy from raw_headers. + # OAuth2 M2M: never take Authorization from the caller (matches + # _prepare_mcp_server_headers for managed MCP). + forwarded_headers: Optional[Dict[str, str]] = None + if mcp_server and mcp_server.extra_headers and raw_headers: + normalized_raw = { + str(k).lower(): v + for k, v in raw_headers.items() + if isinstance(k, str) + } + skip_caller_authorization = bool( + getattr(mcp_server, "has_client_credentials", False) + ) + for header_name in mcp_server.extra_headers: + if not isinstance(header_name, str): + continue + if ( + skip_caller_authorization + and header_name.lower() == "authorization" + ): + continue + value = normalized_raw.get(header_name.lower()) + if value is not None: + if forwarded_headers is None: + forwarded_headers = {} + forwarded_headers[header_name] = value + _auth_token = _request_auth_header.set(auth_header_value) + _extra_token = _request_extra_headers.set(forwarded_headers) try: local_content = await _handle_local_mcp_tool(name, arguments) finally: _request_auth_header.reset(_auth_token) + _request_extra_headers.reset(_extra_token) response = CallToolResult(content=cast(Any, local_content), isError=False) # Try managed MCP server tool (pass the full prefixed name) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py index 957dea22f3c..b91501b8d41 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py @@ -15,6 +15,8 @@ from unittest.mock import AsyncMock, patch import pytest from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( + _request_auth_header, + _request_extra_headers, _resolve_param_list, _resolve_ref, build_input_schema, @@ -1011,3 +1013,140 @@ class TestRegisterToolsFromOpenAPI: assert re.match( r"^[a-zA-Z0-9_-]+$", name ), f"fallback tool name {name!r} not sanitized" + + +class TestRequestExtraHeaders: + """Tests for _request_extra_headers ContextVar forwarding in tool_function.""" + + @pytest.mark.asyncio + async def test_extra_headers_forwarded_to_upstream(self): + """Extra headers set via ContextVar are included in the upstream request.""" + operation = {} + func = create_tool_function( + path="/data", + method="get", + operation=operation, + base_url="https://api.example.com", + ) + + with patch(GET_ASYNC_CLIENT_TARGET) as mock_client: + async_client = _create_mock_client("get", "ok") + mock_client.return_value = async_client + + token = _request_extra_headers.set({"X-TOKEN": "secret-value"}) + try: + result = await func() + finally: + _request_extra_headers.reset(token) + + assert result == "ok" + call_args = async_client.get.call_args + headers_sent = call_args[1]["headers"] + assert headers_sent.get("X-TOKEN") == "secret-value" + + @pytest.mark.asyncio + async def test_no_extra_headers_by_default(self): + """Without setting _request_extra_headers, no extra headers are injected.""" + operation = {} + func = create_tool_function( + path="/data", + method="get", + operation=operation, + base_url="https://api.example.com", + headers={"X-Static": "static-value"}, + ) + + with patch(GET_ASYNC_CLIENT_TARGET) as mock_client: + async_client = _create_mock_client("get", "ok") + mock_client.return_value = async_client + + result = await func() + + assert result == "ok" + call_args = async_client.get.call_args + headers_sent = call_args[1]["headers"] + assert headers_sent == {"X-Static": "static-value"} + assert "X-TOKEN" not in headers_sent + + @pytest.mark.asyncio + async def test_extra_headers_merged_with_static_headers(self): + """Request extra headers are merged on top of static (baked-in) headers.""" + operation = {} + func = create_tool_function( + path="/data", + method="post", + operation=operation, + base_url="https://api.example.com", + headers={"X-Static": "static-value"}, + ) + + with patch(GET_ASYNC_CLIENT_TARGET) as mock_client: + async_client = _create_mock_client("post", "created") + mock_client.return_value = async_client + + token = _request_extra_headers.set({"X-TOKEN": "dynamic-value"}) + try: + result = await func() + finally: + _request_extra_headers.reset(token) + + assert result == "created" + call_args = async_client.post.call_args + headers_sent = call_args[1]["headers"] + assert headers_sent.get("X-Static") == "static-value" + assert headers_sent.get("X-TOKEN") == "dynamic-value" + + @pytest.mark.asyncio + async def test_auth_header_still_overrides_extra_headers(self): + """_request_auth_header takes precedence for Authorization over extra headers.""" + operation = {} + func = create_tool_function( + path="/secure", + method="get", + operation=operation, + base_url="https://api.example.com", + ) + + with patch(GET_ASYNC_CLIENT_TARGET) as mock_client: + async_client = _create_mock_client("get", "secure-data") + mock_client.return_value = async_client + + extra_token = _request_extra_headers.set( + {"Authorization": "Bearer extra", "X-TOKEN": "token-value"} + ) + auth_token = _request_auth_header.set("Bearer byok-credential") + try: + result = await func() + finally: + _request_auth_header.reset(auth_token) + _request_extra_headers.reset(extra_token) + + assert result == "secure-data" + call_args = async_client.get.call_args + headers_sent = call_args[1]["headers"] + assert headers_sent.get("Authorization") == "Bearer byok-credential" + assert headers_sent.get("X-TOKEN") == "token-value" + + @pytest.mark.asyncio + async def test_extra_headers_not_leaked_between_calls(self): + """After resetting the ContextVar, subsequent calls do not see the headers.""" + operation = {} + func = create_tool_function( + path="/data", + method="get", + operation=operation, + base_url="https://api.example.com", + ) + + with patch(GET_ASYNC_CLIENT_TARGET) as mock_client: + async_client = _create_mock_client("get", "ok") + mock_client.return_value = async_client + + token = _request_extra_headers.set({"X-TOKEN": "first-call"}) + _request_extra_headers.reset(token) + + await func() + + call_args = async_client.get.call_args + headers_sent = call_args[1]["headers"] + assert "X-TOKEN" not in headers_sent