scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-04 02:31:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 122

test(mcp): reuse the successful catalog snapshot for toolset controls

Browse source
This commit is contained in:
Joshua Valluru 2026-09-17 17:41:03 -07:00
parent bbab62c359
commit 854266d562
2 changed files with 15 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
tests/e2e/mcp/mcp_client.py
View file

@ -278,8 +278,13 @@ class McpClient:
return self.await_tool_entry(key, server_id, needle).name
def await_tool_entry(self, key: str, server_id: str, needle: str) -> McpToolEntry:
tool = self.await_tool_catalog(key, server_id, needle).tool_containing(server_id, needle)
assert tool is not None
return tool
def await_tool_catalog(self, key: str, server_id: str, needle: str) -> McpToolsListResponse:
"""Poll tools/list until `server_id` serves a tool matching `needle`, and
return the successful tool snapshot. Fails at poll_timeout.
return that catalog snapshot. Fails at poll_timeout.
/v1/mcp/server returns as soon as the DB row is written, but the gateway
runs the initialize + tools/list handshake against the upstream lazily on
@ -293,7 +298,7 @@ class McpClient:
if isinstance(result, Success):
tool = result.data.tool_containing(server_id, needle)
if tool is not None:
return tool
return result.data
if time.monotonic() >= deadline:
raise AssertionError(
f"server {server_id} never served a tool matching {needle!r} within "

12
tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py
View file

@ -75,8 +75,10 @@ class TestMcpToolsetEnforcement:
client.await_registered(server_id)
catalog_key: Final = _key(client, resources, "catalog", server_id=server_id)
known_wire: Final = client.await_tool(catalog_key, server_id, SEARCH_LOGS_TOOL)
catalog: Final = unwrap(client.list_tools(catalog_key)).tool_names_for_server(server_id)
snapshot: Final = client.await_tool_catalog(catalog_key, server_id, SEARCH_LOGS_TOOL)
known_wire: Final = snapshot.tool_name_containing(server_id, SEARCH_LOGS_TOOL)
assert known_wire is not None
catalog: Final = snapshot.tool_names_for_server(server_id)
assert len(catalog) > 2, (
f"the Datadog core toolset must serve more tools than the toolset names, or the "
f"restriction has nothing to hide; got {sorted(catalog)}"
@ -149,8 +151,10 @@ def _assert_principal_toolset(client: McpClient, resources: ResourceManager, pri
for transport in client.proxy.replicas_for("/mcp-rest/tools/list").values():
replica = McpClient(proxy=replace(client.proxy, transport=transport))
granted = replica.await_tool_entry(control_key, server_id, SEARCH_LOGS_TOOL)
catalog = unwrap(replica.list_tools(control_key)).tool_names_for_server(server_id)
snapshot = replica.await_tool_catalog(control_key, server_id, SEARCH_LOGS_TOOL)
granted = snapshot.tool_containing(server_id, SEARCH_LOGS_TOOL)
assert granted is not None
catalog = snapshot.tool_names_for_server(server_id)
outside = sorted(catalog - {granted.name})
assert outside, f"uncapped upstream must expose a tool outside the grant: {catalog}"
expected = frozenset({granted.name})