diff --git a/tests/e2e/mcp/mcp_client.py b/tests/e2e/mcp/mcp_client.py index 6d0f6ddc760..33ec557c339 100644 --- a/tests/e2e/mcp/mcp_client.py +++ b/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 + 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 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 1b53d1ca0b4..f72b75fd43d 100644 --- a/tests/e2e/mcp/test_mcp_access_group_e2e.py +++ b/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()}", diff --git a/tests/e2e/mcp/test_mcp_datadog_e2e.py b/tests/e2e/mcp/test_mcp_datadog_e2e.py index 8a539b86bff..d093e307f99 100644 --- a/tests/e2e/mcp/test_mcp_datadog_e2e.py +++ b/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 diff --git a/tests/e2e/mcp/test_mcp_key_access_e2e.py b/tests/e2e/mcp/test_mcp_key_access_e2e.py index 35c864c07d8..678424e36d1 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 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}"