From dc070752eef016512dad6a1958f01cbd867f474e Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Mon, 27 Jul 2026 12:35:57 -0700 Subject: [PATCH] test(e2e): wait for MCP tool discovery instead of racing it /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. Every MCP test read tools/list immediately after registering, so it raced that handshake. The gateway reports a server it has not discovered yet exactly like a dead one: it catches the per-server handshake exception and returns an empty tool list. The tests asserted on a single read, so the race surfaced as "granted key never saw search_datadog_logs; tools=frozenset()" while a sibling test against the same upstream in the same run passed. Add McpClient.await_tool, which polls tools/list to the suite's existing poll_timeout and returns the qualified tool name, and route the four discovery sites through it. An unreachable upstream or an unapplied grant still fails, and the failure now names the last tools/list result. Refs LIT-4821 --- tests/e2e/mcp/mcp_client.py | 28 +++++++++++++++++++++++- tests/e2e/mcp/test_mcp_datadog_e2e.py | 7 +----- tests/e2e/mcp/test_mcp_guardrail_e2e.py | 9 ++------ tests/e2e/mcp/test_mcp_key_access_e2e.py | 14 ++---------- 4 files changed, 32 insertions(+), 26 deletions(-) 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",