test(e2e): poll MCP tools across multi-worker lag (#35047)

* fix(mcp): resolve call_tool by registry without requiring tool map

Multi-worker reloads put MCP servers in the registry from the DB but do
not re-run tools/list on every process. Gating call_tool on
tool_name_to_mcp_server_name_mapping made cold workers 500 with Tool not
found after another worker had already listed the tool. Treat a registry
match on server id/name/alias as enough; upstream rejects unknown tools

* test(e2e): poll MCP register, tools/list, and tools/call across multi-worker lag

Stage multi-worker gateways only load MCP servers and tool maps on the
process that handled the request. Poll until the server is listed, the
tool appears on tools/list, and tools/call is not a cold-worker 500 so
key-access and Datadog MCP e2e stop racing the LB

* Revert "fix(mcp): resolve call_tool by registry without requiring tool map"

This reverts commit 8b56e51e39.

* test(e2e): tighten MCP multi-worker lag classifier

Only retry tools/call on gateway shapes Tool <name> not found and
server_not_found, not any 500 that mentions tool/server not found, so
upstream failures are not retried until the poll deadline

* test(e2e): drop unit file for MCP lag classifier

The live await_call_tool polls already cover multi-worker lag; a separate
string-match unit module is not worth keeping

(cherry picked from commit c274cf321c)

tests/e2e/mcp/mcp_client.py

@@ -11,13 +11,14 @@ request/response bodies are co-located here because only this suite speaks MCP.
 
 from __future__ import annotations
 
+import re
 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, Success, unwrap
+from e2e_http import Headers, NoBody, Result, Success, UnknownApiError, unwrap
 from models import KeyGenerateBody, ObjectPermission
 from proxy_client import ProxyClient
 
@@ -270,6 +271,60 @@ class McpClient:
                 )
             time.sleep(self.proxy.poll_interval)
 
+    def await_call_tool(
+        self,
+        key: str,
+        *,
+        server_id: str,
+        name: str,
+        arguments: McpToolArguments,
+    ) -> McpCallToolResponse:
+        """Poll tools/call until the result is not a multi-worker registry miss.
+
+        Retries only on the gateway's own cold-worker 500 shapes (Tool <name>
+        not found / server_not_found). Upstream tool errors and other 500s fail
+        immediately so non-idempotent calls are not repeated.
+        """
+        deadline = time.monotonic() + self.proxy.poll_timeout
+        last: Result[McpCallToolResponse] | None = None
+        while True:
+            last = self.call_tool(key, server_id=server_id, name=name, arguments=arguments)
+            if not _is_mcp_not_synced(last, tool_name=name):
+                return unwrap(last)
+            if time.monotonic() >= deadline:
+                raise AssertionError(
+                    f"tools/call for {name!r} on server {server_id} still missing on the "
+                    f"data plane after {self.proxy.poll_timeout}s (multi-worker registry lag); "
+                    f"last result: {last}"
+                )
+            time.sleep(self.proxy.poll_interval)
+
+    def await_call_tool_denied(
+        self,
+        key: str,
+        *,
+        server_id: str,
+        name: str,
+        arguments: McpToolArguments,
+    ) -> UnknownApiError:
+        """Poll tools/call until a cold-worker miss clears and the call is 403 access_denied."""
+        deadline = time.monotonic() + self.proxy.poll_timeout
+        last: Result[McpCallToolResponse] | None = None
+        while True:
+            last = self.call_tool(key, server_id=server_id, name=name, arguments=arguments)
+            if isinstance(last, UnknownApiError) and last.status_code == 403:
+                return last
+            if not _is_mcp_not_synced(last, tool_name=name):
+                raise AssertionError(
+                    f"ungranted key's tools/call was not 403 access_denied: {last}"
+                )
+            if time.monotonic() >= deadline:
+                raise AssertionError(
+                    f"ungranted key never got 403 for {name!r} within {self.proxy.poll_timeout}s; "
+                    f"last result: {last}"
+                )
+            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
@@ -317,5 +372,39 @@ class McpClient:
         )
 
 
