scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

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

Browse source 
* test(e2e): harden harness and tests against data-plane pod churn

A stage autoscaler scale-down produced a 2s window of ALB 502s that killed six
budget tests on their first management call, and a freshly scaled-up pod that
had not run its 30s DB object sync yet failed two MCP tests and one prometheus
cardinality test. Retry transient gateway errors (502/503/504, connection
errors) once at the shared e2e_http dispatch seam, poll MCP server registration
to the poll deadline instead of asserting a single-shot listing, anchor the MCP
guardrail full-sync wait to the later of the guardrail and server writes, and
turn the prometheus alias poll into a drive-and-scrape convergence loop that
re-sends traffic for missing aliases and unions results across scrapes

* test(e2e): drain request body in retry stub handler so keep-alive reuse cannot misparse leftovers as requests

* revert(e2e): drop the transient-502 retry seam

A raw 502 during a pod scale-down is what a real client sees, so the suite
retrying past it hides an availability gap instead of flagging it. The
gateway-side fix is graceful drain on the deployment; until then the failures
are signal

* test(e2e): cap per-alias driver re-drives in the prometheus cardinality poll

Bounds worst-case provider spend to 4 completions per alias while scrapes keep
polling to the deadline; counters persist on whichever pod served them, so the
cap costs no convergence unless that pod dies

* test(e2e): drop driver re-drives from the prometheus cardinality poll

The per-key cardinality contract is process-local and counters persist on
whichever pod served the driver call, so unioning aliases across free scrape
polls converges without re-sending billable traffic. The residual gap, a pod
dying inside the poll window, is deferred to direct per-pod scraping
This commit is contained in:
ryan-crabbe-berri 2026-07-27 19:22:52 -07:00 • committed by GitHub
parent 328e41b1f9
commit daf22ec871
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 75 additions and 33 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
tests/e2e/logging/test_prometheus_cardinality_e2e.py
View file

@ -9,8 +9,11 @@ stamping ``api_key_alias`` (or collapses every key onto one series) would drop
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 @@ class TestPrometheusPerKeyCardinality:
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"
)

21
tests/e2e/mcp/mcp_client.py
View file

@ -195,6 +195,27 @@ class McpClient:
)
).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,
*,

7
tests/e2e/mcp/test_mcp_access_group_e2e.py
View file

@ -44,12 +44,7 @@ class TestMcpAccessGroupToolSelection:
)
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(), (

59
tests/e2e/mcp/test_mcp_guardrail_e2e.py
View file

@ -29,11 +29,11 @@ from mcp_client import McpCallToolResponse, McpClient, McpToolArguments
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 @@ class TestMcpToolCallGuardrail:
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 @@ class TestMcpToolCallGuardrail:
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, (

9
tests/e2e/mcp/test_mcp_key_access_e2e.py
View file

@ -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 @@ class TestMcpKeyWithoutAccessIsDenied:
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 @@ class TestMcpKeyWithoutAccessIsDenied:
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)