Merge pull request #16877 from BerriAI/litellm_mcp_responses_api_header_passing

feat: Add header passing support for MCP tools in Responses API
This commit is contained in:
Sameer Kankute
2025-11-27 08:10:43 +05:30
committed by GitHub
5 changed files with 146 additions and 2 deletions
@@ -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:
+17
View File
@@ -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:
@@ -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,
)
@@ -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
+60
View File
@@ -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