mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-10 16:22:17 +00:00
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:
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user