From c873dc016c2dd14d31a2dba22bf52f821b481bbb Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 28 Jul 2026 22:21:52 -0700 Subject: [PATCH] Revert "test(e2e): poll MCP tools across multi-worker lag (#35047)" This reverts commit c274cf321c5c35c629220a89bb497d15b56f870f. --- tests/e2e/mcp/mcp_client.py | 91 +--------------------- tests/e2e/mcp/test_mcp_access_group_e2e.py | 1 - tests/e2e/mcp/test_mcp_datadog_e2e.py | 28 +++---- tests/e2e/mcp/test_mcp_key_access_e2e.py | 15 ++-- 4 files changed, 24 insertions(+), 111 deletions(-) diff --git a/tests/e2e/mcp/mcp_client.py b/tests/e2e/mcp/mcp_client.py index 33ec557c339..6d0f6ddc760 100644 --- a/tests/e2e/mcp/mcp_client.py +++ b/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 - 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 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) diff --git a/tests/e2e/mcp/test_mcp_access_group_e2e.py b/tests/e2e/mcp/test_mcp_access_group_e2e.py index f72b75fd43d..1b53d1ca0b4 100644 --- a/tests/e2e/mcp/test_mcp_access_group_e2e.py +++ b/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()}", diff --git a/tests/e2e/mcp/test_mcp_datadog_e2e.py b/tests/e2e/mcp/test_mcp_datadog_e2e.py index d093e307f99..8a539b86bff 100644 --- a/tests/e2e/mcp/test_mcp_datadog_e2e.py +++ b/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 diff --git a/tests/e2e/mcp/test_mcp_key_access_e2e.py b/tests/e2e/mcp/test_mcp_key_access_e2e.py index 678424e36d1..35c864c07d8 100644 --- a/tests/e2e/mcp/test_mcp_key_access_e2e.py +++ b/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}")