mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-15 23:31:29 +00:00
fix(mcp): forward extra_headers for OpenAPI MCP tools
OpenAPI-generated tools only applied static closure headers and BYOK Authorization via ContextVar. Copy MCPServer.extra_headers from the incoming MCP request into _request_extra_headers (set in server.py before local tool dispatch), merge in openapi_to_mcp_generator via a small helper. OAuth2 M2M: do not forward caller Authorization from raw_headers (same rule as _prepare_mcp_server_headers for managed MCP). Adds TestRequestExtraHeaders and clarifies mcp_server_manager registration comment. Fixes #26794 Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Milan
parent
a67b7a7e87
commit
2ef77b0591
4 changed files with 195 additions and 9 deletions
|
|
@ -497,7 +497,8 @@ class MCPServerManager:
|
|||
# Add any static headers from server config.
|
||||
#
|
||||
# Note: `extra_headers` on MCPServer is a List[str] of header names to forward
|
||||
# from the client request (not available in this OpenAPI tool generation step).
|
||||
# from each client MCP request; values are applied at call time via
|
||||
# `_request_extra_headers` in server.py (not baked in here).
|
||||
# `static_headers` is a dict of concrete headers to always send.
|
||||
headers = (
|
||||
merge_mcp_headers(
|
||||
|
|
|
|||
|
|
@ -55,6 +55,13 @@ _request_auth_header: contextvars.ContextVar[Optional[str]] = contextvars.Contex
|
|||
"_request_auth_header", default=None
|
||||
)
|
||||
|
||||
# Per-request extra headers forwarded from the client request.
|
||||
# Populated from MCPServer.extra_headers names matched against raw request
|
||||
# headers in server.py before dispatching to a local/OpenAPI tool handler.
|
||||
_request_extra_headers: contextvars.ContextVar[Optional[Dict[str, str]]] = (
|
||||
contextvars.ContextVar("_request_extra_headers", default=None)
|
||||
)
|
||||
|
||||
|
||||
def _sanitize_path_parameter_value(param_value: Any, param_name: str) -> str:
|
||||
"""Ensure path params cannot introduce directory traversal."""
|
||||
|
|
@ -297,6 +304,20 @@ def build_input_schema(operation: Dict[str, Any]) -> Dict[str, Any]:
|
|||
}
|
||||
|
||||
|
||||
def _merge_openapi_tool_request_headers(
|
||||
static_headers: Dict[str, str]
|
||||
) -> Dict[str, str]:
|
||||
"""Merge static closure headers with per-request ContextVar overrides."""
|
||||
effective_headers = dict(static_headers)
|
||||
request_extra = _request_extra_headers.get()
|
||||
if request_extra:
|
||||
effective_headers.update(request_extra)
|
||||
override_auth = _request_auth_header.get()
|
||||
if override_auth:
|
||||
effective_headers["Authorization"] = override_auth
|
||||
return effective_headers
|
||||
|
||||
|
||||
def create_tool_function(
|
||||
path: str,
|
||||
method: str,
|
||||
|
|
@ -334,14 +355,7 @@ def create_tool_function(
|
|||
The function safely handles parameter names that aren't valid Python identifiers
|
||||
by using **kwargs instead of named parameters.
|
||||
"""
|
||||
# Allow per-request auth override (e.g. BYOK credential set via ContextVar).
|
||||
# The ContextVar holds the full Authorization header value, including the
|
||||
# correct prefix (Bearer / ApiKey / Basic) formatted by the caller in
|
||||
# server.py based on the server's configured auth_type.
|
||||
effective_headers = dict(headers)
|
||||
override_auth = _request_auth_header.get()
|
||||
if override_auth:
|
||||
effective_headers["Authorization"] = override_auth
|
||||
effective_headers = _merge_openapi_tool_request_headers(headers)
|
||||
|
||||
# Build URL from base_url and path
|
||||
url = base_url + path
|
||||
|
|
|
|||
|
|
@ -158,6 +158,7 @@ if MCP_AVAILABLE:
|
|||
)
|
||||
from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import (
|
||||
_request_auth_header,
|
||||
_request_extra_headers,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.sse_transport import SseServerTransport
|
||||
from litellm.proxy._experimental.mcp_server.tool_registry import (
|
||||
|
|
@ -2195,11 +2196,42 @@ if MCP_AVAILABLE:
|
|||
auth_header_value = f"Basic {mcp_auth_header}"
|
||||
else:
|
||||
auth_header_value = f"Bearer {mcp_auth_header}"
|
||||
|
||||
# Forward named client headers to OpenAPI tool upstream requests.
|
||||
# MCPServer.extra_headers lists header names to copy from raw_headers.
|
||||
# OAuth2 M2M: never take Authorization from the caller (matches
|
||||
# _prepare_mcp_server_headers for managed MCP).
|
||||
forwarded_headers: Optional[Dict[str, str]] = None
|
||||
if mcp_server and mcp_server.extra_headers and raw_headers:
|
||||
normalized_raw = {
|
||||
str(k).lower(): v
|
||||
for k, v in raw_headers.items()
|
||||
if isinstance(k, str)
|
||||
}
|
||||
skip_caller_authorization = bool(
|
||||
getattr(mcp_server, "has_client_credentials", False)
|
||||
)
|
||||
for header_name in mcp_server.extra_headers:
|
||||
if not isinstance(header_name, str):
|
||||
continue
|
||||
if (
|
||||
skip_caller_authorization
|
||||
and header_name.lower() == "authorization"
|
||||
):
|
||||
continue
|
||||
value = normalized_raw.get(header_name.lower())
|
||||
if value is not None:
|
||||
if forwarded_headers is None:
|
||||
forwarded_headers = {}
|
||||
forwarded_headers[header_name] = value
|
||||
|
||||
_auth_token = _request_auth_header.set(auth_header_value)
|
||||
_extra_token = _request_extra_headers.set(forwarded_headers)
|
||||
try:
|
||||
local_content = await _handle_local_mcp_tool(name, arguments)
|
||||
finally:
|
||||
_request_auth_header.reset(_auth_token)
|
||||
_request_extra_headers.reset(_extra_token)
|
||||
response = CallToolResult(content=cast(Any, local_content), isError=False)
|
||||
|
||||
# Try managed MCP server tool (pass the full prefixed name)
|
||||
|
|
|
|||
|
|
@ -15,6 +15,8 @@ from unittest.mock import AsyncMock, patch
|
|||
import pytest
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import (
|
||||
_request_auth_header,
|
||||
_request_extra_headers,
|
||||
_resolve_param_list,
|
||||
_resolve_ref,
|
||||
build_input_schema,
|
||||
|
|
@ -1011,3 +1013,140 @@ class TestRegisterToolsFromOpenAPI:
|
|||
assert re.match(
|
||||
r"^[a-zA-Z0-9_-]+$", name
|
||||
), f"fallback tool name {name!r} not sanitized"
|
||||
|
||||
|
||||
class TestRequestExtraHeaders:
|
||||
"""Tests for _request_extra_headers ContextVar forwarding in tool_function."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_extra_headers_forwarded_to_upstream(self):
|
||||
"""Extra headers set via ContextVar are included in the upstream request."""
|
||||
operation = {}
|
||||
func = create_tool_function(
|
||||
path="/data",
|
||||
method="get",
|
||||
operation=operation,
|
||||
base_url="https://api.example.com",
|
||||
)
|
||||
|
||||
with patch(GET_ASYNC_CLIENT_TARGET) as mock_client:
|
||||
async_client = _create_mock_client("get", "ok")
|
||||
mock_client.return_value = async_client
|
||||
|
||||
token = _request_extra_headers.set({"X-TOKEN": "secret-value"})
|
||||
try:
|
||||
result = await func()
|
||||
finally:
|
||||
_request_extra_headers.reset(token)
|
||||
|
||||
assert result == "ok"
|
||||
call_args = async_client.get.call_args
|
||||
headers_sent = call_args[1]["headers"]
|
||||
assert headers_sent.get("X-TOKEN") == "secret-value"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_extra_headers_by_default(self):
|
||||
"""Without setting _request_extra_headers, no extra headers are injected."""
|
||||
operation = {}
|
||||
func = create_tool_function(
|
||||
path="/data",
|
||||
method="get",
|
||||
operation=operation,
|
||||
base_url="https://api.example.com",
|
||||
headers={"X-Static": "static-value"},
|
||||
)
|
||||
|
||||
with patch(GET_ASYNC_CLIENT_TARGET) as mock_client:
|
||||
async_client = _create_mock_client("get", "ok")
|
||||
mock_client.return_value = async_client
|
||||
|
||||
result = await func()
|
||||
|
||||
assert result == "ok"
|
||||
call_args = async_client.get.call_args
|
||||
headers_sent = call_args[1]["headers"]
|
||||
assert headers_sent == {"X-Static": "static-value"}
|
||||
assert "X-TOKEN" not in headers_sent
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_extra_headers_merged_with_static_headers(self):
|
||||
"""Request extra headers are merged on top of static (baked-in) headers."""
|
||||
operation = {}
|
||||
func = create_tool_function(
|
||||
path="/data",
|
||||
method="post",
|
||||
operation=operation,
|
||||
base_url="https://api.example.com",
|
||||
headers={"X-Static": "static-value"},
|
||||
)
|
||||
|
||||
with patch(GET_ASYNC_CLIENT_TARGET) as mock_client:
|
||||
async_client = _create_mock_client("post", "created")
|
||||
mock_client.return_value = async_client
|
||||
|
||||
token = _request_extra_headers.set({"X-TOKEN": "dynamic-value"})
|
||||
try:
|
||||
result = await func()
|
||||
finally:
|
||||
_request_extra_headers.reset(token)
|
||||
|
||||
assert result == "created"
|
||||
call_args = async_client.post.call_args
|
||||
headers_sent = call_args[1]["headers"]
|
||||
assert headers_sent.get("X-Static") == "static-value"
|
||||
assert headers_sent.get("X-TOKEN") == "dynamic-value"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auth_header_still_overrides_extra_headers(self):
|
||||
"""_request_auth_header takes precedence for Authorization over extra headers."""
|
||||
operation = {}
|
||||
func = create_tool_function(
|
||||
path="/secure",
|
||||
method="get",
|
||||
operation=operation,
|
||||
base_url="https://api.example.com",
|
||||
)
|
||||
|
||||
with patch(GET_ASYNC_CLIENT_TARGET) as mock_client:
|
||||
async_client = _create_mock_client("get", "secure-data")
|
||||
mock_client.return_value = async_client
|
||||
|
||||
extra_token = _request_extra_headers.set(
|
||||
{"Authorization": "Bearer extra", "X-TOKEN": "token-value"}
|
||||
)
|
||||
auth_token = _request_auth_header.set("Bearer byok-credential")
|
||||
try:
|
||||
result = await func()
|
||||
finally:
|
||||
_request_auth_header.reset(auth_token)
|
||||
_request_extra_headers.reset(extra_token)
|
||||
|
||||
assert result == "secure-data"
|
||||
call_args = async_client.get.call_args
|
||||
headers_sent = call_args[1]["headers"]
|
||||
assert headers_sent.get("Authorization") == "Bearer byok-credential"
|
||||
assert headers_sent.get("X-TOKEN") == "token-value"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_extra_headers_not_leaked_between_calls(self):
|
||||
"""After resetting the ContextVar, subsequent calls do not see the headers."""
|
||||
operation = {}
|
||||
func = create_tool_function(
|
||||
path="/data",
|
||||
method="get",
|
||||
operation=operation,
|
||||
base_url="https://api.example.com",
|
||||
)
|
||||
|
||||
with patch(GET_ASYNC_CLIENT_TARGET) as mock_client:
|
||||
async_client = _create_mock_client("get", "ok")
|
||||
mock_client.return_value = async_client
|
||||
|
||||
token = _request_extra_headers.set({"X-TOKEN": "first-call"})
|
||||
_request_extra_headers.reset(token)
|
||||
|
||||
await func()
|
||||
|
||||
call_args = async_client.get.call_args
|
||||
headers_sent = call_args[1]["headers"]
|
||||
assert "X-TOKEN" not in headers_sent
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue