From 0023194ccd0e1201c94767d14179566cbaefed5f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 19 May 2026 08:46:20 +0000 Subject: [PATCH] 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. --- .../mcp_server/openapi_to_mcp_generator.py | 35 +++++-- .../proxy/_experimental/mcp_server/server.py | 23 +++++ .../mcp_server/test_mcp_user_fields.py | 63 ++++++++++++ .../mcp_server/test_openapi_tool_auth.py | 97 +++++++++++++++++++ 4 files changed, 210 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py index 271517bb1e6..c04d70f4272 100644 --- a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py +++ b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py @@ -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"]: diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 588d02d3b2f..0e9a51014ce 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -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) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_user_fields.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_user_fields.py index 981b9d7eb63..9c4cfbcb833 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_user_fields.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_user_fields.py @@ -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 # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py index 3ad01e9c3ec..71112f75d61 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py @@ -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