diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 25b5211464..2c03cbdae3 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -1735,7 +1735,7 @@ class MCPServerManager: if extra_headers is None: extra_headers = {} for header in mcp_server.extra_headers: - if header in raw_headers: + if isinstance(header, str) and header in raw_headers: extra_headers[header] = raw_headers[header] if mcp_server.static_headers: diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 40b86edef3..950ea7063f 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -279,10 +279,27 @@ async def aresponses_api_with_mcp( user_api_key_auth = kwargs.get("litellm_metadata", {}).get( "user_api_key_auth" ) + + # Extract MCP auth headers from the request to pass to MCP server + secret_fields: Optional[Dict[str, Any]] = kwargs.get("secret_fields") + ( + mcp_auth_header, + mcp_server_auth_headers, + oauth2_headers, + raw_headers_from_request, + ) = ResponsesAPIRequestUtils.extract_mcp_headers_from_request( + secret_fields=secret_fields, + tools=tools, + ) + tool_results = await LiteLLM_Proxy_MCP_Handler._execute_tool_calls( tool_server_map=tool_server_map, tool_calls=tool_calls, user_api_key_auth=user_api_key_auth, + mcp_auth_header=mcp_auth_header, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers_from_request, ) if tool_results: diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index f8f5154e2d..a41d6f4f5a 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -399,7 +399,13 @@ class LiteLLM_Proxy_MCP_Handler: @staticmethod async def _execute_tool_calls( - tool_server_map: dict[str, str], tool_calls: List[Any], user_api_key_auth: Any + tool_server_map: dict[str, str], + tool_calls: List[Any], + user_api_key_auth: Any, + 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, ) -> List[Dict[str, Any]]: """Execute tool calls and return results.""" from fastapi import HTTPException @@ -437,6 +443,10 @@ class LiteLLM_Proxy_MCP_Handler: name=tool_name, arguments=parsed_arguments, user_api_key_auth=user_api_key_auth, + 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, ) diff --git a/litellm/responses/mcp/mcp_streaming_iterator.py b/litellm/responses/mcp/mcp_streaming_iterator.py index ea31f0f7f1..c00c2a2f3b 100644 --- a/litellm/responses/mcp/mcp_streaming_iterator.py +++ b/litellm/responses/mcp/mcp_streaming_iterator.py @@ -299,8 +299,61 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): "custom_llm_provider", None ) + self._extract_mcp_headers_from_params() + # Mark as async iterator self.is_async = True + + def _extract_mcp_headers_from_params(self) -> None: + """Extract MCP headers from original request params to pass to tool calls""" + from typing import Dict, Optional + from starlette.datastructures import Headers + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( + MCPRequestHandler, + ) + + # Extract headers from secret_fields in original_request_params + raw_headers_from_request: Optional[Dict[str, str]] = None + secret_fields = self.original_request_params.get("secret_fields") + if secret_fields and isinstance(secret_fields, dict): + raw_headers_from_request = secret_fields.get("raw_headers") + + # Extract MCP-specific headers + self.mcp_auth_header: Optional[str] = None + self.mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None + self.oauth2_headers: Optional[Dict[str, str]] = None + self.raw_headers: Optional[Dict[str, str]] = raw_headers_from_request + + if raw_headers_from_request: + headers_obj = Headers(raw_headers_from_request) + self.mcp_auth_header = MCPRequestHandler._get_mcp_auth_header_from_headers(headers_obj) + self.mcp_server_auth_headers = MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers_obj) + self.oauth2_headers = MCPRequestHandler._get_oauth2_headers_from_headers(headers_obj) + + # Also check if headers are provided in tools array (from request body) + tools = self.original_request_params.get("tools") + if tools: + for tool in tools: + if isinstance(tool, dict) and tool.get("type") == "mcp": + tool_headers = tool.get("headers", {}) + if tool_headers and isinstance(tool_headers, dict): + # Merge tool headers into mcp_server_auth_headers + headers_obj_from_tool = Headers(tool_headers) + tool_mcp_server_auth_headers = MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers_obj_from_tool) + + if tool_mcp_server_auth_headers: + if self.mcp_server_auth_headers is None: + self.mcp_server_auth_headers = {} + # Merge the headers from tool into existing headers + for server_alias, headers_dict in tool_mcp_server_auth_headers.items(): + if server_alias not in self.mcp_server_auth_headers: + self.mcp_server_auth_headers[server_alias] = {} + self.mcp_server_auth_headers[server_alias].update(headers_dict) + + # Also merge raw headers + if self.raw_headers is None: + self.raw_headers = {} + self.raw_headers.update(tool_headers) def _should_auto_execute_tools(self) -> bool: """Check if tools should be auto-executed""" @@ -511,6 +564,10 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): tool_server_map=self.tool_server_map, tool_calls=tool_calls, user_api_key_auth=self.user_api_key_auth, + mcp_auth_header=self.mcp_auth_header, + mcp_server_auth_headers=self.mcp_server_auth_headers, + oauth2_headers=self.oauth2_headers, + raw_headers=self.raw_headers, ) # Create completion events and output_item.done events for tool execution diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index 7798a7573d..198182cf11 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -2,6 +2,7 @@ import base64 from typing import ( Any, Dict, + Iterable, List, Optional, Type, @@ -350,6 +351,65 @@ class ResponsesAPIRequestUtils: return text return text + @staticmethod + def extract_mcp_headers_from_request( + secret_fields: Optional[Dict[str, Any]], + tools: Optional[Iterable[Any]], + ) -> tuple[ + Optional[str], + Optional[Dict[str, Dict[str, str]]], + Optional[Dict[str, str]], + Optional[Dict[str, str]], + ]: + """ + Extract MCP auth headers from the request to pass to MCP server. + Headers from tools.headers in request body should be passed to MCP server. + """ + from starlette.datastructures import Headers + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( + MCPRequestHandler, + ) + + # Extract headers from secret_fields which contains the original request headers + raw_headers_from_request: Optional[Dict[str, str]] = None + if secret_fields and isinstance(secret_fields, dict): + raw_headers_from_request = secret_fields.get("raw_headers") + + # Extract MCP-specific headers using MCPRequestHandler methods + mcp_auth_header: Optional[str] = None + mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None + oauth2_headers: Optional[Dict[str, str]] = None + + if raw_headers_from_request: + headers_obj = Headers(raw_headers_from_request) + mcp_auth_header = MCPRequestHandler._get_mcp_auth_header_from_headers(headers_obj) + mcp_server_auth_headers = MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers_obj) + oauth2_headers = MCPRequestHandler._get_oauth2_headers_from_headers(headers_obj) + + if tools: + for tool in tools: + if isinstance(tool, dict) and tool.get("type") == "mcp": + tool_headers = tool.get("headers", {}) + if tool_headers and isinstance(tool_headers, dict): + # Merge tool headers into mcp_server_auth_headers + # Extract server-specific headers from tool.headers + headers_obj_from_tool = Headers(tool_headers) + tool_mcp_server_auth_headers = MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers_obj_from_tool) + if tool_mcp_server_auth_headers: + if mcp_server_auth_headers is None: + mcp_server_auth_headers = {} + # Merge the headers from tool into existing headers + for server_alias, headers_dict in tool_mcp_server_auth_headers.items(): + if server_alias not in mcp_server_auth_headers: + mcp_server_auth_headers[server_alias] = {} + mcp_server_auth_headers[server_alias].update(headers_dict) + # Also merge raw headers (non-prefixed headers from tool.headers) + if raw_headers_from_request is None: + raw_headers_from_request = {} + raw_headers_from_request.update(tool_headers) + + return mcp_auth_header, mcp_server_auth_headers, oauth2_headers, raw_headers_from_request + class ResponseAPILoggingUtils: @staticmethod