mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-26 01:12:21 +00:00
fix(mcp): enforce required user_fields in MCPServerManager.call_tool
The streamable-HTTP entrypoint (server.execute_mcp_tool) gated tool calls on a server's admin-declared user_fields, but the Responses API path (LiteLLM_Proxy_MCP_Handler) bypassed that gate by calling MCPServerManager.call_tool directly. With required fields unset, the manager just resolved an empty value map and dispatched anyway, so the upstream call went out without the configured headers/env. Move the enforcement into MCPServerManager.call_tool itself so every dispatch path raises the same friendly 401 (user_fields_missing + config_url) when required values are absent.
This commit is contained in:
parent
b7fd2ec550
commit
72184582f2
2 changed files with 108 additions and 0 deletions
|
|
@ -1371,6 +1371,44 @@ class MCPServerManager:
|
|||
_, values = await _get_user_field_values_cached(mcp_server, user_api_key_auth)
|
||||
return values or {}
|
||||
|
||||
async def _enforce_required_user_fields(
|
||||
self,
|
||||
mcp_server: MCPServer,
|
||||
user_api_key_auth: Optional["UserAPIKeyAuth"],
|
||||
) -> None:
|
||||
"""Raise the user_fields_missing 401 when required values are absent.
|
||||
|
||||
``server.execute_mcp_tool`` runs the same gate for the streamable-HTTP
|
||||
entrypoint, but ``call_tool`` is also called directly by the
|
||||
Responses API path (``LiteLLM_Proxy_MCP_Handler``). Running the check
|
||||
here covers every dispatch path through the manager so a caller
|
||||
cannot skip enforcement by going around ``execute_mcp_tool``.
|
||||
"""
|
||||
from litellm.proxy._experimental.mcp_server.user_fields import (
|
||||
build_user_fields_missing_error,
|
||||
compute_missing_user_fields,
|
||||
server_has_user_fields,
|
||||
)
|
||||
|
||||
if not server_has_user_fields(mcp_server):
|
||||
return
|
||||
|
||||
stored_values = await self._resolve_user_field_values(
|
||||
mcp_server, user_api_key_auth
|
||||
)
|
||||
missing = compute_missing_user_fields(mcp_server, stored_values or None)
|
||||
if not missing:
|
||||
return
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.oauth_utils import (
|
||||
_resolve_proxy_base_url_env,
|
||||
)
|
||||
|
||||
detail = build_user_fields_missing_error(
|
||||
mcp_server, missing, _resolve_proxy_base_url_env()
|
||||
)
|
||||
raise HTTPException(status_code=401, detail=detail)
|
||||
|
||||
async def _create_mcp_client(
|
||||
self,
|
||||
server: MCPServer,
|
||||
|
|
@ -2956,6 +2994,8 @@ class MCPServerManager:
|
|||
if mcp_server is None:
|
||||
raise ValueError(f"Tool {name} not found")
|
||||
|
||||
await self._enforce_required_user_fields(mcp_server, user_api_key_auth)
|
||||
|
||||
#########################################################
|
||||
# Pre MCP Tool Call Hook
|
||||
# Allow validation and modification of tool calls before execution
|
||||
|
|
|
|||
|
|
@ -500,6 +500,74 @@ async def test_enforce_user_fields_no_user_id_raises():
|
|||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_manager_call_tool_enforces_required_user_fields():
|
||||
"""``MCPServerManager.call_tool`` is invoked directly by the Responses
|
||||
API path, so it must raise the same friendly 401 as ``execute_mcp_tool``
|
||||
when a required user-field has no stored value — otherwise a caller
|
||||
can bypass enforcement and dispatch with the value silently dropped.
|
||||
"""
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
|
||||
MCPServerManager,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.server import _user_fields_cache
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
_user_fields_cache.clear()
|
||||
srv = _gmail_server()
|
||||
mgr = MCPServerManager()
|
||||
mgr.registry["s1"] = srv
|
||||
mgr.tool_name_to_mcp_server_name_mapping["gmail-prod-send"] = srv.name
|
||||
mgr.tool_name_to_mcp_server_name_mapping["send"] = srv.name
|
||||
user = UserAPIKeyAuth(api_key="hashed", user_id="responses-user")
|
||||
_user_fields_cache[("responses-user", "s1")] = (None, 1e18)
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await mgr.call_tool(
|
||||
server_name="gmail-prod",
|
||||
name="send",
|
||||
arguments={},
|
||||
user_api_key_auth=user,
|
||||
)
|
||||
assert exc_info.value.status_code == 401
|
||||
assert exc_info.value.detail["error"] == "user_fields_missing"
|
||||
assert "GMAIL_TOKEN" in [
|
||||
m["field_key"] for m in exc_info.value.detail["missing_fields"]
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_manager_call_tool_enforces_when_no_user_id():
|
||||
"""Anonymous callers cannot satisfy required user-fields, so the manager
|
||||
must still surface ``user_fields_missing`` rather than silently dispatch
|
||||
without the configured headers."""
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
|
||||
MCPServerManager,
|
||||
)
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
srv = _gmail_server()
|
||||
mgr = MCPServerManager()
|
||||
mgr.registry["s1"] = srv
|
||||
mgr.tool_name_to_mcp_server_name_mapping["gmail-prod-send"] = srv.name
|
||||
mgr.tool_name_to_mcp_server_name_mapping["send"] = srv.name
|
||||
user = UserAPIKeyAuth(api_key="hashed") # no user_id
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await mgr.call_tool(
|
||||
server_name="gmail-prod",
|
||||
name="send",
|
||||
arguments={},
|
||||
user_api_key_auth=user,
|
||||
)
|
||||
assert exc_info.value.status_code == 401
|
||||
assert exc_info.value.detail["error"] == "user_fields_missing"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_build_stdio_env_merges_user_field_env_over_static():
|
||||
"""Stored user_field env vars must take precedence over static server.env."""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue