scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-10 03:28:53 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 85

fix(mcp): look a called tool up in the listed catalog by its bare name only

Browse source 
get_listed_tool stripped the server prefix a second time when the exact
name was absent from the caller's listing, so a never-listed upstream tool
whose bare name starts with the server prefix resolved to the listed
sibling and that sibling's description and input schema reached the
pre-call hooks for a call to a different tool. Every caller already passes
the once-stripped bare name, so the lookup is now exact.

Tests that looked the catalog up by a prefixed name now use the bare name
the callers pass; two new tests pin the never-listed sibling case at the
manager and at the tools/call path.
This commit is contained in:
Yucheng He 2026-10-01 12:56:38 -07:00
parent 46639297ad
commit ed75c1aef7
3 changed files with 90 additions and 25 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -5691,7 +5691,7 @@ class MCPServerManager:
listed: Final = self._listed_tools_by_server_id.get(server.server_id, MappingProxyType({})).get(identity)
if not listed:
return None
return listed.get(name) or listed.get(strip_known_server_prefix(name, server))
return listed.get(name)
def _create_prefixed_prompts(
self, prompts: Sequence[Prompt], server: MCPServer, add_prefix: bool = True

66
tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py
View file

@ -7131,16 +7131,35 @@ class TestMCPServerManager:
hook_kwargs = proxy_logging_obj._create_mcp_request_object_from_kwargs.call_args.args[0]
assert (hook_kwargs["tool_description"], hook_kwargs["tool_input_schema"]) == (None, None)
def test_get_listed_tool_resolves_prefixed_name_and_latest_listing(self):
def test_get_listed_tool_resolves_the_bare_name_from_the_latest_listing(self):
manager = MCPServerManager()
server = MCPServer(server_id="srv", name="srv", transport=MCPTransport.http, url="http://srv")
manager._create_prefixed_tools([MCPTool(name="echo", description="v1", inputSchema={})], server)
manager._create_prefixed_tools([MCPTool(name="echo", description="v2", inputSchema={})], server)
by_prefixed_name = manager.get_listed_tool(server, "srv-echo")
assert by_prefixed_name is not None and by_prefixed_name.description == "v2"
latest = manager.get_listed_tool(server, "echo")
assert latest is not None and latest.description == "v2"
assert manager.get_listed_tool(server, "missing") is None
def test_get_listed_tool_never_strips_the_bare_name_it_is_given(self):
"""The lookup is exact: a never-listed tool whose bare name starts with the server prefix is not the
listed sibling that stripping the prefix again would name."""
manager = MCPServerManager()
server = MCPServer(server_id="srv-id", name="srv", alias="srv", transport=MCPTransport.http, url="http://srv")
caller = ListedToolsCaller(user_api_key_auth=UserAPIKeyAuth(api_key="sk-user", user_id="alice"))
manager._create_prefixed_tools(
[
MCPTool(name="foo", description="Fetches foo records", inputSchema={"type": "object"}),
MCPTool(name="bar", description="Fetches bar records", inputSchema={"type": "object"}),
],
server,
caller=caller,
)
assert manager.get_listed_tool(server, "srv-foo", caller) is None
listed = manager.get_listed_tool(server, "foo", caller)
assert listed is not None and listed.description == "Fetches foo records"
@pytest.mark.asyncio
async def test_get_listed_tool_uses_admin_description_override_clients_saw(self):
schema = {"type": "object", "properties": {"text": {"type": "string"}}}
@ -7157,7 +7176,7 @@ class TestMCPServerManager:
)
await manager._get_tools_from_server(server, add_prefix=True)
overridden = manager.get_listed_tool(server, "srv-echo")
overridden = manager.get_listed_tool(server, "echo")
assert overridden is not None
assert (overridden.name, overridden.description, overridden.input_schema) == ("echo", "Admin wording", schema)
untouched = manager.get_listed_tool(server, "ping")
@ -7178,7 +7197,7 @@ class TestMCPServerManager:
served = await manager._get_tools_from_server(server, add_prefix=True, proxy_logging_obj=proxy_logging_obj)
assert [tool.description for tool in served] == ["Read a [MASKED] note"]
listed = manager.get_listed_tool(server, "notes-read_note")
listed = manager.get_listed_tool(server, "read_note")
assert listed is not None and listed.description == "Read a [MASKED] note", (
"tools/call must be evaluated against the description tools/list served"
)
@ -7360,15 +7379,15 @@ class TestMCPServerManager:
caller=ListedToolsCaller(user_api_key_auth=bob),
)
alice_tool = manager.get_listed_tool(server, "srv-read", ListedToolsCaller(user_api_key_auth=alice))
bob_tool = manager.get_listed_tool(server, "srv-read", ListedToolsCaller(user_api_key_auth=bob))
alice_tool = manager.get_listed_tool(server, "read", ListedToolsCaller(user_api_key_auth=alice))
bob_tool = manager.get_listed_tool(server, "read", ListedToolsCaller(user_api_key_auth=bob))
assert alice_tool is not None and (alice_tool.description, alice_tool.input_schema) == (
"alice view",
alice_schema,
)
assert bob_tool is not None and (bob_tool.description, bob_tool.input_schema) == ("bob view", bob_schema)
carol = ListedToolsCaller(user_api_key_auth=UserAPIKeyAuth(user_id="carol", token="k"))
assert manager.get_listed_tool(server, "srv-read", carol) is None
assert manager.get_listed_tool(server, "read", carol) is None
shared = MCPServer(server_id="shared", name="shared", transport=MCPTransport.http, url="http://shared")
manager._create_prefixed_tools(
@ -7434,11 +7453,11 @@ class TestMCPServerManager:
[MCPTool(name="turn", description="Catalog B", inputSchema={})], server, caller=caller_b
)
for_a = manager.get_listed_tool(server, "srv-turn", caller_a)
for_b = manager.get_listed_tool(server, "srv-turn", caller_b)
for_a = manager.get_listed_tool(server, "turn", caller_a)
for_b = manager.get_listed_tool(server, "turn", caller_b)
assert for_a is not None and for_a.description == "Catalog A"
assert for_b is not None and for_b.description == "Catalog B"
assert manager.get_listed_tool(server, "srv-turn", ListedToolsCaller()) is None
assert manager.get_listed_tool(server, "turn", ListedToolsCaller()) is None
def test_shared_server_ignores_headers_it_never_forwards(self):
manager = MCPServerManager()
@ -7676,9 +7695,9 @@ class TestMCPServerManager:
manager._create_prefixed_tools(
[MCPTool(name="turn", description="alice view", inputSchema={})], server, caller=alice
)
assert manager.get_listed_tool(server, "srv-turn", bob) is None
assert manager.get_listed_tool(server, "turn", bob) is None
for_alice = manager.get_listed_tool(server, "srv-turn", alice)
for_alice = manager.get_listed_tool(server, "turn", alice)
assert for_alice is not None and for_alice.description == "alice view"
def test_signed_server_slot_splits_on_the_callers_key_not_only_the_user(self):
@ -7696,10 +7715,10 @@ class TestMCPServerManager:
manager._create_prefixed_tools(
[MCPTool(name="turn", description="slot a", inputSchema={})], server, caller=alice
)
assert manager.get_listed_tool(server, "srv-turn", bob) is None
assert manager.get_listed_tool(server, "turn", bob) is None
same_key = ListedToolsCaller(user_api_key_auth=UserAPIKeyAuth(user_id="same-user", api_key="sk-alpha"))
listed = manager.get_listed_tool(server, "srv-turn", same_key)
listed = manager.get_listed_tool(server, "turn", same_key)
assert listed is not None and listed.description == "slot a"
@ -7746,7 +7765,7 @@ class TestMCPServerManager:
proxy_logging_obj.during_call_hook = AsyncMock(return_value=None)
await manager.call_tool(
server_name="catalog",
name="catalog-turn",
name="turn",
arguments={"turn": "A-1"},
user_api_key_auth=UserAPIKeyAuth(api_key="sk-litellm", user_id="shared-key"),
proxy_logging_obj=proxy_logging_obj,
@ -7785,13 +7804,13 @@ class TestMCPServerManager:
[MCPTool(name="read", description="u1 again", inputSchema={})], server, caller=callers[1]
)
assert manager.get_listed_tool(server, "srv-read", callers[0]) is None
second = manager.get_listed_tool(server, "srv-read", callers[1])
assert manager.get_listed_tool(server, "read", callers[0]) is None
second = manager.get_listed_tool(server, "read", callers[1])
assert second is not None and second.description == "u1 again"
newest = manager.get_listed_tool(server, "srv-read", callers[-1])
newest = manager.get_listed_tool(server, "read", callers[-1])
assert newest is not None and newest.description == callers[-1].user_api_key_auth.user_id
assert len(manager._listed_tools_by_server_id[server.server_id]) == _LISTED_TOOLS_CALLERS_PER_SERVER + 1
shared = manager.get_listed_tool(server, "srv-read")
shared = manager.get_listed_tool(server, "read")
assert shared is not None and shared.description == "shared"
@pytest.mark.asyncio
@ -7826,10 +7845,9 @@ class TestMCPServerManager:
global_mcp_tool_registry.unregister_tools_with_prefix("petstore-")
assert [t.name for t in listed] == ["petstore-list_pets" if add_prefix else "list_pets"]
for name in ("list_pets", "petstore-list_pets"):
tool = manager.get_listed_tool(server, name)
assert tool is not None and tool.description == "List pets"
assert tool.input_schema["properties"] == {"limit": {"type": "integer"}}
tool = manager.get_listed_tool(server, "list_pets")
assert tool is not None and tool.description == "List pets"
assert tool.input_schema["properties"] == {"limit": {"type": "integer"}}
@pytest.mark.asyncio
async def test_openapi_listing_ignores_overlapping_server_prefix(self):

