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

This reverts commit c274cf321c.

tests/e2e/mcp/mcp_client.py

@@ -11,14 +11,13 @@ 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, UnknownApiError, unwrap
+from e2e_http import Headers, NoBody, Result, Success, unwrap
 from models import KeyGenerateBody, ObjectPermission
 from proxy_client import ProxyClient
 
@@ -271,60 +270,6 @@ 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
@@ -372,39 +317,5 @@ 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,7 +29,6 @@ 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,7 +60,6 @@ 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(
@@ -79,19 +78,22 @@ class TestDatadogMcpRoundTrip:
         )
 
         tool_name = client.await_tool(key, server_id, SEARCH_LOGS_TOOL)
-        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"
+
+        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"
+                    },
                 },
-            },
+            )
         )
         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 unwrap
+from e2e_http import UnknownApiError, unwrap
 from lifecycle import ResourceManager
 from mcp_client import McpClient
 
@@ -72,12 +72,13 @@ class TestMcpKeyWithoutAccessIsDenied:
             "max_tokens": 1000,
             "telemetry": {"intent": "e2e control call proving granted key can invoke Datadog MCP"},
         }
-        permitted_call = client.await_call_tool(
-            permitted_key, server_id=server_id, name=tool_name, arguments=search_args
+        permitted_call = unwrap(
+            client.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}"
 
-        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}"
+        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}")