From 1c71d4cdbbbc87575ee97fe17429d4c85c006ae2 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 27 Sep 2025 17:30:46 -0700 Subject: [PATCH] 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 2f4c6c2d8d..56aeb77f42 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 3417dad2e4..4b2c3385bb 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 891adf8328..c971abe430 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 de1c3e67fc..62ad1cc910 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 0dec1b23c6..dbd849646d 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