47
tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py
View file

@ -8285,6 +8285,53 @@ async def test_execute_mcp_tool_runs_the_longer_colliding_operation_and_hands_ho
assert result.content[0].text == "long"
@pytest.mark.asyncio
async def test_execute_mcp_tool_hands_hooks_nothing_for_a_never_listed_operation_named_after_a_listed_one():
"""After the caller listed ``get_pet``, a call to the never-listed ``petstore-get_pet`` operation hands the
pre-call hooks name and arguments only, not the listed sibling's description and schema."""
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",
)
registry = mcp_module.global_mcp_tool_registry
registry.register_tool(
name="petstore-petstore-get_pet", description="long", input_schema={}, handler=lambda: "long"
)
manager = mcp_module.global_mcp_server_manager
alice = UserAPIKeyAuth(api_key="sk-user", user_id="alice")
manager._record_listed_tools(
petstore,
[MCPTool(name="get_pet", description="Fetches pet records. FLAGWORD", inputSchema={"type": "object"})],
ListedToolsCaller(user_api_key_auth=alice),
)
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=alice,
)
finally:
registry.unregister_tools_with_prefix("petstore-")
manager._listed_tools_by_server_id.pop(petstore.server_id, None)
assert pre_call_tool_check.call_args.kwargs["tool"] is None
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.