From 0ff7177373d1d83fa15005c4115f0e54f2b08ef1 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 20 Sep 2025 20:00:25 -0700 Subject: [PATCH 01/12] 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 2e02f460b65..1176248d4f1 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 56a22040f0d..2f4c6c2d8d5 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 281edf38adb..3417dad2e4f 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 9fb39e74980..891adf83289 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 38b0cfd99a7..de1c3e67fcc 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 afafa510a41..7bdefe84c1d 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 From 1c71d4cdbbbc87575ee97fe17429d4c85c006ae2 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 27 Sep 2025 17:30:46 -0700 Subject: [PATCH 02/12] feat(mcp/): allow specifying forwardable headers allows admin to specify which clientside headers to forward to the backend mcp server easier than requiring client to specify `x-mcp-{server_alias}-key` --- .../mcp_server/auth/litellm_auth_handler.py | 3 +++ .../mcp_server/auth/user_api_key_auth_mcp.py | 6 ++++- .../mcp_server/mcp_server_manager.py | 16 +++++++++++++ .../proxy/_experimental/mcp_server/server.py | 24 ++++++++++++++++++- .../types/mcp_server/mcp_server_manager.py | 3 +++ 5 files changed, 50 insertions(+), 2 deletions(-) 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 2f4c6c2d8d5..56aeb77f422 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/litellm_auth_handler.py +++ b/litellm/proxy/_experimental/mcp_server/auth/litellm_auth_handler.py @@ -15,6 +15,7 @@ class MCPAuthenticatedUser(AuthenticatedUser): 3. MCP server configuration (can include access groups) 4. Server-specific authentication headers 5. OAuth2 headers + 6. Raw headers - allows forwarding specific headers to the MCP server, specified by the admin. """ def __init__( @@ -25,6 +26,7 @@ class MCPAuthenticatedUser(AuthenticatedUser): mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None, oauth2_headers: Optional[Dict[str, str]] = None, mcp_protocol_version: Optional[str] = None, + raw_headers: Optional[Dict[str, str]] = None, ): self.user_api_key_auth = user_api_key_auth self.mcp_auth_header = mcp_auth_header @@ -32,3 +34,4 @@ class MCPAuthenticatedUser(AuthenticatedUser): self.mcp_server_auth_headers = mcp_server_auth_headers or {} self.mcp_protocol_version = mcp_protocol_version self.oauth2_headers = oauth2_headers + self.raw_headers = raw_headers 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 3417dad2e4f..4b2c3385bbc 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 @@ -41,6 +41,7 @@ class MCPRequestHandler: Optional[List[str]], Optional[Dict[str, Dict[str, str]]], Optional[Dict[str, str]], + Optional[Dict[str, str]], ]: """ Process and validate MCP request headers from the ASGI scope. @@ -49,6 +50,7 @@ class MCPRequestHandler: 2. Processing MCP server configuration 3. Handling MCP-specific headers 4. Handling oauth2 headers + 5. Raw headers - allows forwarding specific headers to the MCP server, specified by the admin. Args: scope: ASGI scope containing request information @@ -58,7 +60,8 @@ class MCPRequestHandler: mcp_auth_header: Optional[str] MCP auth header to be passed to the MCP server (deprecated) mcp_servers: Optional[List[str]] List of MCP servers and access groups to use mcp_server_auth_headers: Optional[Dict[str, str]] Server-specific auth headers in format {server_alias: auth_value} - + oauth2_headers: Optional[Dict[str, str]] OAuth2 headers + raw_headers: Optional[Dict[str, str]] Raw headers to be forwarded to the MCP server Raises: HTTPException: If headers are invalid or missing required headers """ @@ -116,6 +119,7 @@ class MCPRequestHandler: mcp_servers, mcp_server_auth_headers, oauth2_headers, + dict(headers), ) @staticmethod diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 891adf83289..c971abe4303 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -212,6 +212,7 @@ class MCPServerManager: "authentication_token", server_config.get("auth_value", None) ), mcp_info=mcp_info, + forwardable_headers=server_config.get("forwardable_headers", None), access_groups=server_config.get("access_groups", None), ) self.config_mcp_servers[server_id] = new_server @@ -264,6 +265,13 @@ class MCPServerManager: transport=cast(MCPTransportType, mcp_server.transport), auth_type=cast(MCPAuthType, mcp_server.auth_type), mcp_info=mcp_info, + forwardable_headers=getattr(mcp_server, "forwardable_headers", None), + # oauth specific fields + client_id=getattr(mcp_server, "client_id", None), + client_secret=getattr(mcp_server, "client_secret", None), + scopes=getattr(mcp_server, "scopes", None), + authorization_url=getattr(mcp_server, "authorization_url", None), + token_url=getattr(mcp_server, "token_url", None), # Stdio-specific fields command=getattr(mcp_server, "command", None), args=getattr(mcp_server, "args", None) or [], @@ -641,6 +649,7 @@ class MCPServerManager: mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None, proxy_logging_obj: Optional[ProxyLogging] = None, oauth2_headers: Optional[Dict[str, str]] = None, + raw_headers: Optional[Dict[str, str]] = None, ) -> CallToolResult: """ Call a tool with the given name and arguments (handles prefixed tool names) @@ -709,6 +718,13 @@ class MCPServerManager: if mcp_server.auth_type == MCPAuth.oauth2: extra_headers = oauth2_headers + if mcp_server.forwardable_headers and raw_headers: + if extra_headers is None: + extra_headers = {} + for header in mcp_server.forwardable_headers: + if header in raw_headers: + extra_headers[header] = raw_headers[header] + client = self._create_mcp_client( server=mcp_server, mcp_auth_header=server_auth_header, diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index de1c3e67fcc..62ad1cc9107 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -180,6 +180,7 @@ if MCP_AVAILABLE: mcp_servers, mcp_server_auth_headers, oauth2_headers, + raw_headers, ) = get_auth_context() verbose_logger.debug( f"MCP list_tools - User API Key Auth from context: {user_api_key_auth}" @@ -198,6 +199,7 @@ if MCP_AVAILABLE: mcp_servers=mcp_servers, mcp_server_auth_headers=mcp_server_auth_headers, oauth2_headers=oauth2_headers, + raw_headers=raw_headers, ) verbose_logger.info( f"MCP list_tools - Successfully returned {len(tools)} tools" @@ -239,6 +241,7 @@ if MCP_AVAILABLE: _, mcp_server_auth_headers, oauth2_headers, + raw_headers, ) = get_auth_context() verbose_logger.debug( @@ -271,6 +274,7 @@ if MCP_AVAILABLE: mcp_auth_header=mcp_auth_header, mcp_server_auth_headers=mcp_server_auth_headers, oauth2_headers=oauth2_headers, + raw_headers=raw_headers, **data, # for logging ) except BlockedPiiEntityError as e: @@ -363,6 +367,7 @@ if MCP_AVAILABLE: mcp_servers: Optional[List[str]], mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None, oauth2_headers: Optional[Dict[str, str]] = None, + raw_headers: Optional[Dict[str, str]] = None, ) -> List[MCPTool]: """ Helper method to fetch tools from MCP servers based on server filtering criteria. @@ -440,6 +445,7 @@ if MCP_AVAILABLE: mcp_servers: Optional[List[str]] = None, mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None, oauth2_headers: Optional[Dict[str, str]] = None, + raw_headers: Optional[Dict[str, str]] = None, ) -> List[MCPTool]: """ List all available MCP tools. @@ -464,6 +470,7 @@ if MCP_AVAILABLE: mcp_servers=mcp_servers, mcp_server_auth_headers=mcp_server_auth_headers, oauth2_headers=oauth2_headers, + raw_headers=raw_headers, ) verbose_logger.debug( f"Successfully fetched {len(managed_tools)} tools from managed MCP servers" @@ -507,6 +514,7 @@ if MCP_AVAILABLE: mcp_auth_header: Optional[str] = None, mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None, oauth2_headers: Optional[Dict[str, str]] = None, + raw_headers: Optional[Dict[str, str]] = None, **kwargs: Any, ) -> List[Union[TextContent, ImageContent, EmbeddedResource]]: """ @@ -555,6 +563,7 @@ if MCP_AVAILABLE: mcp_auth_header=mcp_auth_header, mcp_server_auth_headers=mcp_server_auth_headers, oauth2_headers=oauth2_headers, + raw_headers=raw_headers, litellm_logging_obj=litellm_logging_obj, ) @@ -608,6 +617,7 @@ if MCP_AVAILABLE: mcp_auth_header: Optional[str] = None, mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None, oauth2_headers: Optional[Dict[str, str]] = None, + raw_headers: Optional[Dict[str, str]] = None, litellm_logging_obj: Optional[Any] = None, ) -> List[Union[TextContent, ImageContent, EmbeddedResource]]: """Handle tool execution for managed server tools""" @@ -621,6 +631,7 @@ if MCP_AVAILABLE: mcp_auth_header=mcp_auth_header, mcp_server_auth_headers=mcp_server_auth_headers, oauth2_headers=oauth2_headers, + raw_headers=raw_headers, proxy_logging_obj=proxy_logging_obj, ) verbose_logger.debug("CALL TOOL RESULT: %s", call_tool_result) @@ -702,6 +713,7 @@ if MCP_AVAILABLE: _, mcp_server_auth_headers, oauth2_headers, + raw_headers, ) = await MCPRequestHandler.process_mcp_request(scope) mcp_servers = mcp_servers_from_path else: @@ -711,6 +723,7 @@ if MCP_AVAILABLE: mcp_servers, mcp_server_auth_headers, oauth2_headers, + raw_headers, ) = await MCPRequestHandler.process_mcp_request(scope) return ( user_api_key_auth, @@ -718,6 +731,7 @@ if MCP_AVAILABLE: mcp_servers, mcp_server_auth_headers, oauth2_headers, + raw_headers, ) async def handle_streamable_http_mcp( @@ -732,6 +746,7 @@ if MCP_AVAILABLE: mcp_servers, mcp_server_auth_headers, oauth2_headers, + raw_headers, ) = await extract_mcp_auth_context(scope, path) verbose_logger.debug( f"MCP request mcp_servers (header/path): {mcp_servers}" @@ -746,6 +761,7 @@ if MCP_AVAILABLE: mcp_servers=mcp_servers, mcp_server_auth_headers=mcp_server_auth_headers, oauth2_headers=oauth2_headers, + raw_headers=raw_headers, ) # Ensure session managers are initialized @@ -785,6 +801,7 @@ if MCP_AVAILABLE: mcp_servers, mcp_server_auth_headers, oauth2_headers, + raw_headers, ) = await extract_mcp_auth_context(scope, path) verbose_logger.debug( f"MCP request mcp_servers (header/path): {mcp_servers}" @@ -798,6 +815,7 @@ if MCP_AVAILABLE: mcp_servers=mcp_servers, mcp_server_auth_headers=mcp_server_auth_headers, oauth2_headers=oauth2_headers, + raw_headers=raw_headers, ) if not _SESSION_MANAGERS_INITIALIZED: @@ -860,6 +878,7 @@ if MCP_AVAILABLE: mcp_servers: Optional[List[str]] = None, mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None, oauth2_headers: Optional[Dict[str, str]] = None, + raw_headers: Optional[Dict[str, str]] = None, ) -> None: """ Set the UserAPIKeyAuth in the auth context variable. @@ -876,6 +895,7 @@ if MCP_AVAILABLE: mcp_servers=mcp_servers, mcp_server_auth_headers=mcp_server_auth_headers, oauth2_headers=oauth2_headers, + raw_headers=raw_headers, ) auth_context_var.set(auth_user) @@ -885,6 +905,7 @@ if MCP_AVAILABLE: Optional[List[str]], Optional[Dict[str, Dict[str, str]]], Optional[Dict[str, str]], + Optional[Dict[str, str]], ]: """ Get the UserAPIKeyAuth from the auth context variable. @@ -901,8 +922,9 @@ if MCP_AVAILABLE: auth_user.mcp_servers, auth_user.mcp_server_auth_headers, auth_user.oauth2_headers, + auth_user.raw_headers, ) - return None, None, None, None, None + return None, None, None, None, None, None ######################################################## ############ End of Auth Context Functions ############# diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 0dec1b23c6c..dbd849646d8 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -20,6 +20,9 @@ class MCPServer(BaseModel): auth_type: Optional[MCPAuthType] = None authentication_token: Optional[str] = None mcp_info: Optional[MCPInfo] = None + forwardable_headers: Optional[List[str]] = ( + None # allow admin to specify which headers to forward to the MCP server + ) # OAuth-specific fields client_id: Optional[str] = None client_secret: Optional[str] = None From c6a7e676f6d833464f4ebeae885f8a3e660165dc Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 27 Sep 2025 17:55:04 -0700 Subject: [PATCH 03/12] fix(server.py): allow user to specify custom headers to forward to mcp server --- litellm/proxy/_experimental/mcp_server/server.py | 7 +++++++ litellm/proxy/_new_secret_config.yaml | 1 + 2 files changed, 8 insertions(+) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 62ad1cc9107..913673f2811 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -414,6 +414,13 @@ if MCP_AVAILABLE: if server.auth_type == MCPAuth.oauth2: extra_headers = oauth2_headers + if server.forwardable_headers and raw_headers: + if extra_headers is None: + extra_headers = {} + for header in server.forwardable_headers: + if header in raw_headers: + extra_headers[header] = raw_headers[header] + # Fall back to deprecated mcp_auth_header if no server-specific header found if server_auth_header is None: server_auth_header = mcp_auth_header diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index c50d391e058..533d28395e7 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -25,6 +25,7 @@ mcp_servers: client_id: os.environ/GITHUB_OAUTH_CLIENT_ID client_secret: os.environ/GITHUB_OAUTH_CLIENT_SECRET scopes: ["public_repo", "user:email"] + forwardable_headers: ["custom_key"] # allowed_tools: ["list_tools"] # disallowed_tools: ["repo_delete"] From 327447cb175f1cc75617f30d9cd0d72a459e7e54 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 27 Sep 2025 17:58:11 -0700 Subject: [PATCH 04/12] fix(mcp/): rename param to be 'extra_headers', instead of 'forwardable_headers' --- .../proxy/_experimental/mcp_server/mcp_server_manager.py | 8 ++++---- litellm/proxy/_experimental/mcp_server/server.py | 4 ++-- litellm/proxy/_new_secret_config.yaml | 2 +- litellm/types/mcp_server/mcp_server_manager.py | 2 +- 4 files changed, 8 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index c971abe4303..a3a5d93a3fd 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -212,7 +212,7 @@ class MCPServerManager: "authentication_token", server_config.get("auth_value", None) ), mcp_info=mcp_info, - forwardable_headers=server_config.get("forwardable_headers", None), + extra_headers=server_config.get("extra_headers", None), access_groups=server_config.get("access_groups", None), ) self.config_mcp_servers[server_id] = new_server @@ -265,7 +265,7 @@ class MCPServerManager: transport=cast(MCPTransportType, mcp_server.transport), auth_type=cast(MCPAuthType, mcp_server.auth_type), mcp_info=mcp_info, - forwardable_headers=getattr(mcp_server, "forwardable_headers", None), + extra_headers=getattr(mcp_server, "extra_headers", None), # oauth specific fields client_id=getattr(mcp_server, "client_id", None), client_secret=getattr(mcp_server, "client_secret", None), @@ -718,10 +718,10 @@ class MCPServerManager: if mcp_server.auth_type == MCPAuth.oauth2: extra_headers = oauth2_headers - if mcp_server.forwardable_headers and raw_headers: + if mcp_server.extra_headers and raw_headers: if extra_headers is None: extra_headers = {} - for header in mcp_server.forwardable_headers: + for header in mcp_server.extra_headers: if header in raw_headers: extra_headers[header] = raw_headers[header] diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 913673f2811..a3e51103ed5 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -414,10 +414,10 @@ if MCP_AVAILABLE: if server.auth_type == MCPAuth.oauth2: extra_headers = oauth2_headers - if server.forwardable_headers and raw_headers: + if server.extra_headers and raw_headers: if extra_headers is None: extra_headers = {} - for header in server.forwardable_headers: + for header in server.extra_headers: if header in raw_headers: extra_headers[header] = raw_headers[header] diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index 533d28395e7..e7494bd0dab 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -25,7 +25,7 @@ mcp_servers: client_id: os.environ/GITHUB_OAUTH_CLIENT_ID client_secret: os.environ/GITHUB_OAUTH_CLIENT_SECRET scopes: ["public_repo", "user:email"] - forwardable_headers: ["custom_key"] + extra_headers: ["custom_key"] # allowed_tools: ["list_tools"] # disallowed_tools: ["repo_delete"] diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index dbd849646d8..4327bd5afe7 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -20,7 +20,7 @@ class MCPServer(BaseModel): auth_type: Optional[MCPAuthType] = None authentication_token: Optional[str] = None mcp_info: Optional[MCPInfo] = None - forwardable_headers: Optional[List[str]] = ( + extra_headers: Optional[List[str]] = ( None # allow admin to specify which headers to forward to the MCP server ) # OAuth-specific fields From db27600ce9d3754e3da2db31ad73c7b195b8187d Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 27 Sep 2025 18:13:19 -0700 Subject: [PATCH 05/12] docs(mcp.md): update docs with new 'extra_headers' tutorial --- docs/my-website/docs/mcp.md | 212 ++++++++++++++++++++++++++++++++++-- 1 file changed, 205 insertions(+), 7 deletions(-) diff --git a/docs/my-website/docs/mcp.md b/docs/my-website/docs/mcp.md index 80b4c32d0ab..0ee63b0e852 100644 --- a/docs/my-website/docs/mcp.md +++ b/docs/my-website/docs/mcp.md @@ -137,6 +137,7 @@ mcp_servers: | `basic` | `Authorization: Basic ` | | `authorization` | `Authorization: ` | +- **Extra Headers**: Optional list of additional header names that should be forwarded from client to the MCP server - **Spec Version**: Optional MCP specification version (defaults to `2025-06-18`) Examples for each auth type: @@ -162,6 +163,13 @@ mcp_servers: url: "https://my-mcp-server.com/mcp" auth_type: "authorization" auth_value: "Token example123" # headers={"Authorization": "Token example123"} + + # Example with extra headers forwarding + github_mcp: + url: "https://api.githubcopilot.com/mcp" + auth_type: "bearer_token" + auth_value: "ghp_example_token" + extra_headers: ["custom_key", "x-custom-header"] # These headers will be forwarded from client ``` @@ -771,6 +779,203 @@ When creating API keys, you can assign them to specific access groups for permis /> +## Forwarding Custom Headers to MCP Servers + +LiteLLM supports forwarding additional custom headers from MCP clients to backend MCP servers using the `extra_headers` configuration parameter. This allows you to pass custom authentication tokens, API keys, or other headers that your MCP server requires. + +### Configuration + + + + +Configure `extra_headers` in your MCP server configuration to specify which header names should be forwarded: + +```yaml title="config.yaml with extra_headers" showLineNumbers +mcp_servers: + github_mcp: + url: "https://api.githubcopilot.com/mcp" + auth_type: "bearer_token" + auth_value: "ghp_default_token" + extra_headers: ["custom_key", "x-custom-header", "Authorization"] + description: "GitHub MCP server with custom header forwarding" +``` + + + +Use this when giving users access to a [group of MCP servers](#grouping-mcps-access-groups). + +**Format:** `x-mcp-{server_alias}-{header_name}: value` + +This allows you to use different authentication for different MCP servers. + + +**Examples:** +- `x-mcp-github-authorization: Bearer ghp_xxxxxxxxx` - GitHub MCP server with Bearer token +- `x-mcp-zapier-x-api-key: sk-xxxxxxxxx` - Zapier MCP server with API key +- `x-mcp-deepwiki-authorization: Basic base64_encoded_creds` - DeepWiki MCP server with Basic auth + +```python title="Python Client with Server-Specific Auth" showLineNumbers +from fastmcp import Client +import asyncio + +# Standard MCP configuration with multiple servers +config = { + "mcpServers": { + "mcp_group": { + "url": "http://localhost:4000/mcp", + "headers": { + "x-mcp-servers": "dev_group", # assume this gives access to github, zapier and deepwiki + "x-litellm-api-key": "Bearer sk-1234", + "x-mcp-github-authorization": "Bearer gho_token", + "x-mcp-zapier-x-api-key": "sk-xxxxxxxxx", + "x-mcp-deepwiki-authorization": "Basic base64_encoded_creds", + "custom_key": "value" + } + } + } +} + +# Create a client that connects to all servers +client = Client(config) + + +async def main(): + async with client: + tools = await client.list_tools() + print(f"Available tools: {tools}") + + # call mcp + await client.call_tool( + name="github_mcp-search_issues", + arguments={'query': 'created:>2024-01-01', 'sort': 'created', 'order': 'desc', 'perPage': 30} + ) + +if __name__ == "__main__": + asyncio.run(main()) + +``` + + + +**Benefits:** +- **Server-specific authentication**: Each MCP server can use different auth methods +- **Better security**: No need to share the same auth token across all servers +- **Flexible header names**: Support for different auth header types (authorization, x-api-key, etc.) +- **Clean separation**: Each server's auth is clearly identified + + + + + + + +### Client Usage + +When connecting from MCP clients, include the custom headers that match the `extra_headers` configuration: + + + + +```python title="FastMCP Client with Custom Headers" showLineNumbers +from fastmcp import Client +import asyncio + +# MCP client configuration with custom headers +config = { + "mcpServers": { + "github": { + "url": "http://localhost:4000/github_mcp/mcp", + "headers": { + "x-litellm-api-key": "Bearer sk-1234", + "Authorization": "Bearer gho_token", + "custom_key": "custom_value", + "x-custom-header": "additional_data" + } + } + } +} + +# Create a client that connects to the server +client = Client(config) + +async def main(): + async with client: + # List available tools + tools = await client.list_tools() + print(f"Available tools: {tools}") + + # Call a tool if available + if tools: + result = await client.call_tool(tools[0].name, {}) + print(f"Tool result: {result}") + +# Run the client +asyncio.run(main()) +``` + + + + + +```json title="Cursor MCP Configuration with Custom Headers" showLineNumbers +{ + "mcpServers": { + "GitHub": { + "url": "http://localhost:4000/github_mcp/mcp", + "headers": { + "x-litellm-api-key": "Bearer $LITELLM_API_KEY", + "Authorization": "Bearer $GITHUB_TOKEN", + "custom_key": "custom_value", + "x-custom-header": "additional_data" + } + } + } +} +``` + + + + + +```bash title="cURL with Custom Headers" showLineNumbers +curl --location 'http://localhost:4000/github_mcp/mcp' \ +--header 'Content-Type: application/json' \ +--header 'x-litellm-api-key: Bearer sk-1234' \ +--header 'Authorization: Bearer gho_token' \ +--header 'custom_key: custom_value' \ +--header 'x-custom-header: additional_data' \ +--data '{ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/list" +}' +``` + + + + +### How It Works + +1. **Configuration**: Define `extra_headers` in your MCP server config with the header names you want to forward +2. **Client Headers**: Include the corresponding headers in your MCP client requests +3. **Header Forwarding**: LiteLLM automatically forwards matching headers to the backend MCP server +4. **Authentication**: The backend MCP server receives both the configured auth headers and the custom headers + +### Use Cases + +- **Custom Authentication**: Forward custom API keys or tokens required by specific MCP servers +- **Request Context**: Pass user identification, session data, or request tracking headers +- **Third-party Integration**: Include headers required by external services that your MCP server integrates with +- **Multi-tenant Systems**: Forward tenant-specific headers for proper request routing + +### Security Considerations + +- Only headers listed in `extra_headers` are forwarded to maintain security +- Sensitive headers should be passed through environment variables when possible +- Consider using server-specific auth headers for better security isolation + +--- + ## Using your MCP with client side credentials Use this if you want to pass a client side authentication token to LiteLLM to then pass to your MCP to auth to your MCP. @@ -780,13 +985,6 @@ Use this if you want to pass a client side authentication token to LiteLLM to th You can specify MCP auth tokens using server-specific headers in the format `x-mcp-{server_alias}-{header_name}`. This allows you to use different authentication for different MCP servers. -**Format:** `x-mcp-{server_alias}-{header_name}: value` - -**Examples:** -- `x-mcp-github-authorization: Bearer ghp_xxxxxxxxx` - GitHub MCP server with Bearer token -- `x-mcp-zapier-x-api-key: sk-xxxxxxxxx` - Zapier MCP server with API key -- `x-mcp-deepwiki-authorization: Basic base64_encoded_creds` - DeepWiki MCP server with Basic auth - **Benefits:** - **Server-specific authentication**: Each MCP server can use different auth methods - **Better security**: No need to share the same auth token across all servers From 526156ed9d065721b57be5490a28df6ff7f2640b Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 27 Sep 2025 19:36:11 -0700 Subject: [PATCH 06/12] feat(mcp/): allows admin to prevent llm's from accidentally deleting github repo's even if user is allowed to do this --- .../mcp_server/mcp_server_manager.py | 26 ++ litellm/proxy/_new_secret_config.yaml | 3 +- .../types/mcp_server/mcp_server_manager.py | 2 + .../mcp_server/test_mcp_server_manager.py | 237 +++++++++++++++++- 4 files changed, 265 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index a3a5d93a3fd..6423a9ae153 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -213,6 +213,8 @@ class MCPServerManager: ), mcp_info=mcp_info, extra_headers=server_config.get("extra_headers", None), + allowed_tools=server_config.get("allowed_tools", None), + disallowed_tools=server_config.get("disallowed_tools", None), access_groups=server_config.get("access_groups", None), ) self.config_mcp_servers[server_id] = new_server @@ -277,6 +279,8 @@ class MCPServerManager: args=getattr(mcp_server, "args", None) or [], env=env_dict, access_groups=getattr(mcp_server, "mcp_access_groups", None), + allowed_tools=getattr(mcp_server, "allowed_tools", None), + disallowed_tools=getattr(mcp_server, "disallowed_tools", None), ) self.registry[mcp_server.server_id] = new_server verbose_logger.debug(f"Added MCP Server: {name_for_prefix}") @@ -569,6 +573,16 @@ class MCPServerManager: ) return prefixed_tools + def check_allowed_or_banned_tools(self, tool_name: str, server: MCPServer) -> bool: + """ + Check if the tool is allowed or banned for the given server + """ + if server.allowed_tools: + return tool_name in server.allowed_tools + if server.disallowed_tools: + return tool_name not in server.disallowed_tools + return True + async def pre_call_tool_check( self, name: str, @@ -576,7 +590,18 @@ class MCPServerManager: server_name_from_prefix: str, user_api_key_auth: Optional[UserAPIKeyAuth], proxy_logging_obj: ProxyLogging, + server: MCPServer, ): + + ## check if the tool is allowed or banned for the given server + if not self.check_allowed_or_banned_tools(name, server): + raise HTTPException( + status_code=403, + detail={ + "error": f"Tool {name} is not allowed for server {server.name}. Contact proxy admin to allow this tool." + }, + ) + pre_hook_kwargs = { "name": name, "arguments": arguments, @@ -700,6 +725,7 @@ class MCPServerManager: server_name_from_prefix=server_name_from_prefix, user_api_key_auth=user_api_key_auth, proxy_logging_obj=proxy_logging_obj, + server=mcp_server, ) # Get server-specific auth header if available diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index e7494bd0dab..804cf2cf2cf 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -25,7 +25,6 @@ mcp_servers: client_id: os.environ/GITHUB_OAUTH_CLIENT_ID client_secret: os.environ/GITHUB_OAUTH_CLIENT_SECRET scopes: ["public_repo", "user:email"] - extra_headers: ["custom_key"] - # allowed_tools: ["list_tools"] + allowed_tools: ["list_tools"] # disallowed_tools: ["repo_delete"] diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 4327bd5afe7..3e0c2b20e39 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -23,6 +23,8 @@ class MCPServer(BaseModel): extra_headers: Optional[List[str]] = ( None # allow admin to specify which headers to forward to the MCP server ) + allowed_tools: Optional[List[str]] = None + disallowed_tools: Optional[List[str]] = None # OAuth-specific fields client_id: Optional[str] = None client_secret: Optional[str] = None 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 3237de37636..126a4ebb896 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 @@ -1,8 +1,9 @@ import sys from datetime import datetime -from unittest.mock import MagicMock, AsyncMock +from unittest.mock import AsyncMock, MagicMock import pytest +from fastapi import HTTPException # Add the parent directory to the path so we can import litellm sys.path.insert(0, "../../../../../") @@ -420,6 +421,240 @@ class TestMCPServerManager: assert result["status"] == "healthy" assert result["tools_count"] == 1 + @pytest.mark.asyncio + async def test_pre_call_tool_check_allowed_tools_list_allows_tool(self): + """Test pre_call_tool_check allows tool when it's in allowed_tools list""" + manager = MCPServerManager() + + # Create server with allowed_tools list + server = MCPServer( + server_id="test-server", + name="test-server", + transport=MCPTransport.stdio, + allowed_tools=["allowed_tool", "another_allowed_tool"], + disallowed_tools=None, + ) + + # Mock dependencies + user_api_key_auth = MagicMock() + proxy_logging_obj = MagicMock() + + # Mock the async methods that pre_call_tool_check calls + proxy_logging_obj._create_mcp_request_object_from_kwargs = MagicMock( + return_value={} + ) + proxy_logging_obj._convert_mcp_to_llm_format = MagicMock(return_value={}) + proxy_logging_obj.pre_call_hook = AsyncMock(return_value={}) + + # This should not raise an exception + await manager.pre_call_tool_check( + name="allowed_tool", + arguments={"param": "value"}, + server_name_from_prefix="test-server", + user_api_key_auth=user_api_key_auth, + proxy_logging_obj=proxy_logging_obj, + server=server, + ) + + @pytest.mark.asyncio + async def test_pre_call_tool_check_allowed_tools_list_blocks_tool(self): + """Test pre_call_tool_check blocks tool when it's not in allowed_tools list""" + manager = MCPServerManager() + + # Create server with allowed_tools list + server = MCPServer( + server_id="test-server", + name="test-server", + transport=MCPTransport.stdio, + allowed_tools=["allowed_tool", "another_allowed_tool"], + disallowed_tools=None, + ) + + # Mock dependencies + user_api_key_auth = MagicMock() + proxy_logging_obj = MagicMock() + + # This should raise an HTTPException + with pytest.raises(HTTPException) as exc_info: + await manager.pre_call_tool_check( + name="blocked_tool", + arguments={"param": "value"}, + server_name_from_prefix="test-server", + user_api_key_auth=user_api_key_auth, + proxy_logging_obj=proxy_logging_obj, + server=server, + ) + + assert exc_info.value.status_code == 403 + assert ( + "Tool blocked_tool is not allowed for server test-server" + in exc_info.value.detail["error"] + ) + assert ( + "Contact proxy admin to allow this tool" in exc_info.value.detail["error"] + ) + + @pytest.mark.asyncio + async def test_pre_call_tool_check_disallowed_tools_list_allows_tool(self): + """Test pre_call_tool_check allows tool when it's not in disallowed_tools list""" + manager = MCPServerManager() + + # Create server with disallowed_tools list + server = MCPServer( + server_id="test-server", + name="test-server", + transport=MCPTransport.stdio, + allowed_tools=None, + disallowed_tools=["banned_tool", "another_banned_tool"], + ) + + # Mock dependencies + user_api_key_auth = MagicMock() + proxy_logging_obj = MagicMock() + + # Mock the async methods that pre_call_tool_check calls + proxy_logging_obj._create_mcp_request_object_from_kwargs = MagicMock( + return_value={} + ) + proxy_logging_obj._convert_mcp_to_llm_format = MagicMock(return_value={}) + proxy_logging_obj.pre_call_hook = AsyncMock(return_value={}) + + # This should not raise an exception + await manager.pre_call_tool_check( + name="allowed_tool", + arguments={"param": "value"}, + server_name_from_prefix="test-server", + user_api_key_auth=user_api_key_auth, + proxy_logging_obj=proxy_logging_obj, + server=server, + ) + + @pytest.mark.asyncio + async def test_pre_call_tool_check_disallowed_tools_list_blocks_tool(self): + """Test pre_call_tool_check blocks tool when it's in disallowed_tools list""" + manager = MCPServerManager() + + # Create server with disallowed_tools list + server = MCPServer( + server_id="test-server", + name="test-server", + transport=MCPTransport.stdio, + allowed_tools=None, + disallowed_tools=["banned_tool", "another_banned_tool"], + ) + + # Mock dependencies + user_api_key_auth = MagicMock() + proxy_logging_obj = MagicMock() + + # This should raise an HTTPException + with pytest.raises(HTTPException) as exc_info: + await manager.pre_call_tool_check( + name="banned_tool", + arguments={"param": "value"}, + server_name_from_prefix="test-server", + user_api_key_auth=user_api_key_auth, + proxy_logging_obj=proxy_logging_obj, + server=server, + ) + + assert exc_info.value.status_code == 403 + assert ( + "Tool banned_tool is not allowed for server test-server" + in exc_info.value.detail["error"] + ) + assert ( + "Contact proxy admin to allow this tool" in exc_info.value.detail["error"] + ) + + @pytest.mark.asyncio + async def test_pre_call_tool_check_no_restrictions_allows_any_tool(self): + """Test pre_call_tool_check allows any tool when no restrictions are set""" + manager = MCPServerManager() + + # Create server with no tool restrictions + server = MCPServer( + server_id="test-server", + name="test-server", + transport=MCPTransport.stdio, + allowed_tools=None, + disallowed_tools=None, + ) + + # Mock dependencies + user_api_key_auth = MagicMock() + proxy_logging_obj = MagicMock() + + # Mock the async methods that pre_call_tool_check calls + proxy_logging_obj._create_mcp_request_object_from_kwargs = MagicMock( + return_value={} + ) + proxy_logging_obj._convert_mcp_to_llm_format = MagicMock(return_value={}) + proxy_logging_obj.pre_call_hook = AsyncMock(return_value={}) + + # This should not raise an exception + await manager.pre_call_tool_check( + name="any_tool", + arguments={"param": "value"}, + server_name_from_prefix="test-server", + user_api_key_auth=user_api_key_auth, + proxy_logging_obj=proxy_logging_obj, + server=server, + ) + + @pytest.mark.asyncio + async def test_pre_call_tool_check_allowed_tools_takes_precedence(self): + """Test that allowed_tools list takes precedence over disallowed_tools list""" + manager = MCPServerManager() + + # Create server with both allowed_tools and disallowed_tools + # Note: The logic in check_allowed_or_banned_tools prioritizes allowed_tools + server = MCPServer( + server_id="test-server", + name="test-server", + transport=MCPTransport.stdio, + allowed_tools=["tool1", "tool2"], + disallowed_tools=["tool2", "tool3"], # tool2 is in both lists + ) + + # Mock dependencies + user_api_key_auth = MagicMock() + proxy_logging_obj = MagicMock() + + # Mock the async methods that pre_call_tool_check calls + proxy_logging_obj._create_mcp_request_object_from_kwargs = MagicMock( + return_value={} + ) + proxy_logging_obj._convert_mcp_to_llm_format = MagicMock(return_value={}) + proxy_logging_obj.pre_call_hook = AsyncMock(return_value={}) + + # tool2 should be allowed since it's in allowed_tools (takes precedence) + await manager.pre_call_tool_check( + name="tool2", + arguments={"param": "value"}, + server_name_from_prefix="test-server", + user_api_key_auth=user_api_key_auth, + proxy_logging_obj=proxy_logging_obj, + server=server, + ) + + # tool3 should be blocked since it's not in allowed_tools + with pytest.raises(HTTPException) as exc_info: + await manager.pre_call_tool_check( + name="tool3", + arguments={"param": "value"}, + server_name_from_prefix="test-server", + user_api_key_auth=user_api_key_auth, + proxy_logging_obj=proxy_logging_obj, + server=server, + ) + + assert exc_info.value.status_code == 403 + assert ( + "Tool tool3 is not allowed for server test-server" + in exc_info.value.detail["error"] + ) + if __name__ == "__main__": pytest.main([__file__]) From 340262c85f1c231d8662bb383aa4321d44f534f4 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 27 Sep 2025 20:57:12 -0700 Subject: [PATCH 07/12] feat(server.py): only show allowed mcp tools --- .../proxy/_experimental/mcp_server/server.py | 20 +- tests/mcp_tests/test_mcp_server.py | 296 ++++++++++++++++++ 2 files changed, 315 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index a3e51103ed5..94f81ccef8d 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -361,6 +361,24 @@ if MCP_AVAILABLE: return allowed_mcp_servers + def filter_tools_by_allowed_tools( + tools: List[MCPTool], + mcp_server: MCPServer, + ) -> List[MCPTool]: + """ + Filter tools by allowed tools + """ + tools_to_return = tools + if mcp_server.allowed_tools: + tools_to_return = [ + tool for tool in tools if tool.name in mcp_server.allowed_tools + ] + if mcp_server.disallowed_tools: + tools_to_return = [ + tool for tool in tools if tool.name not in mcp_server.disallowed_tools + ] + return tools_to_return + async def _get_tools_from_mcp_servers( user_api_key_auth: Optional[UserAPIKeyAuth], mcp_auth_header: Optional[str], @@ -431,7 +449,7 @@ if MCP_AVAILABLE: mcp_auth_header=server_auth_header, extra_headers=extra_headers, ) - all_tools.extend(tools) + all_tools.extend(filter_tools_by_allowed_tools(tools, server)) verbose_logger.debug( f"Successfully fetched {len(tools)} tools from server {server.name}" ) diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py index d0489bac580..759e642062e 100644 --- a/tests/mcp_tests/test_mcp_server.py +++ b/tests/mcp_tests/test_mcp_server.py @@ -1982,6 +1982,302 @@ async def test_list_tool_rest_api_all_servers_with_auth(): assert calls[1][0][1] == "Bearer slack_token" # server_auth_header +@pytest.mark.asyncio +async def test_filter_tools_by_allowed_tools_integration(): + """Test that filter_tools_by_allowed_tools works correctly via _get_tools_from_mcp_servers""" + from litellm.proxy._experimental.mcp_server.server import ( + _get_tools_from_mcp_servers, + ) + from litellm.proxy._types import UserAPIKeyAuth + from mcp.types import Tool as MCPTool + + # Create a mock user auth + mock_user_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user") + + # Create mock tools that will be returned by the server + mock_tools = [ + MCPTool( + name="allowed_tool_1", + description="This tool should be allowed", + inputSchema={"type": "object"}, + ), + MCPTool( + name="allowed_tool_2", + description="This tool should also be allowed", + inputSchema={"type": "object"}, + ), + MCPTool( + name="blocked_tool_1", + description="This tool should be blocked", + inputSchema={"type": "object"}, + ), + MCPTool( + name="blocked_tool_2", + description="This tool should also be blocked", + inputSchema={"type": "object"}, + ), + ] + + # Create a mock server with allowed_tools restriction + mock_server = MCPServer( + server_id="test-server-123", + name="test_server_with_allowed_tools", + url="https://test-server.com/mcp", + transport=MCPTransport.http, + allowed_tools=[ + "allowed_tool_1", + "allowed_tool_2", + ], # Only these tools should be returned + disallowed_tools=None, + ) + + # Create a mock MCPClient that returns all tools + mock_client = AsyncMock() + mock_client.list_tools = AsyncMock(return_value=mock_tools) + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) + + def mock_client_constructor(*args, **kwargs): + return mock_client + + # Mock the global MCP server manager + with patch( + "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager" + ) as mock_manager: + # Mock manager methods + mock_manager.get_allowed_mcp_servers = AsyncMock( + return_value=["test-server-123"] + ) + mock_manager.get_mcp_server_by_id = MagicMock(return_value=mock_server) + + # Mock the _get_tools_from_server method to return all tools + mock_manager._get_tools_from_server = AsyncMock(return_value=mock_tools) + + # Mock the MCPClient constructor + with patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPClient", + mock_client_constructor, + ): + # Call _get_tools_from_mcp_servers which should apply the filtering + filtered_tools = await _get_tools_from_mcp_servers( + user_api_key_auth=mock_user_auth, + mcp_auth_header="Bearer test_token", + mcp_servers=None, # Get from all servers + ) + + # Verify that only allowed tools are returned + assert ( + len(filtered_tools) == 2 + ), f"Expected 2 tools, got {len(filtered_tools)}" + + tool_names = [tool.name for tool in filtered_tools] + assert ( + "allowed_tool_1" in tool_names + ), "allowed_tool_1 should be in filtered results" + assert ( + "allowed_tool_2" in tool_names + ), "allowed_tool_2 should be in filtered results" + assert ( + "blocked_tool_1" not in tool_names + ), "blocked_tool_1 should be filtered out" + assert ( + "blocked_tool_2" not in tool_names + ), "blocked_tool_2 should be filtered out" + + # Verify the manager methods were called correctly + mock_manager.get_allowed_mcp_servers.assert_called_once_with(mock_user_auth) + mock_manager.get_mcp_server_by_id.assert_called_once_with("test-server-123") + mock_manager._get_tools_from_server.assert_called_once() + + +@pytest.mark.asyncio +async def test_filter_tools_by_disallowed_tools_integration(): + """Test that filter_tools_by_allowed_tools works correctly with disallowed_tools via _get_tools_from_mcp_servers""" + from litellm.proxy._experimental.mcp_server.server import ( + _get_tools_from_mcp_servers, + ) + from litellm.proxy._types import UserAPIKeyAuth + from mcp.types import Tool as MCPTool + + # Create a mock user auth + mock_user_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user") + + # Create mock tools that will be returned by the server + mock_tools = [ + MCPTool( + name="safe_tool_1", + description="This tool should be allowed", + inputSchema={"type": "object"}, + ), + MCPTool( + name="safe_tool_2", + description="This tool should also be allowed", + inputSchema={"type": "object"}, + ), + MCPTool( + name="dangerous_tool_1", + description="This tool should be blocked", + inputSchema={"type": "object"}, + ), + MCPTool( + name="dangerous_tool_2", + description="This tool should also be blocked", + inputSchema={"type": "object"}, + ), + ] + + # Create a mock server with disallowed_tools restriction + mock_server = MCPServer( + server_id="test-server-456", + name="test_server_with_disallowed_tools", + url="https://test-server.com/mcp", + transport=MCPTransport.http, + allowed_tools=None, + disallowed_tools=[ + "dangerous_tool_1", + "dangerous_tool_2", + ], # These tools should be filtered out + ) + + # Create a mock MCPClient that returns all tools + mock_client = AsyncMock() + mock_client.list_tools = AsyncMock(return_value=mock_tools) + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) + + def mock_client_constructor(*args, **kwargs): + return mock_client + + # Mock the global MCP server manager + with patch( + "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager" + ) as mock_manager: + # Mock manager methods + mock_manager.get_allowed_mcp_servers = AsyncMock( + return_value=["test-server-456"] + ) + mock_manager.get_mcp_server_by_id = MagicMock(return_value=mock_server) + + # Mock the _get_tools_from_server method to return all tools + mock_manager._get_tools_from_server = AsyncMock(return_value=mock_tools) + + # Mock the MCPClient constructor + with patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPClient", + mock_client_constructor, + ): + # Call _get_tools_from_mcp_servers which should apply the filtering + filtered_tools = await _get_tools_from_mcp_servers( + user_api_key_auth=mock_user_auth, + mcp_auth_header="Bearer test_token", + mcp_servers=None, # Get from all servers + ) + + # Verify that only safe tools are returned (dangerous tools filtered out) + assert ( + len(filtered_tools) == 2 + ), f"Expected 2 tools, got {len(filtered_tools)}" + + tool_names = [tool.name for tool in filtered_tools] + assert ( + "safe_tool_1" in tool_names + ), "safe_tool_1 should be in filtered results" + assert ( + "safe_tool_2" in tool_names + ), "safe_tool_2 should be in filtered results" + assert ( + "dangerous_tool_1" not in tool_names + ), "dangerous_tool_1 should be filtered out" + assert ( + "dangerous_tool_2" not in tool_names + ), "dangerous_tool_2 should be filtered out" + + # Verify the manager methods were called correctly + mock_manager.get_allowed_mcp_servers.assert_called_once_with(mock_user_auth) + mock_manager.get_mcp_server_by_id.assert_called_once_with("test-server-456") + mock_manager._get_tools_from_server.assert_called_once() + + +@pytest.mark.asyncio +async def test_filter_tools_no_restrictions_integration(): + """Test that filter_tools_by_allowed_tools returns all tools when no restrictions are set""" + from litellm.proxy._experimental.mcp_server.server import ( + _get_tools_from_mcp_servers, + ) + from litellm.proxy._types import UserAPIKeyAuth + from mcp.types import Tool as MCPTool + + # Create a mock user auth + mock_user_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user") + + # Create mock tools that will be returned by the server + mock_tools = [ + MCPTool( + name="tool_1", + description="Tool 1", + inputSchema={"type": "object"}, + ), + MCPTool( + name="tool_2", + description="Tool 2", + inputSchema={"type": "object"}, + ), + ] + + # Create a mock server with no tool restrictions + mock_server = MCPServer( + server_id="test-server-000", + name="test_server_no_restrictions", + url="https://test-server.com/mcp", + transport=MCPTransport.http, + allowed_tools=None, # No restrictions + disallowed_tools=None, # No restrictions + ) + + # Create a mock MCPClient that returns all tools + mock_client = AsyncMock() + mock_client.list_tools = AsyncMock(return_value=mock_tools) + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) + + def mock_client_constructor(*args, **kwargs): + return mock_client + + # Mock the global MCP server manager + with patch( + "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager" + ) as mock_manager: + # Mock manager methods + mock_manager.get_allowed_mcp_servers = AsyncMock( + return_value=["test-server-000"] + ) + mock_manager.get_mcp_server_by_id = MagicMock(return_value=mock_server) + + # Mock the _get_tools_from_server method to return all tools + mock_manager._get_tools_from_server = AsyncMock(return_value=mock_tools) + + # Mock the MCPClient constructor + with patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPClient", + mock_client_constructor, + ): + # Call _get_tools_from_mcp_servers which should apply the filtering + filtered_tools = await _get_tools_from_mcp_servers( + user_api_key_auth=mock_user_auth, + mcp_auth_header="Bearer test_token", + mcp_servers=None, # Get from all servers + ) + + # Should return all tools when no restrictions + assert ( + len(filtered_tools) == 2 + ), f"Expected 2 tools, got {len(filtered_tools)}" + + tool_names = [tool.name for tool in filtered_tools] + assert "tool_1" in tool_names, "tool_1 should be in filtered results" + assert "tool_2" in tool_names, "tool_2 should be in filtered results" + + @pytest.mark.asyncio async def test_mcp_access_group_permission_inheritance_integration(): """Integration test for MCP access group permission inheritance""" From e6810b5f87e99fd721d582a90395377aebf220c0 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sun, 28 Sep 2025 08:50:40 -0700 Subject: [PATCH 08/12] docs(mcp.md): add allow/disallowed tools to docs --- docs/my-website/docs/mcp.md | 59 +++++++++++++++++++++++++++++++++++++ 1 file changed, 59 insertions(+) diff --git a/docs/my-website/docs/mcp.md b/docs/my-website/docs/mcp.md index 0ee63b0e852..18c7c6b8508 100644 --- a/docs/my-website/docs/mcp.md +++ b/docs/my-website/docs/mcp.md @@ -199,6 +199,65 @@ litellm_settings: +## MCP Tool Filtering + +Control which tools are available from your MCP servers. You can either allow only specific tools or block dangerous ones. + + + + +Use `allowed_tools` to specify exactly which tools users can access. All other tools will be blocked. + +```yaml title="config.yaml" showLineNumbers +mcp_servers: + github_mcp: + url: "https://api.githubcopilot.com/mcp" + auth_type: oauth2 + authorization_url: https://github.com/login/oauth/authorize + token_url: https://github.com/login/oauth/access_token + client_id: os.environ/GITHUB_OAUTH_CLIENT_ID + client_secret: os.environ/GITHUB_OAUTH_CLIENT_SECRET + scopes: ["public_repo", "user:email"] + allowed_tools: ["list_tools"] + # only list_tools will be available +``` + +**Use this when:** +- You want strict control over which tools are available +- You're in a high-security environment +- You're testing a new MCP server with limited tools + + + + +Use `disallowed_tools` to block specific tools. All other tools will be available. + +```yaml title="config.yaml" showLineNumbers +mcp_servers: + github_mcp: + url: "https://api.githubcopilot.com/mcp" + auth_type: oauth2 + authorization_url: https://github.com/login/oauth/authorize + token_url: https://github.com/login/oauth/access_token + client_id: os.environ/GITHUB_OAUTH_CLIENT_ID + client_secret: os.environ/GITHUB_OAUTH_CLIENT_SECRET + scopes: ["public_repo", "user:email"] + disallowed_tools: ["repo_delete"] + # only repo_delete will be blocked +``` + +**Use this when:** +- Most tools are safe, but you want to block a few dangerous ones +- You want to prevent expensive API calls +- You're gradually adding restrictions to an existing server + + + + +### Important Notes + +- If you specify both `allowed_tools` and `disallowed_tools`, the allowed list takes priority +- Tool names are case-sensitive ## Using your MCP From 1256dd53d9bc2da350adfc614193471a4b8f8a65 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sun, 28 Sep 2025 08:53:11 -0700 Subject: [PATCH 09/12] test: update tests --- tests/mcp_tests/test_mcp_server.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py index 759e642062e..48237bd9444 100644 --- a/tests/mcp_tests/test_mcp_server.py +++ b/tests/mcp_tests/test_mcp_server.py @@ -1349,6 +1349,11 @@ def test_add_update_server_with_alias(): mock_mcp_server.command = None mock_mcp_server.args = [] mock_mcp_server.env = None + # OAuth fields - set explicitly to None to avoid MagicMock objects + mock_mcp_server.client_id = None + mock_mcp_server.client_secret = None + mock_mcp_server.authorization_url = None + mock_mcp_server.token_url = None # Add server to manager test_manager.add_update_server(mock_mcp_server) @@ -1411,6 +1416,11 @@ def test_add_update_server_fallback_to_server_id(): mock_mcp_server.command = None mock_mcp_server.args = [] mock_mcp_server.env = None + # OAuth fields - set explicitly to None to avoid MagicMock objects + mock_mcp_server.client_id = None + mock_mcp_server.client_secret = None + mock_mcp_server.authorization_url = None + mock_mcp_server.token_url = None # Add server to manager test_manager.add_update_server(mock_mcp_server) From 9403efa511ff2f6dcd28ef895a87ba8f8b8987d1 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sun, 28 Sep 2025 08:54:15 -0700 Subject: [PATCH 10/12] test: fix tests --- tests/mcp_tests/test_mcp_server.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py index 48237bd9444..2dfe39777fd 100644 --- a/tests/mcp_tests/test_mcp_server.py +++ b/tests/mcp_tests/test_mcp_server.py @@ -1385,6 +1385,11 @@ def test_add_update_server_without_alias(): mock_mcp_server.command = None mock_mcp_server.args = [] mock_mcp_server.env = None + # OAuth fields - set explicitly to None to avoid MagicMock objects + mock_mcp_server.client_id = None + mock_mcp_server.client_secret = None + mock_mcp_server.authorization_url = None + mock_mcp_server.token_url = None # Add server to manager test_manager.add_update_server(mock_mcp_server) From 53d0cbb1b7cb340ae1b549b43387ec9e1f40ac19 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sun, 28 Sep 2025 09:09:11 -0700 Subject: [PATCH 11/12] fix: update tests + logic for passing multiple headers --- .../mcp_server/auth/user_api_key_auth_mcp.py | 2 +- .../auth/test_user_api_key_auth_mcp.py | 48 ++++++++++++------- 2 files changed, 31 insertions(+), 19 deletions(-) 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 4b2c3385bbc..581da3cec2e 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 @@ -186,7 +186,7 @@ class MCPRequestHandler: if len(parts) == 2: server_alias, auth_header_name = parts - # Convert header name to proper case (e.g., "authorization" -> "Authorization") + # Convert common header names to proper case if auth_header_name == "authorization": auth_header_name = "Authorization" 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 7bdefe84c1d..bf051085ccb 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 @@ -286,7 +286,10 @@ class TestMCPRequestHandler: ], "test-api-key-123", None, - {"github": "Bearer github-token", "zapier_x_api": "zapier-api-key"}, + { + "github": {"Authorization": "Bearer github-token"}, + "zapier_x_api": {"key": "zapier-api-key"}, + }, ), # Test case 10: Both legacy and server-specific auth headers ( @@ -297,7 +300,7 @@ class TestMCPRequestHandler: ], "test-api-key-123", "legacy-token", - {"github": "Bearer github-token"}, + {"github": {"Authorization": "Bearer github-token"}}, ), # Test case 11: Server-specific auth headers with different header types ( @@ -308,7 +311,10 @@ class TestMCPRequestHandler: ], "test-api-key-123", None, - {"deepwiki": "Basic base64-encoded", "custom_x_custom": "custom-value"}, + { + "deepwiki": {"Authorization": "Basic base64-encoded"}, + "custom_x_custom": {"header": "custom-value"}, + }, ), # Test case 12: Case insensitive server-specific headers ( @@ -318,7 +324,7 @@ class TestMCPRequestHandler: ], "test-api-key-123", None, - {"github": "Bearer github-token"}, + {"github": {"Authorization": "Bearer github-token"}}, ), ], ) @@ -365,6 +371,7 @@ class TestMCPRequestHandler: mcp_servers, mcp_server_auth_headers, oauth2_headers, + raw_headers, ) = await MCPRequestHandler.process_mcp_request(scope) # Assert the results @@ -751,7 +758,7 @@ class TestMCPCustomHeaderName: } ) result = MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers) - assert result == {"github": "Bearer github-token"} + assert result == {"github": {"Authorization": "Bearer github-token"}} # Test case 3: Multiple server-specific headers headers = Headers( @@ -764,9 +771,9 @@ class TestMCPCustomHeaderName: ) result = MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers) expected = { - "github": "Bearer github-token", - "zapier_x_api": "zapier-api-key", - "deepwiki": "Basic base64-encoded", + "github": {"Authorization": "Bearer github-token"}, + "zapier_x_api": {"key": "zapier-api-key"}, + "deepwiki": {"Authorization": "Basic base64-encoded"}, } assert result == expected @@ -775,11 +782,14 @@ class TestMCPCustomHeaderName: { "x-litellm-api-key": "test-key", "X-MCP-GITHUB-AUTHORIZATION": "Bearer github-token", - "x-mcp-ZAPIER_x_api-key": "zapier-api-key", + "x-mcp-ZAPIER-x-api-key": "zapier-api-key", } ) result = MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers) - expected = {"github": "Bearer github-token", "zapier_x_api": "zapier-api-key"} + expected = { + "github": {"Authorization": "Bearer github-token"}, + "zapier": {"x-api-key": "zapier-api-key"}, + } assert result == expected # Test case 5: Invalid format headers (should be ignored) @@ -792,7 +802,7 @@ class TestMCPCustomHeaderName: } ) result = MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers) - assert result == {"github": "Bearer github-token"} + assert result == {"github": {"Authorization": "Bearer github-token"}} # Test case 6: Edge case - header with multiple hyphens in server alias headers = Headers( @@ -804,8 +814,8 @@ class TestMCPCustomHeaderName: ) result = MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers) expected = { - "github_mcp": "Bearer github-mcp-token", - "gh_mcp2": "Bearer gh-mcp2-token", + "github_mcp": {"Authorization": "Bearer github-mcp-token"}, + "gh_mcp2": {"Authorization": "Bearer gh-mcp2-token"}, } assert result == expected @@ -817,14 +827,14 @@ class TestMCPCustomHeaderName: } ) result = MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers) - assert result == {"github_mcp": "Bearer github-mcp-token"} + assert result == {"github_mcp": {"Authorization": "Bearer github-mcp-token"}} # Test case 8: Edge case - empty header value headers = Headers( {"x-litellm-api-key": "test-key", "x-mcp-github-authorization": ""} ) result = MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers) - assert result == {"github": ""} + assert result == {"github": {"Authorization": ""}} # Test case 9: Edge case - very long header value long_token = "Bearer " + "x" * 1000 @@ -832,7 +842,7 @@ class TestMCPCustomHeaderName: {"x-litellm-api-key": "test-key", "x-mcp-github-authorization": long_token} ) result = MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers) - assert result == {"github": long_token} + assert result == {"github": {"Authorization": long_token}} # Test case 10: Edge case - special characters in server alias headers = Headers( @@ -844,8 +854,8 @@ class TestMCPCustomHeaderName: ) result = MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers) expected = { - "github-123": "Bearer github-123-token", - "github_test": "Bearer github-test-token", + "github": {"123-authorization": "Bearer github-123-token"}, + "github_test": {"Authorization": "Bearer github-test-token"}, } assert result == expected @@ -890,6 +900,7 @@ class TestMCPAccessGroupsE2E: mcp_servers, mcp_server_auth_headers, oauth2_headers, + raw_headers, ) = await MCPRequestHandler.process_mcp_request(scope) # Assert the results @@ -940,6 +951,7 @@ class TestMCPAccessGroupsE2E: mcp_servers, mcp_server_auth_headers, oauth2_headers, + raw_headers, ) = await MCPRequestHandler.process_mcp_request(scope) # Assert the results From 433624706e803873b960e1944c35d5bfe846b91b Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sun, 28 Sep 2025 09:18:02 -0700 Subject: [PATCH 12/12] fix: fix linting errors --- .../proxy/_experimental/mcp_server/auth/litellm_auth_handler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 56aeb77f422..081d83dd1c8 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, Union +from typing import Dict, List, Optional from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser