diff --git a/tests/e2e/mcp/mcp_client.py b/tests/e2e/mcp/mcp_client.py index f758a41cae6..4b1725bb205 100644 --- a/tests/e2e/mcp/mcp_client.py +++ b/tests/e2e/mcp/mcp_client.py @@ -11,12 +11,13 @@ request/response bodies are co-located here because only this suite speaks MCP. from __future__ import annotations +import time from collections.abc import Mapping from dataclasses import dataclass from pydantic import BaseModel, ConfigDict, Field, RootModel -from e2e_http import Headers, NoBody, Result, unwrap +from e2e_http import Headers, NoBody, Result, Success, unwrap from models import KeyGenerateBody, ObjectPermission from proxy_client import ProxyClient @@ -223,6 +224,31 @@ class McpClient: response_type=McpToolsListResponse, ) + def await_tool(self, key: str, server_id: str, needle: str) -> str: + """Poll tools/list until `server_id` serves a tool matching `needle`, and + return its fully-qualified name. 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 + the first request that needs it, and reports a server it has not + discovered yet exactly like a dead one: an empty tool list. Waiting is + what separates the two. + """ + deadline = time.monotonic() + self.proxy.poll_timeout + while True: + result = self.list_tools(key) + if isinstance(result, Success): + tool_name = result.data.tool_name_containing(server_id, needle) + if tool_name is not None: + return tool_name + if time.monotonic() >= deadline: + raise AssertionError( + f"server {server_id} never served a tool matching {needle!r} within " + f"{self.proxy.poll_timeout}s of registration (upstream unreachable, or " + f"the key's grant was not applied); last tools/list: {result}" + ) + time.sleep(self.proxy.poll_interval) + def register_mcp_content_filter(self, *, name: str, blocked_keyword: str) -> str: """Register a default-on content-filter guardrail that runs on the MCP tool-call hook (pre_mcp_call) and blocks a single keyword. The keyword is diff --git a/tests/e2e/mcp/test_mcp_datadog_e2e.py b/tests/e2e/mcp/test_mcp_datadog_e2e.py index c772c4d3899..8a539b86bff 100644 --- a/tests/e2e/mcp/test_mcp_datadog_e2e.py +++ b/tests/e2e/mcp/test_mcp_datadog_e2e.py @@ -77,12 +77,7 @@ class TestDatadogMcpRoundTrip: "within the poll deadline; MCP search would have nothing to find" ) - tools = unwrap(client.list_tools(key)) - tool_name = tools.tool_name_containing(server_id, SEARCH_LOGS_TOOL) - assert tool_name is not None, ( - f"granted key never saw {SEARCH_LOGS_TOOL} on server {server_id}; " - f"tools={tools.tool_names_for_server(server_id)}" - ) + tool_name = client.await_tool(key, server_id, SEARCH_LOGS_TOOL) call = unwrap( client.call_tool( diff --git a/tests/e2e/mcp/test_mcp_guardrail_e2e.py b/tests/e2e/mcp/test_mcp_guardrail_e2e.py index 63239444454..dcab235465f 100644 --- a/tests/e2e/mcp/test_mcp_guardrail_e2e.py +++ b/tests/e2e/mcp/test_mcp_guardrail_e2e.py @@ -22,7 +22,7 @@ import pytest from datadog_mcp import SEARCH_LOGS_TOOL, assert_dd_mcp_creds, register_datadog_mcp from e2e_config import DD_SEARCH_FROM, unique_marker -from e2e_http import Result, Success, UnknownApiError, unwrap +from e2e_http import Result, Success, UnknownApiError from lifecycle import ResourceManager from mcp_client import McpCallToolResponse, McpClient, McpToolArguments @@ -75,12 +75,7 @@ class TestMcpToolCallGuardrail: key = client.generate_key(user_id=f"e2e-mcp-guard-{marker}", mcp_servers=[server_id]) resources.defer(lambda: client.proxy.delete_key(key)) - tools = unwrap(client.list_tools(key)) - tool_name = tools.tool_name_containing(server_id, SEARCH_LOGS_TOOL) - assert tool_name is not None, ( - f"granted key never saw {SEARCH_LOGS_TOOL} on server {server_id}; " - f"tools={tools.tool_names_for_server(server_id)}" - ) + tool_name = client.await_tool(key, server_id, SEARCH_LOGS_TOOL) def search(query: str) -> Result[McpCallToolResponse]: arguments: McpToolArguments = { diff --git a/tests/e2e/mcp/test_mcp_key_access_e2e.py b/tests/e2e/mcp/test_mcp_key_access_e2e.py index 412b33d244a..4aeb811a64f 100644 --- a/tests/e2e/mcp/test_mcp_key_access_e2e.py +++ b/tests/e2e/mcp/test_mcp_key_access_e2e.py @@ -48,12 +48,7 @@ class TestMcpKeyWithoutAccessIsDenied: permitted_key = _key(client, resources, mcp_servers=[server_id]) denied_key = _key(client, resources, mcp_servers=None) - permitted = unwrap(client.list_tools(permitted_key)) - tool_name = permitted.tool_name_containing(server_id, SEARCH_LOGS_TOOL) - assert tool_name is not None, ( - f"granted key did not see {SEARCH_LOGS_TOOL} (upstream dead or grant not applied): " - f"{permitted.tool_names_for_server(server_id)}" - ) + _ = client.await_tool(permitted_key, server_id, SEARCH_LOGS_TOOL) denied_tools = unwrap(client.list_tools(denied_key)).tool_names_for_server(server_id) assert denied_tools == frozenset(), ( @@ -73,12 +68,7 @@ class TestMcpKeyWithoutAccessIsDenied: permitted_key = _key(client, resources, mcp_servers=[server_id]) denied_key = _key(client, resources, mcp_servers=None) - permitted = unwrap(client.list_tools(permitted_key)) - tool_name = permitted.tool_name_containing(server_id, SEARCH_LOGS_TOOL) - assert tool_name is not None, ( - f"granted key did not discover {SEARCH_LOGS_TOOL} (upstream dead or grant not applied): " - f"{permitted.tool_names_for_server(server_id)}" - ) + tool_name = client.await_tool(permitted_key, server_id, SEARCH_LOGS_TOOL) search_args = { "query": "service:litellm",