mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
Merge pull request #27768 from milan-berri/litellm_cherry-pick-27383-onto-1.84.0rc2
cherry-pick: OpenAPI MCP extra_headers (#27383) onto litellm_1.84.0rc2
This commit is contained in:
yuneng-jiang
GitHub
commit
4046cb1690
4 changed files with 276 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(
|
||||
|
|
|
|||
|
|
@ -31,6 +31,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."""
|
||||
|
|
@ -273,6 +280,46 @@ 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.
|
||||
|
||||
Precedence (highest to lowest):
|
||||
1. ``_request_auth_header`` — BYOK override of ``Authorization``
|
||||
2. ``static_headers`` — operator-configured headers baked into the
|
||||
tool closure at registration time
|
||||
3. ``_request_extra_headers`` — per-request headers forwarded from
|
||||
the MCP caller (allowlisted by ``MCPServer.extra_headers``)
|
||||
|
||||
This matches the existing MCP invariant in
|
||||
:func:`litellm.proxy._experimental.mcp_server.utils.merge_mcp_headers`
|
||||
and the managed MCP path, where ``static_headers`` always wins over
|
||||
caller-forwarded headers. Keeping the same precedence here prevents an
|
||||
authenticated caller from overriding an operator-configured value
|
||||
(e.g. a tenant id or upstream API key) by sending the same header name.
|
||||
|
||||
Header names are compared case-insensitively so different casing cannot
|
||||
bypass the precedence rules.
|
||||
"""
|
||||
request_extra = _request_extra_headers.get() or {}
|
||||
static = static_headers or {}
|
||||
|
||||
static_lower_names = {k.lower() for k in static}
|
||||
effective_headers: Dict[str, str] = {
|
||||
k: v for k, v in request_extra.items() if k.lower() not in static_lower_names
|
||||
}
|
||||
effective_headers.update(static)
|
||||
|
||||
override_auth = _request_auth_header.get()
|
||||
if override_auth:
|
||||
for existing in [k for k in effective_headers if k.lower() == "authorization"]:
|
||||
del effective_headers[existing]
|
||||
effective_headers["Authorization"] = override_auth
|
||||
|
||||
return effective_headers
|
||||
|
||||
|
||||
def create_tool_function(
|
||||
path: str,
|
||||
method: str,
|
||||
|
|
@ -310,14 +357,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,40 @@ 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(mcp_server.has_client_credentials)
|
||||
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,
|
||||
|
|
@ -868,3 +870,197 @@ class TestResolveOperationParams:
|
|||
assert "per_page" in names
|
||||
assert "sha" in names
|
||||
assert len(names) == 4 # no duplicates
|
||||
|
||||
|
||||
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):
|
||||
"""Forwarded headers are passed through alongside non-conflicting static 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_static_headers_win_over_forwarded_on_conflict(self):
|
||||
"""Static (operator) headers must override forwarded (caller) headers on name conflict."""
|
||||
operation = {}
|
||||
func = create_tool_function(
|
||||
path="/data",
|
||||
method="get",
|
||||
operation=operation,
|
||||
base_url="https://api.example.com",
|
||||
headers={"X-Tenant": "operator-tenant"},
|
||||
)
|
||||
|
||||
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-Tenant": "caller-spoofed"})
|
||||
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-Tenant") == "operator-tenant"
|
||||
assert "caller-spoofed" not in headers_sent.values()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_static_headers_win_case_insensitively(self):
|
||||
"""Forwarded header with different casing must not bypass the static-wins rule."""
|
||||
operation = {}
|
||||
func = create_tool_function(
|
||||
path="/data",
|
||||
method="get",
|
||||
operation=operation,
|
||||
base_url="https://api.example.com",
|
||||
headers={"X-Tenant": "operator-tenant"},
|
||||
)
|
||||
|
||||
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-tenant": "caller-spoofed"})
|
||||
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-Tenant") == "operator-tenant"
|
||||
assert "x-tenant" not in headers_sent
|
||||
assert "caller-spoofed" not in headers_sent.values()
|
||||
|
||||
@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