From 0ff7177373d1d83fa15005c4115f0e54f2b08ef1 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 20 Sep 2025 20:00:25 -0700 Subject: [PATCH] feat(user_api_key_auth_mcp.py): pass extra headers from clientside straight through - allow multiple clientside headers Closes LIT-952 --- litellm/experimental_mcp_client/client.py | 38 +++++++++++-------- .../mcp_server/auth/litellm_auth_handler.py | 6 +-- .../mcp_server/auth/user_api_key_auth_mcp.py | 25 +++++++++--- .../mcp_server/mcp_server_manager.py | 10 ++--- .../proxy/_experimental/mcp_server/server.py | 16 ++++---- .../auth/test_user_api_key_auth_mcp.py | 32 ++++++++++++++++ 6 files changed, 89 insertions(+), 38 deletions(-) diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index 2e02f460b6..1176248d4f 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -5,7 +5,7 @@ LiteLLM Proxy uses this MCP Client to connnect to other MCP servers. import asyncio import base64 from datetime import timedelta -from typing import Dict, List, Optional +from typing import Dict, List, Optional, Union from mcp import ClientSession, StdioServerParameters from mcp.client.sse import sse_client @@ -44,7 +44,7 @@ class MCPClient: server_url: str = "", transport_type: MCPTransportType = MCPTransport.http, auth_type: MCPAuthType = None, - auth_value: Optional[str] = None, + auth_value: Optional[Union[str, Dict[str, str]]] = None, timeout: float = 60.0, stdio_config: Optional[MCPStdioConfig] = None, extra_headers: Optional[Dict[str, str]] = None, @@ -53,7 +53,7 @@ class MCPClient: self.transport_type: MCPTransport = transport_type self.auth_type: MCPAuthType = auth_type self.timeout: float = timeout - self._mcp_auth_value: Optional[str] = None + self._mcp_auth_value: Optional[Union[str, Dict[str, str]]] = None self._session: Optional[ClientSession] = None self._context = None self._transport_ctx = None @@ -180,28 +180,34 @@ class MCPClient: pass self._context = None - def update_auth_value(self, mcp_auth_value: str): + def update_auth_value(self, mcp_auth_value: Union[str, Dict[str, str]]): """ Set the authentication header for the MCP client. """ - if self.auth_type == MCPAuth.basic: - # Assuming mcp_auth_value is in format "username:password", convert it when updating - mcp_auth_value = to_basic_auth(mcp_auth_value) - self._mcp_auth_value = mcp_auth_value + if isinstance(mcp_auth_value, dict): + self._mcp_auth_value = mcp_auth_value + else: + if self.auth_type == MCPAuth.basic: + # Assuming mcp_auth_value is in format "username:password", convert it when updating + mcp_auth_value = to_basic_auth(mcp_auth_value) + self._mcp_auth_value = mcp_auth_value def _get_auth_headers(self) -> dict: """Generate authentication headers based on auth type.""" headers = {"MCP-Protocol-Version": "2025-06-18"} if self._mcp_auth_value: - if self.auth_type == MCPAuth.bearer_token: - headers["Authorization"] = f"Bearer {self._mcp_auth_value}" - elif self.auth_type == MCPAuth.basic: - headers["Authorization"] = f"Basic {self._mcp_auth_value}" - elif self.auth_type == MCPAuth.api_key: - headers["X-API-Key"] = self._mcp_auth_value - elif self.auth_type == MCPAuth.authorization: - headers["Authorization"] = self._mcp_auth_value + if isinstance(self._mcp_auth_value, str): + if self.auth_type == MCPAuth.bearer_token: + headers["Authorization"] = f"Bearer {self._mcp_auth_value}" + elif self.auth_type == MCPAuth.basic: + headers["Authorization"] = f"Basic {self._mcp_auth_value}" + elif self.auth_type == MCPAuth.api_key: + headers["X-API-Key"] = self._mcp_auth_value + elif self.auth_type == MCPAuth.authorization: + headers["Authorization"] = self._mcp_auth_value + elif isinstance(self._mcp_auth_value, dict): + headers.update(self._mcp_auth_value) # update the headers with the extra headers if self.extra_headers: diff --git a/litellm/proxy/_experimental/mcp_server/auth/litellm_auth_handler.py b/litellm/proxy/_experimental/mcp_server/auth/litellm_auth_handler.py index 56a22040f0..2f4c6c2d8d 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/litellm_auth_handler.py +++ b/litellm/proxy/_experimental/mcp_server/auth/litellm_auth_handler.py @@ -1,4 +1,4 @@ -from typing import Dict, List, Optional +from typing import Dict, List, Optional, Union from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser @@ -22,9 +22,9 @@ class MCPAuthenticatedUser(AuthenticatedUser): user_api_key_auth: UserAPIKeyAuth, mcp_auth_header: Optional[str] = None, mcp_servers: Optional[List[str]] = None, - mcp_server_auth_headers: Optional[Dict[str, str]] = None, - mcp_protocol_version: Optional[str] = None, + mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None, oauth2_headers: Optional[Dict[str, str]] = None, + mcp_protocol_version: Optional[str] = None, ): self.user_api_key_auth = user_api_key_auth self.mcp_auth_header = mcp_auth_header diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 281edf38ad..3417dad2e4 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -39,7 +39,7 @@ class MCPRequestHandler: UserAPIKeyAuth, Optional[str], Optional[List[str]], - Optional[Dict[str, str]], + Optional[Dict[str, Dict[str, str]]], Optional[Dict[str, str]], ]: """ @@ -145,7 +145,9 @@ class MCPRequestHandler: return auth_header @staticmethod - def _get_mcp_server_auth_headers_from_headers(headers: Headers) -> Dict[str, str]: + def _get_mcp_server_auth_headers_from_headers( + headers: Headers, + ) -> Dict[str, Dict[str, str]]: """ Parse server-specific MCP auth headers from the request headers. @@ -156,7 +158,7 @@ class MCPRequestHandler: - x-mcp-deepwiki-authorization: Basic base64_encoded_creds Returns: - Dict[str, str]: Mapping of server alias to auth value + Dict[str, Dict[str, str]]: Mapping of server alias to header dict """ server_auth_headers = {} prefix = "x-mcp-" @@ -175,11 +177,22 @@ class MCPRequestHandler: # Extract server_alias and header_name from x-mcp-{server_alias}-{header_name} remaining = header_name[len(prefix) :].lower() if "-" in remaining: - # Split on the last dash to separate server_alias from header_name - parts = remaining.rsplit("-", 1) + # Split on the first dash to separate server_alias from header_name + parts = remaining.split("-", 1) if len(parts) == 2: server_alias, auth_header_name = parts - server_auth_headers[server_alias] = header_value + + # Convert header name to proper case (e.g., "authorization" -> "Authorization") + if auth_header_name == "authorization": + auth_header_name = "Authorization" + + # Initialize server dict if not exists + if server_alias not in server_auth_headers: + server_auth_headers[server_alias] = {} + + server_auth_headers[server_alias][ + auth_header_name + ] = header_value verbose_logger.debug( f"Found server auth header: {server_alias} -> {auth_header_name}: {header_value[:10]}..." ) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 9fb39e7498..891adf8328 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -10,7 +10,7 @@ import asyncio import datetime import hashlib import json -from typing import Any, Dict, List, Optional, cast +from typing import Any, Dict, List, Optional, Union, cast from fastapi import HTTPException from mcp.types import CallToolRequestParams as MCPCallToolRequestParams @@ -319,7 +319,7 @@ class MCPServerManager: self, user_api_key_auth: Optional[UserAPIKeyAuth] = None, mcp_auth_header: Optional[str] = None, - mcp_server_auth_headers: Optional[Dict[str, str]] = None, + mcp_server_auth_headers: Optional[Dict[str, Union[str, Dict[str, str]]]] = None, ) -> List[MCPTool]: """ List all tools available across all MCP Servers. @@ -381,7 +381,7 @@ class MCPServerManager: def _create_mcp_client( self, server: MCPServer, - mcp_auth_header: Optional[str] = None, + mcp_auth_header: Optional[Union[str, Dict[str, str]]] = None, extra_headers: Optional[Dict[str, str]] = None, ) -> MCPClient: """ @@ -429,7 +429,7 @@ class MCPServerManager: async def _get_tools_from_server( self, server: MCPServer, - mcp_auth_header: Optional[str] = None, + mcp_auth_header: Optional[Union[str, Dict[str, str]]] = None, extra_headers: Optional[Dict[str, str]] = None, ) -> List[MCPTool]: """ @@ -638,7 +638,7 @@ class MCPServerManager: arguments: Dict[str, Any], user_api_key_auth: Optional[UserAPIKeyAuth] = None, mcp_auth_header: Optional[str] = None, - mcp_server_auth_headers: Optional[Dict[str, str]] = None, + mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None, proxy_logging_obj: Optional[ProxyLogging] = None, oauth2_headers: Optional[Dict[str, str]] = None, ) -> CallToolResult: diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 38b0cfd99a..de1c3e67fc 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -361,7 +361,7 @@ if MCP_AVAILABLE: user_api_key_auth: Optional[UserAPIKeyAuth], mcp_auth_header: Optional[str], mcp_servers: Optional[List[str]], - mcp_server_auth_headers: Optional[Dict[str, str]] = None, + mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None, oauth2_headers: Optional[Dict[str, str]] = None, ) -> List[MCPTool]: """ @@ -371,8 +371,8 @@ if MCP_AVAILABLE: user_api_key_auth: User authentication info for access control mcp_auth_header: Optional auth header for MCP server (deprecated) mcp_servers: Optional list of server names/aliases to filter by - mcp_server_auth_headers: Optional dict of server-specific auth headers - oauth2_headers: Optional dict of oauth2 headers + mcp_server_auth_headers: Optional dict of server-specific auth headers + oauth2_headers: Optional dict of oauth2 headers Returns: List[MCPTool]: Combined list of tools from filtered servers @@ -438,7 +438,7 @@ if MCP_AVAILABLE: user_api_key_auth: Optional[UserAPIKeyAuth] = None, mcp_auth_header: Optional[str] = None, mcp_servers: Optional[List[str]] = None, - mcp_server_auth_headers: Optional[Dict[str, str]] = None, + mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None, oauth2_headers: Optional[Dict[str, str]] = None, ) -> List[MCPTool]: """ @@ -505,7 +505,7 @@ if MCP_AVAILABLE: arguments: Optional[Dict[str, Any]] = None, user_api_key_auth: Optional[UserAPIKeyAuth] = None, mcp_auth_header: Optional[str] = None, - mcp_server_auth_headers: Optional[Dict[str, str]] = None, + mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None, oauth2_headers: Optional[Dict[str, str]] = None, **kwargs: Any, ) -> List[Union[TextContent, ImageContent, EmbeddedResource]]: @@ -606,7 +606,7 @@ if MCP_AVAILABLE: arguments: Dict[str, Any], user_api_key_auth: Optional[UserAPIKeyAuth] = None, mcp_auth_header: Optional[str] = None, - mcp_server_auth_headers: Optional[Dict[str, str]] = None, + mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None, oauth2_headers: Optional[Dict[str, str]] = None, litellm_logging_obj: Optional[Any] = None, ) -> List[Union[TextContent, ImageContent, EmbeddedResource]]: @@ -858,7 +858,7 @@ if MCP_AVAILABLE: user_api_key_auth: UserAPIKeyAuth, mcp_auth_header: Optional[str] = None, mcp_servers: Optional[List[str]] = None, - mcp_server_auth_headers: Optional[Dict[str, str]] = None, + mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None, oauth2_headers: Optional[Dict[str, str]] = None, ) -> None: """ @@ -883,7 +883,7 @@ if MCP_AVAILABLE: Optional[UserAPIKeyAuth], Optional[str], Optional[List[str]], - Optional[Dict[str, str]], + Optional[Dict[str, Dict[str, str]]], Optional[Dict[str, str]], ]: """ diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index afafa510a4..7bdefe84c1 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -1014,3 +1014,35 @@ def test_mcp_path_based_server_segregation(monkeypatch): # The context should have mcp_servers set to ["zapier", "group1"] assert list(captured_mcp_servers.values())[0] == ["zapier", "group1"] + + +@pytest.mark.parametrize( + "headers,expected_result", + [ + ( + Headers( + { + "x-litellm-api-key": "test-key", + "x-mcp-github-authorization": "Bearer github-token", + } + ), + {"github": {"Authorization": "Bearer github-token"}}, + ), + ( + Headers( + { + "x-litellm-api-key": "test-key", + "x-mcp-github-x-api-key": "Basic base64-encoded-creds", + } + ), + {"github": {"x-api-key": "Basic base64-encoded-creds"}}, + ), + ], +) +def test_get_mcp_server_auth_headers_from_headers(headers, expected_result): + """Test _get_mcp_server_auth_headers_from_headers method""" + from starlette.datastructures import Headers + + headers = Headers(headers) + result = MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers) + assert result == expected_result