scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-19 00:01:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 88

fix(mcp): hand pre-call hooks the metadata of the registered tool that actually runs

Browse source 
Building the guardrail's tool metadata from the local registry entry that
dispatch resolved, instead of re-deriving it from the tool name, keeps an
OpenAPI operation whose name starts with its own server prefix from being
reported with the shorter operation's description and schema. The registry
branch in get_listed_tool is gone with it, and the test doubles for the
local registry now carry a string description and dict schema like the real
entries do

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yucheng 2026-09-11 07:54:36 +00:00
parent 6c865d1413
commit df739bdd9e
5 changed files with 71 additions and 52 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
View file

@ -127,7 +127,6 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import (
from litellm.proxy._experimental.mcp_server.sampling_handler import (
MCP_SAMPLING_AVAILABLE,
)
from litellm.proxy._experimental.mcp_server.tool_registry import global_mcp_tool_registry
from litellm.proxy._experimental.mcp_server.utils import (
MCP_TOOL_PREFIX_SEPARATOR,
MCPMissingUserEnvVarsError,
@ -5140,14 +5139,6 @@ class MCPServerManager:
return prefixed_tools
def get_listed_tool(self, server: MCPServer, name: str) -> MCPTool | None:
if server.spec_path:
bare_name: Final = strip_known_server_prefix(name, server)
registered: Final = global_mcp_tool_registry.get_tool(
f"{get_server_prefix(server)}{MCP_TOOL_PREFIX_SEPARATOR}{bare_name}"
) or global_mcp_tool_registry.get_tool(bare_name)
if registered is None:
return None
return MCPTool(name=bare_name, description=registered.description, inputSchema=registered.input_schema)
listed: Final = self._listed_tools_by_server_id.get(server.server_id)
if not listed:
return None

11
litellm/proxy/_experimental/mcp_server/server.py
View file

@ -86,6 +86,7 @@ from litellm.proxy.litellm_pre_call_utils import (
)
from litellm.types.mcp import MCPAuth, MCPSpecVersion
from litellm.types.mcp_server.mcp_server_manager import MCPInfo, MCPServer
from litellm.types.mcp_server.tool_registry import MCPTool as RegisteredTool
from litellm.types.utils import CallTypes, StandardLoggingMCPToolCall
from litellm.utils import Rules, client, function_setup
@ -2777,6 +2778,9 @@ if MCP_AVAILABLE:
return managed_resource_templates
def _registered_tool_metadata(name: str, registered: RegisteredTool) -> MCPTool:
return MCPTool(name=name, description=registered.description, inputSchema=registered.input_schema)
def _resolve_display_name_to_original(
name: str,
allowed_mcp_servers: list[MCPServer],
@ -3119,7 +3123,7 @@ if MCP_AVAILABLE:
server=mcp_server,
raw_headers=raw_headers,
litellm_logging_obj=litellm_logging_obj,
tool=global_mcp_server_manager.get_listed_tool(mcp_server, original_tool_name),
tool=_registered_tool_metadata(original_tool_name, local_tool),
)
# `pre_call_tool_check` may return guardrail-modified
# arguments; honor them on the local path too.
@ -3185,7 +3189,8 @@ if MCP_AVAILABLE:
# not in the registry either, `_handle_local_mcp_tool` below reports
# 404 and nothing runs, so demanding a server here would turn every
# unknown tool name into a misleading 503.
if global_mcp_tool_registry.get_tool(original_tool_name) is not None:
registered_local_tool: Final = global_mcp_tool_registry.get_tool(original_tool_name)
if registered_local_tool is not None:
# `mcp_server` is None here because the tool name is not in the
# tool -> server mapping, but the name still carries a prefix
# that the server-level check above compared against the
@ -3226,7 +3231,7 @@ if MCP_AVAILABLE:
server=prefix_server,
raw_headers=raw_headers,
litellm_logging_obj=litellm_logging_obj,
tool=global_mcp_server_manager.get_listed_tool(prefix_server, original_tool_name),
tool=_registered_tool_metadata(original_tool_name, registered_local_tool),
)
if "arguments" in hook_result:
arguments = hook_result["arguments"] # pyright: ignore[reportAny] # hook returns untyped args

60
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
View file

@ -7109,6 +7109,8 @@ async def test_execute_mcp_tool_sets_model_in_model_call_details():
fake_tool = MagicMock()
fake_tool.name = "list_pets"
fake_tool.description = "test tool"
fake_tool.input_schema = {"type": "object"}
start_time = datetime.now(timezone.utc)
litellm_logging_obj, _ = function_setup(
@ -7175,7 +7177,7 @@ async def test_execute_mcp_tool_hands_openapi_registered_tool_metadata_to_pre_ca
)
schema = {"type": "object", "properties": {"limit": {"type": "integer"}}}
mcp_module.global_mcp_tool_registry.register_tool(
name="petstore-list_pets", description="List the pets", input_schema=schema, handler=lambda: None
name="petstore-list_pets", description="List the pets", input_schema=schema, handler=lambda limit: "ok"
)
manager = mcp_module.global_mcp_server_manager
manager._listed_tools_by_server_id.pop(petstore.server_id, None)
@ -7185,14 +7187,6 @@ async def test_execute_mcp_tool_hands_openapi_registered_tool_metadata_to_pre_ca
with (
patch.object(manager, "_get_mcp_server_from_tool_name", return_value=petstore),
patch.object(manager, "pre_call_tool_check", new=pre_call_tool_check),
patch(
"litellm.proxy._experimental.mcp_server.server._handle_local_mcp_tool",
new=AsyncMock(return_value=[]),
),
patch(
"litellm.proxy._experimental.mcp_server.server.MCPRequestHandler.is_tool_allowed",
return_value=True,
),
):
await mcp_module.execute_mcp_tool(
name="petstore-list_pets",
@ -7212,6 +7206,54 @@ async def test_execute_mcp_tool_hands_openapi_registered_tool_metadata_to_pre_ca
)
@pytest.mark.asyncio
async def test_execute_mcp_tool_hands_hooks_the_metadata_of_the_operation_it_runs_when_names_collide():
"""An OpenAPI operation whose name starts with its own server prefix must not be reported to the
pre-call hooks with the metadata of the shorter operation, since that is not the one that runs."""
from litellm.proxy._experimental.mcp_server import server as mcp_module
petstore = MCPServer(
server_id="petstore-id",
name="petstore",
server_name="petstore",
transport=MCPTransport.http,
url=None,
spec_path="https://example.com/petstore.yaml",
)
registry = mcp_module.global_mcp_tool_registry
registry.register_tool(name="petstore-get_pet", description="short", input_schema={}, handler=lambda: "short")
registry.register_tool(
name="petstore-petstore-get_pet",
description="long",
input_schema={"type": "object", "properties": {"petId": {"type": "integer"}}},
handler=lambda: "long",
)
manager = mcp_module.global_mcp_server_manager
pre_call_tool_check = AsyncMock(return_value={})
try:
with (
patch.object(manager, "_get_mcp_server_from_tool_name", return_value=petstore),
patch.object(manager, "pre_call_tool_check", new=pre_call_tool_check),
):
result = await mcp_module.execute_mcp_tool(
name="petstore-petstore-get_pet",
arguments={},
allowed_mcp_servers=[petstore],
start_time=datetime.now(),
user_api_key_auth=UserAPIKeyAuth(api_key="sk-user", user_id="alice"),
)
finally:
registry.unregister_tools_with_prefix("petstore-")
handed_tool = pre_call_tool_check.call_args.kwargs["tool"]
assert (handed_tool.description, handed_tool.inputSchema) == (
"long",
{"type": "object", "properties": {"petId": {"type": "integer"}}},
)
assert result.content[0].text == "long"
@pytest.mark.asyncio
async def test_execute_mcp_tool_rest_unresolved_prefixed_name_routes_to_requested_server():
"""A prefixed REST name that resolves to no tool must still dispatch to the server_id.

31
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py
View file

@ -6504,37 +6504,6 @@ class TestMCPServerManager:
assert by_prefixed_name is not None and by_prefixed_name.description == "v2"
assert manager.get_listed_tool(server, "missing") is None
def test_get_listed_tool_reads_openapi_registry_without_a_prior_listing(self):
"""OpenAPI tools live in the local registry from registration on, so their metadata must resolve
before any tools/list has run and must disappear with the registration."""
from litellm.proxy._experimental.mcp_server.tool_registry import global_mcp_tool_registry
server = MCPServer(
server_id="petstore-id",
name="petstore",
server_name="petstore",
transport=MCPTransport.http,
url=None,
spec_path="https://example.com/petstore.yaml",
)
schema = {"type": "object", "properties": {"petId": {"type": "integer"}}}
manager = MCPServerManager()
global_mcp_tool_registry.register_tool(
name="petstore-get_pet", description="Fetch a pet", input_schema=schema, handler=lambda: None
)
try:
for spelling in ("get_pet", "petstore-get_pet"):
tool = manager.get_listed_tool(server, spelling)
assert tool is not None and (tool.name, tool.description, tool.inputSchema) == (
"get_pet",
"Fetch a pet",
schema,
)
finally:
global_mcp_tool_registry.unregister_tools_with_prefix("petstore-")
assert manager.get_listed_tool(server, "petstore-get_pet") is None
@pytest.mark.asyncio
async def test_get_allowed_mcp_servers_with_user_api_key_auth(self):
"""

12
tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py
View file

@ -43,6 +43,8 @@ async def test_openapi_local_tool_runs_pre_call_tool_check():
fake_tool = MagicMock()
fake_tool.name = "list_pets"
fake_tool.description = "test tool"
fake_tool.input_schema = {"type": "object"}
pre_call = AsyncMock(return_value={})
handle_local = AsyncMock(return_value=[])
@ -124,6 +126,8 @@ async def test_openapi_local_tool_blocked_when_pre_call_check_raises():
fake_tool = MagicMock()
fake_tool.name = "delete_pet"
fake_tool.description = "test tool"
fake_tool.input_schema = {"type": "object"}
pre_call = AsyncMock(
side_effect=HTTPException(status_code=403, detail="not allowed")
@ -186,6 +190,8 @@ async def test_openapi_local_tool_denied_when_server_not_resolvable():
fake_tool = MagicMock()
fake_tool.name = "list_pets"
fake_tool.description = "test tool"
fake_tool.input_schema = {"type": "object"}
pre_call = AsyncMock(return_value={})
handle_local = AsyncMock(return_value=[])
@ -270,6 +276,8 @@ async def test_openapi_local_tool_injects_resolved_oauth_token():
fake_tool = MagicMock()
fake_tool.name = "get_values"
fake_tool.description = "test tool"
fake_tool.input_schema = {"type": "object"}
captured: dict = {}
async def handle_local(_name, _arguments):
@ -616,6 +624,8 @@ async def test_per_server_auth_header_reaches_both_openapi_dispatch_arms(dispatc
if dispatch_arm == "local_registry":
fake_tool = MagicMock()
fake_tool.name = "list_reports"
fake_tool.description = "test tool"
fake_tool.input_schema = {"type": "object"}
with (
patch.object(manager, "_get_mcp_server_from_tool_name", return_value=server),
patch.object(mcp_module.global_mcp_tool_registry, "get_tool", return_value=fake_tool),
@ -687,6 +697,8 @@ async def test_local_dispatch_reports_the_outcome_instead_of_success(failure: st
fake_tool = MagicMock()
fake_tool.name = "list_reports"
fake_tool.description = "test tool"
fake_tool.input_schema = {"type": "object"}
fake_tool.handler = raising_handler
server = MCPServer(
server_id="srv-openapi",