mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-10 03:28:53 +00:00
fix(mcp): hand never-listed tools/call hooks name and arguments only
The local-registry call path fell back to the registry entry with the admin description override when no tools/list had been recorded for the caller, so a pre_mcp_call guardrail scanned a description the caller was never served and blocked OpenAPI calls that passed before, and base's own selected-guardrail REST test failed on the two-text redaction. _registered_tool_metadata now returns the listed entry or None, so a tools/call with no prior listing sends name and arguments only as promised, and that REST test double goes back to its base shape
This commit is contained in:
Yucheng He
parent
4fc40056ae
commit
12947900b6
3 changed files with 42 additions and 82 deletions
|
|
@ -141,7 +141,6 @@ from litellm.types.mcp import (
|
|||
without_header,
|
||||
)
|
||||
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
|
||||
|
||||
|
|
@ -1609,17 +1608,10 @@ async def _list_mcp_resource_templates(
|
|||
return managed_resource_templates
|
||||
|
||||
|
||||
def _registered_tool_metadata(
|
||||
name: str, registered: RegisteredTool, server: MCPServer, caller: ListedToolsCaller
|
||||
) -> MCPTool:
|
||||
"""The tool as ``tools/list`` served it to this caller (pinned, overridden, guardrail-masked) when a
|
||||
listing was recorded, else the registry entry with the admin description override applied."""
|
||||
listed: Final = global_mcp_server_manager.get_listed_tool(server, name, caller)
|
||||
if listed is not None:
|
||||
return listed
|
||||
overrides: Final = server.tool_name_to_description
|
||||
description: Final = overrides.get(name, registered.description) if overrides else registered.description
|
||||
return MCPTool(name=name, description=description, input_schema=registered.input_schema)
|
||||
def _registered_tool_metadata(name: str, server: MCPServer, caller: ListedToolsCaller) -> MCPTool | None:
|
||||
"""The tool as ``tools/list`` served it to this caller (pinned, overridden, guardrail-masked), or None when
|
||||
no listing was recorded so the call hands the hooks name and arguments only."""
|
||||
return global_mcp_server_manager.get_listed_tool(server, name, caller)
|
||||
|
||||
|
||||
def _resolve_display_name_to_original(
|
||||
|
|
@ -2093,7 +2085,6 @@ async def _execute_mcp_tool(
|
|||
guardrail_context=guardrail_context,
|
||||
tool=_registered_tool_metadata(
|
||||
original_tool_name,
|
||||
local_tool,
|
||||
mcp_server,
|
||||
listed_tools_caller_for(
|
||||
mcp_server, user_api_key_auth, mcp_auth_header, mcp_server_auth_headers, raw_headers, oauth2_headers
|
||||
|
|
@ -2212,7 +2203,6 @@ async def _execute_mcp_tool(
|
|||
guardrail_context=guardrail_context,
|
||||
tool=_registered_tool_metadata(
|
||||
original_tool_name,
|
||||
registered_local_tool,
|
||||
prefix_server,
|
||||
listed_tools_caller_for(
|
||||
prefix_server,
|
||||
|
|
|
|||
|
|
@ -8070,10 +8070,13 @@ async def test_execute_mcp_tool_sets_model_in_model_call_details():
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_mcp_tool_hands_openapi_registered_tool_metadata_to_pre_call_hooks():
|
||||
"""OpenAPI-generated tools dispatch through the local registry, so the pre-call hooks must get the
|
||||
registered description and input schema on that path too, even when no tools/list ran first."""
|
||||
async def test_execute_mcp_tool_hands_openapi_hooks_the_listed_entry_and_nothing_before_a_listing():
|
||||
"""A local-registry tools/call with no prior tools/list hands the pre-call hooks name and arguments
|
||||
only, as before this metadata existed, so a pre_mcp_call policy never scans a description the caller was
|
||||
not served. Once the caller has listed, the same call hands the entry that listing served."""
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.proxy._experimental.mcp_server import operations as mcp_module
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
|
||||
petstore = MCPServer(
|
||||
server_id="petstore-id",
|
||||
|
|
@ -8082,6 +8085,7 @@ async def test_execute_mcp_tool_hands_openapi_registered_tool_metadata_to_pre_ca
|
|||
transport=MCPTransport.http,
|
||||
url=None,
|
||||
spec_path="https://example.com/petstore.yaml",
|
||||
tool_name_to_description={"list_pets": "ADMIN DESC"},
|
||||
)
|
||||
schema = {"type": "object", "properties": {"limit": {"type": "integer"}}}
|
||||
mcp_module.global_mcp_tool_registry.register_tool(
|
||||
|
|
@ -8089,71 +8093,42 @@ async def test_execute_mcp_tool_hands_openapi_registered_tool_metadata_to_pre_ca
|
|||
)
|
||||
manager = mcp_module.global_mcp_server_manager
|
||||
manager._listed_tools_by_server_id.pop(petstore.server_id, None)
|
||||
pre_call_tool_check = AsyncMock(return_value={})
|
||||
alice = UserAPIKeyAuth(api_key="sk-user", user_id="alice")
|
||||
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
|
||||
proxy_logging.pre_call_hook = AsyncMock(return_value={})
|
||||
pre_call_tool_check = AsyncMock(wraps=manager.pre_call_tool_check)
|
||||
|
||||
async def call() -> tuple[MCPTool | None, dict]:
|
||||
await mcp_module.execute_mcp_tool(
|
||||
name="petstore-list_pets",
|
||||
arguments={"limit": 10},
|
||||
allowed_mcp_servers=[petstore],
|
||||
start_time=datetime.now(),
|
||||
user_api_key_auth=alice,
|
||||
)
|
||||
return pre_call_tool_check.call_args.kwargs["tool"], proxy_logging.pre_call_hook.call_args.kwargs["data"]
|
||||
|
||||
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),
|
||||
patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging),
|
||||
):
|
||||
await mcp_module.execute_mcp_tool(
|
||||
name="petstore-list_pets",
|
||||
arguments={"limit": 10},
|
||||
allowed_mcp_servers=[petstore],
|
||||
start_time=datetime.now(),
|
||||
user_api_key_auth=UserAPIKeyAuth(api_key="sk-user", user_id="alice"),
|
||||
never_listed_tool, never_listed_data = await call()
|
||||
manager._record_listed_tools(
|
||||
petstore,
|
||||
[MCPTool(name="list_pets", description="ADMIN DESC", inputSchema=schema)],
|
||||
ListedToolsCaller(user_api_key_auth=alice),
|
||||
)
|
||||
listed_tool, listed_data = await call()
|
||||
finally:
|
||||
mcp_module.global_mcp_tool_registry.unregister_tools_with_prefix("petstore-")
|
||||
manager._listed_tools_by_server_id.pop(petstore.server_id, None)
|
||||
|
||||
handed_tool = pre_call_tool_check.call_args.kwargs["tool"]
|
||||
assert (handed_tool.name, handed_tool.description, handed_tool.input_schema) == (
|
||||
"list_pets",
|
||||
"List the pets",
|
||||
schema,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_mcp_tool_hands_openapi_hooks_the_admin_description_clients_saw():
|
||||
"""tools/list shows the admin's tool_name_to_description wording, so the local-registry call path
|
||||
must hand the pre-call hooks that same wording rather than the generated one."""
|
||||
from litellm.proxy._experimental.mcp_server import operations 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",
|
||||
tool_name_to_description={"getpetbyid": "ADMIN DESC"},
|
||||
)
|
||||
schema = {"type": "object", "properties": {"petId": {"type": "integer"}}}
|
||||
mcp_module.global_mcp_tool_registry.register_tool(
|
||||
name="petstore-getpetbyid", description="Find pet by ID", input_schema=schema, handler=lambda petId: "ok"
|
||||
)
|
||||
manager = mcp_module.global_mcp_server_manager
|
||||
manager._listed_tools_by_server_id.pop(petstore.server_id, None)
|
||||
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),
|
||||
):
|
||||
await mcp_module.execute_mcp_tool(
|
||||
name="petstore-getpetbyid",
|
||||
arguments={"petId": 1},
|
||||
allowed_mcp_servers=[petstore],
|
||||
start_time=datetime.now(),
|
||||
user_api_key_auth=UserAPIKeyAuth(api_key="sk-user", user_id="alice"),
|
||||
)
|
||||
finally:
|
||||
mcp_module.global_mcp_tool_registry.unregister_tools_with_prefix("petstore-")
|
||||
|
||||
handed_tool = pre_call_tool_check.call_args.kwargs["tool"]
|
||||
assert (handed_tool.description, handed_tool.input_schema) == ("ADMIN DESC", schema)
|
||||
assert never_listed_tool is None
|
||||
assert (never_listed_data.get("mcp_tool_description"), never_listed_data.get("mcp_input_schema")) == (None, None)
|
||||
assert listed_tool is not None and (listed_tool.description, listed_tool.input_schema) == ("ADMIN DESC", schema)
|
||||
assert (listed_data["mcp_tool_description"], listed_data["mcp_input_schema"]) == ("ADMIN DESC", schema)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -8267,9 +8242,9 @@ async def test_execute_mcp_tool_hands_openapi_hooks_each_callers_own_listed_entr
|
|||
|
||||
|
||||
@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."""
|
||||
async def test_execute_mcp_tool_runs_the_longer_colliding_operation_and_hands_hooks_no_registry_metadata():
|
||||
"""An OpenAPI operation whose name starts with its own server prefix runs instead of the shorter one, and
|
||||
with no prior listing the pre-call hooks get name and arguments only, never either registry entry."""
|
||||
from litellm.proxy._experimental.mcp_server import operations as mcp_module
|
||||
|
||||
petstore = MCPServer(
|
||||
|
|
@ -8306,11 +8281,7 @@ async def test_execute_mcp_tool_hands_hooks_the_metadata_of_the_operation_it_run
|
|||
finally:
|
||||
registry.unregister_tools_with_prefix("petstore-")
|
||||
|
||||
handed_tool = pre_call_tool_check.call_args.kwargs["tool"]
|
||||
assert (handed_tool.description, handed_tool.input_schema) == (
|
||||
"long",
|
||||
{"type": "object", "properties": {"petId": {"type": "integer"}}},
|
||||
)
|
||||
assert pre_call_tool_check.call_args.kwargs["tool"] is None
|
||||
assert result.content[0].text == "long"
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -3326,8 +3326,7 @@ async def test_request_selected_tool_specific_guardrail_applies_to_virtual_execu
|
|||
default_on=False,
|
||||
custom_code="def apply_guardrail(inputs, request_data, input_type):\n"
|
||||
' if inputs.get("tools", [{}])[0].get("function", {}).get("name") == "execute":\n'
|
||||
' texts = [t.replace("confidential", "redacted") for t in inputs.get("texts", [])]\n'
|
||||
f' return {{"action": "{action}", "reason": "resolved tool blocked", "texts": texts}}\n'
|
||||
f' return {{"action": "{action}", "reason": "resolved tool blocked", "texts": ["redacted"]}}\n'
|
||||
" return allow()\n",
|
||||
)
|
||||
manager: Final = mcp_server_manager.MCPServerManager()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue