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

@@ -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

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(

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 = {

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",