mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-27 01:22:18 +00:00
fix(mcp): inject user-field headers for OpenAPI/local tool dispatch
The OpenAPI/local-tool dispatch path validated required user-field values via _enforce_user_fields but never injected them into the upstream request. The managed MCP path already did this in _call_regular_mcp_tool, so users configuring a user-fields-enabled server backed by an OpenAPI spec would see the dashboard turn green but still hit upstream auth failures. Add a _request_user_field_headers ContextVar mirroring the existing _request_auth_header/_request_extra_headers pattern, populate it in the local-tool branch of execute_mcp_tool, and merge it after static_headers in _merge_openapi_tool_request_headers so admin-declared per-user values win over operator-configured static headers (matching managed-MCP precedence). BYOK Authorization override still wins over the user-field Authorization to mirror the managed path's separate mcp_auth_header parameter.
This commit is contained in:
parent
5437285451
commit
0023194ccd
4 changed files with 210 additions and 8 deletions
|
|
@ -62,6 +62,14 @@ _request_extra_headers: contextvars.ContextVar[Optional[Dict[str, str]]] = (
|
|||
contextvars.ContextVar("_request_extra_headers", default=None)
|
||||
)
|
||||
|
||||
# Per-request user-field headers resolved from MCPServer.user_fields and the
|
||||
# calling user's stored values. Set this ContextVar before calling a local
|
||||
# tool handler so admin-declared per-user values reach the upstream API even
|
||||
# when the tool dispatch goes through the OpenAPI/local registry path.
|
||||
_request_user_field_headers: contextvars.ContextVar[Optional[Dict[str, str]]] = (
|
||||
contextvars.ContextVar("_request_user_field_headers", default=None)
|
||||
)
|
||||
|
||||
|
||||
def _sanitize_path_parameter_value(param_value: Any, param_name: str) -> str:
|
||||
"""Ensure path params cannot introduce directory traversal."""
|
||||
|
|
@ -311,23 +319,26 @@ def _merge_openapi_tool_request_headers(
|
|||
|
||||
Precedence (highest to lowest):
|
||||
1. ``_request_auth_header`` — BYOK override of ``Authorization``
|
||||
2. ``static_headers`` — operator-configured headers baked into the
|
||||
2. ``_request_user_field_headers`` — admin-declared per-user values
|
||||
resolved from ``MCPServer.user_fields`` and the calling user's
|
||||
stored values
|
||||
3. ``static_headers`` — operator-configured headers baked into the
|
||||
tool closure at registration time
|
||||
3. ``_request_extra_headers`` — per-request headers forwarded from
|
||||
4. ``_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.
|
||||
User-field headers win over ``static_headers`` because the admin
|
||||
explicitly declared a field requiring a per-user value for that header
|
||||
name. This matches the managed MCP dispatch path
|
||||
(:meth:`MCPServerManager._call_regular_mcp_tool`), where user-field
|
||||
headers are merged after ``static_headers``.
|
||||
|
||||
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 {}
|
||||
user_field = _request_user_field_headers.get() or {}
|
||||
|
||||
static_lower_names = {k.lower() for k in static}
|
||||
effective_headers: Dict[str, str] = {
|
||||
|
|
@ -335,6 +346,14 @@ def _merge_openapi_tool_request_headers(
|
|||
}
|
||||
effective_headers.update(static)
|
||||
|
||||
if user_field:
|
||||
existing_lower_names = {k.lower(): k for k in effective_headers}
|
||||
for header_name, value in user_field.items():
|
||||
collision = existing_lower_names.get(header_name.lower())
|
||||
if collision is not None and collision != header_name:
|
||||
del effective_headers[collision]
|
||||
effective_headers[header_name] = value
|
||||
|
||||
override_auth = _request_auth_header.get()
|
||||
if override_auth:
|
||||
for existing in [k for k in effective_headers if k.lower() == "authorization"]:
|
||||
|
|
|
|||
|
|
@ -231,6 +231,7 @@ if MCP_AVAILABLE:
|
|||
from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import (
|
||||
_request_auth_header,
|
||||
_request_extra_headers,
|
||||
_request_user_field_headers,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.sse_transport import SseServerTransport
|
||||
from litellm.proxy._experimental.mcp_server.tool_registry import (
|
||||
|
|
@ -2363,13 +2364,35 @@ if MCP_AVAILABLE:
|
|||
forwarded_headers = {}
|
||||
forwarded_headers[header_name] = value
|
||||
|
||||
# User-fields headers: enforce_user_fields above guarantees the
|
||||
# required values are present; resolve them once and inject so
|
||||
# OpenAPI/local tools see the same user-provided values that the
|
||||
# managed MCP dispatch path injects.
|
||||
user_field_headers: Optional[Dict[str, str]] = None
|
||||
if server_has_user_fields(mcp_server):
|
||||
from litellm.proxy._experimental.mcp_server.user_fields import (
|
||||
resolve_user_field_headers,
|
||||
)
|
||||
|
||||
_, stored_field_values = await _get_user_field_values_cached(
|
||||
mcp_server, user_api_key_auth
|
||||
)
|
||||
if stored_field_values:
|
||||
resolved = resolve_user_field_headers(
|
||||
mcp_server, stored_field_values
|
||||
)
|
||||
if resolved:
|
||||
user_field_headers = resolved
|
||||
|
||||
_auth_token = _request_auth_header.set(auth_header_value)
|
||||
_extra_token = _request_extra_headers.set(forwarded_headers)
|
||||
_user_field_token = _request_user_field_headers.set(user_field_headers)
|
||||
try:
|
||||
local_content = await _handle_local_mcp_tool(name, arguments)
|
||||
finally:
|
||||
_request_auth_header.reset(_auth_token)
|
||||
_request_extra_headers.reset(_extra_token)
|
||||
_request_user_field_headers.reset(_user_field_token)
|
||||
response = CallToolResult(content=cast(Any, local_content), isError=False)
|
||||
|
||||
# Try managed MCP server tool (pass the full prefixed name)
|
||||
|
|
|
|||
|
|
@ -493,6 +493,69 @@ async def test_build_stdio_env_returns_user_env_when_no_static_env():
|
|||
assert merged == {"X": "1"}
|
||||
|
||||
|
||||
def test_openapi_tool_merge_injects_user_field_headers_over_static():
|
||||
"""User-field headers override static_headers in the OpenAPI/local-tool
|
||||
dispatch path, matching managed-MCP precedence.
|
||||
|
||||
Regression: stored user-field values were validated at enforcement time
|
||||
but silently dropped for OpenAPI tools, so upstream calls failed even
|
||||
after users provided required values.
|
||||
"""
|
||||
from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import (
|
||||
_merge_openapi_tool_request_headers,
|
||||
_request_auth_header,
|
||||
_request_extra_headers,
|
||||
_request_user_field_headers,
|
||||
)
|
||||
|
||||
extra_token = _request_extra_headers.set({"X-Trace": "abc"})
|
||||
user_token = _request_user_field_headers.set(
|
||||
{"Authorization": "Bearer user-tok", "X-Workspace": "ws1"}
|
||||
)
|
||||
auth_token = _request_auth_header.set(None)
|
||||
try:
|
||||
merged = _merge_openapi_tool_request_headers(
|
||||
{"Authorization": "Bearer static-default", "X-Static": "keep"}
|
||||
)
|
||||
finally:
|
||||
_request_auth_header.reset(auth_token)
|
||||
_request_user_field_headers.reset(user_token)
|
||||
_request_extra_headers.reset(extra_token)
|
||||
|
||||
# User-field Authorization wins over static.
|
||||
assert merged["Authorization"] == "Bearer user-tok"
|
||||
# Other user-field headers are added.
|
||||
assert merged["X-Workspace"] == "ws1"
|
||||
# Static headers without a user-field override survive.
|
||||
assert merged["X-Static"] == "keep"
|
||||
# Forwarded extras still come through when not colliding with static.
|
||||
assert merged["X-Trace"] == "abc"
|
||||
|
||||
|
||||
def test_openapi_tool_merge_byok_auth_overrides_user_field_authorization():
|
||||
"""BYOK `_request_auth_header` is the final word on Authorization,
|
||||
even when a user_field also targets Authorization. Matches the managed
|
||||
path's `mcp_auth_header` parameter precedence."""
|
||||
from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import (
|
||||
_merge_openapi_tool_request_headers,
|
||||
_request_auth_header,
|
||||
_request_extra_headers,
|
||||
_request_user_field_headers,
|
||||
)
|
||||
|
||||
extra_token = _request_extra_headers.set(None)
|
||||
user_token = _request_user_field_headers.set({"Authorization": "Bearer user-tok"})
|
||||
auth_token = _request_auth_header.set("Bearer byok-tok")
|
||||
try:
|
||||
merged = _merge_openapi_tool_request_headers({})
|
||||
finally:
|
||||
_request_auth_header.reset(auth_token)
|
||||
_request_user_field_headers.reset(user_token)
|
||||
_request_extra_headers.reset(extra_token)
|
||||
|
||||
assert merged["Authorization"] == "Bearer byok-tok"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# HTTP endpoints — POST / GET / DELETE /v1/mcp/server/{id}/user-field-values
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ MCP server tools.
|
|||
"""
|
||||
|
||||
from datetime import datetime, timezone
|
||||
from typing import Dict, Optional
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
|
@ -159,6 +160,102 @@ async def test_openapi_local_tool_blocked_when_pre_call_check_raises():
|
|||
handle_local.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_openapi_local_tool_injects_user_field_headers():
|
||||
"""OpenAPI/local-tool dispatch must inject the calling user's stored
|
||||
user-field values via `_request_user_field_headers` so upstream calls
|
||||
see them. Pre-fix this path validated the values at enforcement time
|
||||
but never set the ContextVar, so the upstream request went out without
|
||||
them and the user saw "missing credential" errors after configuring
|
||||
the dashboard.
|
||||
"""
|
||||
from litellm.proxy._experimental.mcp_server import server as mcp_module
|
||||
from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import (
|
||||
_request_user_field_headers,
|
||||
)
|
||||
|
||||
user = UserAPIKeyAuth(
|
||||
api_key="sk-user",
|
||||
user_id="alice",
|
||||
user_role=LitellmUserRoles.INTERNAL_USER.value,
|
||||
)
|
||||
|
||||
fake_server = MagicMock()
|
||||
fake_server.name = "openapi-petstore"
|
||||
fake_server.is_byok = False
|
||||
fake_server.auth_type = None
|
||||
fake_server.mcp_info = None
|
||||
fake_server.server_id = "srv-1"
|
||||
fake_server.server_name = "openapi-petstore"
|
||||
fake_server.user_fields = [
|
||||
{
|
||||
"field_key": "API_KEY",
|
||||
"display_name": "API Key",
|
||||
"header_name": "X-Api-Key",
|
||||
"required": True,
|
||||
}
|
||||
]
|
||||
fake_server.extra_headers = None
|
||||
fake_server.has_client_credentials = False
|
||||
|
||||
fake_tool = MagicMock()
|
||||
fake_tool.name = "list_pets"
|
||||
|
||||
captured: Dict[str, Optional[Dict[str, str]]] = {"headers": None}
|
||||
|
||||
async def fake_handle_local(name, arguments):
|
||||
captured["headers"] = _request_user_field_headers.get()
|
||||
return []
|
||||
|
||||
pre_call = AsyncMock(return_value={})
|
||||
|
||||
# Simulate "user has saved their values" by seeding the in-memory cache
|
||||
# the enforcement / dispatch path consult.
|
||||
mcp_module._user_fields_cache.clear()
|
||||
mcp_module._user_fields_cache[("alice", "srv-1")] = (
|
||||
{"API_KEY": "secret-token"},
|
||||
1e18,
|
||||
)
|
||||
|
||||
try:
|
||||
with (
|
||||
patch.object(
|
||||
mcp_module.global_mcp_server_manager,
|
||||
"_get_mcp_server_from_tool_name",
|
||||
return_value=fake_server,
|
||||
),
|
||||
patch.object(
|
||||
mcp_module.global_mcp_server_manager,
|
||||
"pre_call_tool_check",
|
||||
new=pre_call,
|
||||
),
|
||||
patch.object(
|
||||
mcp_module.global_mcp_tool_registry,
|
||||
"get_tool",
|
||||
return_value=fake_tool,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.server._handle_local_mcp_tool",
|
||||
new=fake_handle_local,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.server.MCPRequestHandler.is_tool_allowed",
|
||||
return_value=True,
|
||||
),
|
||||
):
|
||||
await mcp_module.execute_mcp_tool(
|
||||
name="list_pets",
|
||||
arguments={},
|
||||
allowed_mcp_servers=[fake_server],
|
||||
start_time=datetime.now(timezone.utc),
|
||||
user_api_key_auth=user,
|
||||
)
|
||||
finally:
|
||||
mcp_module._user_fields_cache.clear()
|
||||
|
||||
assert captured["headers"] == {"X-Api-Key": "secret-token"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_openapi_local_tool_denied_when_server_not_resolvable():
|
||||
"""If the local-registry tool is found but no MCP server resolves
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue