test(e2e): make MCP and prometheus e2e tests robust to data-plane sync lag

tests/e2e/logging/test_prometheus_cardinality_e2e.py

@@ -9,8 +9,11 @@
 the aliases and fail here.
 
 Scraping goes through ``transport.probe`` (raw text) and is parsed with
-prometheus_client; the metric is eventually consistent (it increments on the
-success-logging callback), so the scrape polls to a deadline.
+prometheus_client. ``/metrics`` is per-pod behind a round-robin LB and the metric
+is eventually consistent (it increments on the success-logging callback), so the
+poll unions the aliases seen across scrapes until the deadline: counters persist
+on whichever pod served the driver call, so repeated scrapes converge without
+re-sending any billable traffic.
 """
 
 from __future__ import annotations
@@ -58,13 +61,14 @@ def test_distinct_key_aliases_produce_distinct_series(
         deadline = time.monotonic() + client.proxy.poll_timeout
         seen: frozenset[str] = frozenset()
         while time.monotonic() < deadline:
-            seen = _aliases_in_metric(client.scrape_metrics(), REQUESTS_METRIC, ALIAS_LABEL)
+            seen = seen | _aliases_in_metric(client.scrape_metrics(), REQUESTS_METRIC, ALIAS_LABEL)
             if wanted <= seen:
                 break
             time.sleep(client.proxy.poll_interval)
 
         missing = wanted - seen
         assert not missing, (
-            f"{REQUESTS_METRIC} is missing a per-key series for aliases {sorted(missing)}; "
+            f"{REQUESTS_METRIC} never exposed a per-key series for aliases {sorted(missing)} "
+            f"on any scraped pod within the deadline; "
             f"each distinct {ALIAS_LABEL} must grow its own series"
         )

tests/e2e/mcp/mcp_client.py

@@ -195,6 +195,27 @@ def registered_servers(self) -> list[McpServerRow]:
             )
         ).root
 
+    def await_registered(self, server_id: str) -> None:
+        """Poll /v1/mcp/server until `server_id` is listed. Fails at poll_timeout.
+
+        The DB row exists the moment registration returns, but a data-plane pod
+        answers the listing from a registry it refreshes on a periodic DB sync, so a
+        pod that joined the load balancer after the write reports the server as
+        absent until its first sync.
+        """
+        deadline = time.monotonic() + self.proxy.poll_timeout
+        while True:
+            registered = frozenset(row.server_id for row in self.registered_servers())
+            if server_id in registered:
+                return
+            if time.monotonic() >= deadline:
+                raise AssertionError(
+                    f"registered server {server_id} still absent from /v1/mcp/server "
+                    f"{self.proxy.poll_timeout}s after registration (the data plane never synced "
+                    f"the row): {registered}"
+                )
+            time.sleep(self.proxy.poll_interval)
+
     def generate_key(
         self,
         *,

tests/e2e/mcp/test_mcp_access_group_e2e.py

@@ -44,12 +44,7 @@ def test_access_group_scopes_tool_selection(
         )
         resources.defer(lambda: client.proxy.delete_key(other))
 
-        granted_tools = unwrap(client.list_tools(granted))
-        assert granted_tools.tool_name_containing(server_id, SEARCH_LOGS_TOOL) is not None, (
-            f"key granted access group {group} did not see the tagged server's tool "
-            f"(upstream dead or access-group grant not applied): "
-            f"{granted_tools.tool_names_for_server(server_id)}"
-        )
+        _ = client.await_tool(granted, server_id, SEARCH_LOGS_TOOL)
 
         other_tools = unwrap(client.list_tools(other)).tool_names_for_server(server_id)
         assert other_tools == frozenset(), (

tests/e2e/mcp/test_mcp_guardrail_e2e.py

@@ -29,11 +29,11 @@
 pytestmark = pytest.mark.e2e
 
 # Stage runs several data-plane pods behind the shared key, and each picks up a
-# newly registered guardrail only on its next periodic DB sync (~30s in
+# newly registered guardrail or MCP server only on its next periodic DB sync (~30s in
 # proxy_server.py). Every pod is guaranteed to have refreshed only once a full sync
-# interval has elapsed since the create; before then a banned call routed to a
-# lagging pod passes through as legitimate in-flight propagation, not a leak.
-GUARDRAIL_FULL_SYNC_SECONDS = 40.0
+# interval has elapsed since the later of those two writes; before then a banned call
+# routed to a lagging pod passes through as legitimate in-flight propagation, not a leak.
+FULL_SYNC_SECONDS = 40.0
 POST_SYNC_VERIFICATION_CALLS = 4
 
 
@@ -53,6 +53,30 @@ def _poll_until_blocked(
     return last
 
 
+def _pod_lacks_mcp_server(result: Result[McpCallToolResponse]) -> bool:
+    """True when the pod that served the call answered as though the MCP server or its
+    tool does not exist (500 "Tool ... not found"), i.e. its MCP registry has not synced
+    yet and the request never reached the guardrail at all."""
+    if not isinstance(result, UnknownApiError) or result.status_code != 500:
+        return False
+    body = result.body.lower()
+    return "not found" in body and ("tool" in body or "server" in body)
+
+
+def _search_on_synced_pod(
+    search: Callable[[str], Result[McpCallToolResponse]], query: str, client: McpClient
+) -> Result[McpCallToolResponse]:
+    """Issue `query`, retrying to the poll deadline only while the serving pod does not
+    know the MCP server yet. Every other outcome, guardrail block or pass-through, comes
+    back untouched so the caller's assertion still decides it."""
+    deadline = time.monotonic() + client.proxy.poll_timeout
+    last = search(query)
+    while _pod_lacks_mcp_server(last) and time.monotonic() < deadline:
+        time.sleep(client.proxy.poll_interval)
+        last = search(query)
+    return last
+
+
 class TestMcpToolCallGuardrail:
     @pytest.mark.covers(
         "guardrail.litellm_content_filter.pre_mcp_call.blocks",
@@ -72,6 +96,7 @@ def test_content_filter_blocks_banned_keyword_in_tool_args(
         resources.defer(lambda: client.delete_guardrail(guardrail_id))
 
         server_id = register_datadog_mcp(client, resources)
+        server_registered_at = time.monotonic()
         key = client.generate_key(user_id=f"e2e-mcp-guard-{marker}", mcp_servers=[server_id])
         resources.defer(lambda: client.proxy.delete_key(key))
 
@@ -105,31 +130,33 @@ def search(query: str) -> Result[McpCallToolResponse]:
             case _:
                 pytest.fail(
                     "content_filter never blocked the banned keyword on the MCP tool call within "
-                    f"{client.proxy.poll_timeout}s (guardrail sync to the data plane never landed); "
-                    f"last result: {blocked}"
+                    f"{client.proxy.poll_timeout}s (the guardrail or the MCP server never synced to "
+                    f"the data plane); last result: {blocked}"
                 )
 
         # The block above only proves the one pod that served it has synced; another
         # pod could still lack the guardrail and let the banned call reach Datadog.
-        # Wait out the full sync interval from the create so every pod has refreshed
-        # from the DB, then require the banned call to stay blocked across several
-        # attempts. A pass-through now is a genuine partial-propagation leak, not a
-        # race. Client load balancing still can't guarantee every pod is hit, so this
+        # Wait out the full sync interval from the later of the guardrail create and the
+        # MCP server registration (each syncs on its own clock, so the earlier write's
+        # deadline can elapse while a pod still lacks the other) so every pod has
+        # refreshed from the DB, then require the banned call to stay blocked across
+        # several attempts. A pass-through now is a genuine partial-propagation leak, not
+        # a race. Client load balancing still can't guarantee every pod is hit, so this
         # samples several worker selections rather than proving all pods synced.
-        sync_remaining = guardrail_created_at + GUARDRAIL_FULL_SYNC_SECONDS - time.monotonic()
+        sync_remaining = max(guardrail_created_at, server_registered_at) + FULL_SYNC_SECONDS - time.monotonic()
         if sync_remaining > 0:
             time.sleep(sync_remaining)
         for attempt in range(1, POST_SYNC_VERIFICATION_CALLS + 1):
-            reblocked = search(f"still about {banned_keyword} #{attempt}")
+            reblocked = _search_on_synced_pod(search, f"still about {banned_keyword} #{attempt}", client)
             assert isinstance(reblocked, UnknownApiError) and reblocked.status_code == 400, (
-                "after the guardrail sync interval every data-plane pod must block the banned "
-                f"keyword, but attempt {attempt} of {POST_SYNC_VERIFICATION_CALLS} was allowed "
-                f"through (a pod still lacks the guardrail): {reblocked}"
+                "after the sync interval every data-plane pod must block the banned keyword, but "
+                f"attempt {attempt} of {POST_SYNC_VERIFICATION_CALLS} was not blocked (a pod still "
+                f"lacks the guardrail, or never synced the MCP server): {reblocked}"
             )
             if attempt < POST_SYNC_VERIFICATION_CALLS:
                 time.sleep(client.proxy.poll_interval)
 
-        allowed = search(f"e2e-clean-{marker}")
+        allowed = _search_on_synced_pod(search, f"e2e-clean-{marker}", client)
         match allowed:
             case Success(data=result):
                 assert result.is_error is not True, (

tests/e2e/mcp/test_mcp_key_access_e2e.py

@@ -30,11 +30,6 @@ def _key(client: McpClient, resources: ResourceManager, *, mcp_servers: list[str
     return key
 
 
-def _assert_registered(client: McpClient, server_id: str) -> None:
-    registered = {row.server_id for row in client.registered_servers()}
-    assert server_id in registered, f"registered server {server_id} absent from /v1/mcp/server: {registered}"
-
-
 class TestMcpKeyWithoutAccessIsDenied:
     @pytest.mark.covers("mcp.list_tools.api_key.denied_without_permission")
     def test_list_tools_denied_without_permission(
@@ -43,7 +38,7 @@ def test_list_tools_denied_without_permission(
         resources: ResourceManager,
     ) -> None:
         server_id = register_datadog_mcp(client, resources)
-        _assert_registered(client, server_id)
+        client.await_registered(server_id)
 
         permitted_key = _key(client, resources, mcp_servers=[server_id])
         denied_key = _key(client, resources, mcp_servers=None)
@@ -63,7 +58,7 @@ def test_call_tool_denied_without_permission(
         resources: ResourceManager,
     ) -> None:
         server_id = register_datadog_mcp(client, resources)
-        _assert_registered(client, server_id)
+        client.await_registered(server_id)
 
         permitted_key = _key(client, resources, mcp_servers=[server_id])
         denied_key = _key(client, resources, mcp_servers=None)