From 9966b8855cabd4884725ff86147169bb758feb77 Mon Sep 17 00:00:00 2001 From: Yucheng Low Date: Fri, 28 Aug 2026 10:34:46 -0700 Subject: [PATCH] test: poll the exact tool set to the shared deadline instead of trusting one replica's sync Co-Authored-By: Claude Fable 5 --- .../mcp/test_mcp_toolset_enforcement_e2e.py | 38 +++++++++++++------ 1 file changed, 27 insertions(+), 11 deletions(-) diff --git a/tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py b/tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py index b6c1f66328d..fe10bf7c0e9 100644 --- a/tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py +++ b/tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py @@ -17,8 +17,9 @@ whose toolset ceils the servers their own key grants. from __future__ import annotations -import pytest +import time +import pytest from datadog_mcp import SEARCH_LOGS_TOOL, register_datadog_mcp from e2e_config import DD_SEARCH_FROM, unique_marker from e2e_http import unwrap @@ -59,6 +60,25 @@ def _control_tools(client: McpClient, resources: ResourceManager, server_id: str return granted_tool, outside[0] +def _await_exact_tools(client: McpClient, key: str, server_id: str, expected: frozenset[str]) -> None: + """Poll tools/list until the key's view of `server_id` is exactly `expected`. + + A single successful listing proves one replica has synced the grant rows; the + next request may land on another. Polling the exact set to the shared deadline + is the propagation barrier, and a genuine enforcement leak stays red because + the leaked set never converges to `expected`.""" + deadline = time.monotonic() + client.proxy.poll_timeout + while True: + listed = unwrap(client.list_tools(key)).tool_names_for_server(server_id) + if listed == expected: + return + if time.monotonic() >= deadline: + raise AssertionError( + f"toolset-scoped key must list exactly {sorted(expected)}; the toolset ceiling leaked: {sorted(listed)}" + ) + time.sleep(client.proxy.poll_interval) + + def _assert_toolset_ceiling( client: McpClient, scoped_key: str, @@ -68,12 +88,10 @@ def _assert_toolset_ceiling( outside_tool: str, ) -> None: """The scoped key sees exactly the toolset's tool, is 403-refused on a tool - outside it, and can still execute the granted one.""" - _ = client.await_tool(scoped_key, server_id, SEARCH_LOGS_TOOL) - listed = unwrap(client.list_tools(scoped_key)).tool_names_for_server(server_id) - assert listed == frozenset({granted_tool}), ( - f"toolset-scoped key must list exactly {granted_tool!r}; the toolset ceiling leaked: {sorted(listed)}" - ) + outside it, and can still execute the granted one. Every leg polls to the + shared deadline so it cannot flake on a replica that has not synced the + just-created grant rows yet.""" + _await_exact_tools(client, scoped_key, server_id, frozenset({granted_tool})) denied = client.await_call_tool_denied(scoped_key, server_id=server_id, name=outside_tool, arguments={}) assert denied.status_code == 403 @@ -133,8 +151,7 @@ class TestMcpToolsetEnforcementPerLevel: object_permission=ObjectPermission(mcp_servers=[server_id], mcp_toolsets=[toolset_id]), ) resources.defer(lambda: client.delete_org(org_id)) - # The team declares no MCP grants of its own; whatever its key can reach - # comes from the org, so the org's toolset must cap it. + # the team declares no MCP grants of its own, so only the org's cap applies team_id = client.create_team(alias=f"e2e-ts-org-team-{unique_marker()}", organization_id=org_id) resources.defer(lambda: client.delete_team(team_id)) team_key = client.generate_key( @@ -162,8 +179,7 @@ class TestMcpToolsetEnforcementPerLevel: server_id, toolset_id = _open_server_and_toolset(client, resources) granted_tool, outside_tool = _control_tools(client, resources, server_id) - # The user's row holds only the toolset; their key grants the whole - # server. The user-level toolset is a ceiling over the key's grant. + # the key grants the whole server; the user row's toolset must ceil it user_id = client.create_user( user_email=f"e2e-ts-user-{unique_marker()}@example.com", object_permission=ObjectPermission(mcp_toolsets=[toolset_id]),