+def _is_mcp_not_synced(
+    result: Result[McpCallToolResponse],
+    *,
+    tool_name: str | None = None,
+) -> bool:
+    """True only for gateway multi-worker registry misses, not upstream errors.
+
+    Matches the proxy's own shapes:
+    - ValueError ``Tool <name> not found`` wrapped as HTTP 500 (cold tool map /
+      unresolved server on this process)
+    - REST ``server_not_found`` when this worker has not loaded the MCP server row
+
+    Does not treat arbitrary 500 bodies that merely mention "tool" and "not found"
+    (e.g. upstream MCP payload text) as lag, so await_call_tool does not retry
+    real failures or non-idempotent calls.
+    """
+    if not isinstance(result, UnknownApiError) or result.status_code != 500:
+        return False
+    body = result.body
+    body_l = body.lower()
+
+    if "server_not_found" in body_l:
+        return True
+    if re.search(r"mcp server ['\"][^'\"]+['\"] was not found", body_l):
+        return True
+
+    # Gateway: "Tool search_datadog_logs not found" (optionally inside a longer message)
+    if tool_name is not None:
+        return (
+            re.search(rf"\btool\s+{re.escape(tool_name)}\s+not found\b", body_l) is not None
+        )
+    return re.search(r"\btool\s+\S+\s+not found\b", body_l) is not None
+
+
 def build_client(proxy: ProxyClient) -> McpClient:
     return McpClient(proxy=proxy)

tests/e2e/mcp/test_mcp_access_group_e2e.py

@@ -29,6 +29,7 @@ class TestMcpAccessGroupToolSelection:
     ) -> None:
         group = f"e2e-mcp-grp-{unique_marker()}"
         server_id = register_datadog_mcp(client, resources, mcp_access_groups=[group])
+        client.await_registered(server_id)
 
         granted = client.generate_key(
             user_id=f"e2e-mcp-ag-granted-{unique_marker()}",

tests/e2e/mcp/test_mcp_datadog_e2e.py

@@ -60,6 +60,7 @@ class TestDatadogMcpRoundTrip:
         _assert_datadog_logger_active(client.proxy)
 
         server_id = register_datadog_mcp(client, resources)
+        client.await_registered(server_id)
         marker = f"{MARKER_PREFIX}{unique_marker()}"
 
         key = client.generate_key(
@@ -78,22 +79,19 @@ class TestDatadogMcpRoundTrip:
         )
 
         tool_name = client.await_tool(key, server_id, SEARCH_LOGS_TOOL)
-
-        call = unwrap(
-            client.call_tool(
-                key,
-                server_id=server_id,
-                name=tool_name,
-                arguments={
-                    "query": marker,
-                    "from": DD_SEARCH_FROM,
-                    "to": "now",
-                    "max_tokens": 5000,
-                    "telemetry": {
-                        "intent": "e2e assert seeded litellm completion log is searchable via MCP"
-                    },
+        call = client.await_call_tool(
+            key,
+            server_id=server_id,
+            name=tool_name,
+            arguments={
+                "query": marker,
+                "from": DD_SEARCH_FROM,
+                "to": "now",
+                "max_tokens": 5000,
+                "telemetry": {
+                    "intent": "e2e assert seeded litellm completion log is searchable via MCP"
                 },
-            )
+            },
         )
         assert call.is_error is not True, f"search_datadog_logs errored: {call}"
         body = call.all_text

tests/e2e/mcp/test_mcp_key_access_e2e.py

@@ -16,7 +16,7 @@ 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 UnknownApiError, unwrap
+from e2e_http import unwrap
 from lifecycle import ResourceManager
 from mcp_client import McpClient
 
@@ -72,13 +72,12 @@ class TestMcpKeyWithoutAccessIsDenied:
             "max_tokens": 1000,
             "telemetry": {"intent": "e2e control call proving granted key can invoke Datadog MCP"},
         }
-        permitted_call = unwrap(
-            client.call_tool(permitted_key, server_id=server_id, name=tool_name, arguments=search_args)
+        permitted_call = client.await_call_tool(
+            permitted_key, server_id=server_id, name=tool_name, arguments=search_args
         )
         assert permitted_call.is_error is not True, f"granted key's tool call errored: {permitted_call}"
 
-        match client.call_tool(denied_key, server_id=server_id, name=tool_name, arguments=search_args):
-            case UnknownApiError(status_code=403, body=body):
-                assert "access_denied" in body, f"403 was not an MCP access denial: {body}"
-            case other:
-                pytest.fail(f"ungranted key's tool call was not refused with 403 access_denied: {other}")
+        denied = client.await_call_tool_denied(
+            denied_key, server_id=server_id, name=tool_name, arguments=search_args
+        )
+        assert "access_denied" in denied.body, f"403 was not an MCP access denial: {denied.body}"