From e1afe2e29cee700710faa85063a9a0f7927104f6 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 28 Jul 2026 16:42:15 -0700 Subject: [PATCH 1/6] test(e2e): bound the post-/model/new servable wait at 40s _await_model_servable used poll_timeout (120s), the spend/log read-back budget. A stuck model reload therefore stalled every suite that creates a deployment for two minutes before failing Give create_model a fixed harness middle ground: model_servable_timeout=40s, polled every 2s, with each /v1/models call capped at 5s and clamped to the remaining deadline so one slow GET cannot overrun the wait. Happy path still returns on the first listing. Not derived from proxy general_settings or env Transport.get accepts an optional per-call timeout for that clamp. Unit tests cover the deadline arithmetic and clamp without a live proxy (cherry picked from commit c082a0e6488f50978bf5255f6b5298ba7e8fd8da) --- tests/e2e/proxy_client.py | 134 ++++++++++-- tests/e2e/test_proxy_client_model_servable.py | 190 ++++++++++++++++++ tests/e2e/transport.py | 13 +- 3 files changed, 313 insertions(+), 24 deletions(-) create mode 100644 tests/e2e/test_proxy_client_model_servable.py diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 6c6b948e29c..87693175e62 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -75,12 +75,95 @@ from transport import HttpTransport, SplitTransport, Transport RowsPredicate = Callable[[list[SpendLogRow]], bool] +# After /model/new, poll the data plane until the model is listed (or fail). +# Shorter than poll_timeout (spend/log read-backs ~120s); longer than a single +# request. 40s is the harness middle ground: happy path returns on the first +# poll, a stuck reload fails in under a minute instead of two. +MODEL_SERVABLE_TIMEOUT = 40.0 +MODEL_SERVABLE_INTERVAL = 2.0 +# Cap each /v1/models poll so one slow request cannot outlast the budget. +# Clamped further to remaining deadline inside await_servable. +MODEL_SERVABLE_REQUEST_TIMEOUT = 5.0 + + +@dataclass(frozen=True, slots=True) +class Servable: + """The data plane listed the model within the deadline.""" + + +@dataclass(frozen=True, slots=True) +class NotServable: + """The deadline passed without the data plane listing the model. + + `last_result` is the final /v1/models read, so the caller can tell "the proxy + answered but omitted the model" (propagation) from "the read itself failed" + (network/auth) when reporting.""" + + last_result: Result[ModelsListResponse] | None + + +ServableOutcome = Servable | NotServable + + +def await_servable( + list_models: Callable[[float], Result[ModelsListResponse]], + *, + model_name: str, + timeout: float, + interval: float, + request_timeout: float, + now: Callable[[], float], + sleep: Callable[[float], None], +) -> ServableOutcome: + """Poll `list_models` until the data plane lists `model_name` or `timeout` elapses. + + `list_models` receives the per-poll request timeout, clamped to the remaining + deadline so a slow final poll cannot overrun the overall budget. Clock and sleep + are injected so this is exercised without wall-clock waits. Always polls at least + once when the loop starts with a positive budget.""" + deadline = now() + timeout + last_result: Result[ModelsListResponse] | None = None + while True: + remaining = deadline - now() + if remaining <= 0 and last_result is not None: + return NotServable(last_result=last_result) + poll_timeout = min(request_timeout, remaining) if remaining > 0 else request_timeout + last_result = list_models(poll_timeout) + if isinstance(last_result, Success) and any( + entry.id == model_name for entry in last_result.data.data + ): + return Servable() + if now() + interval >= deadline: + return NotServable(last_result=last_result) + sleep(interval) + + +def servable_timeout_message( + *, + model_name: str, + timeout: float, + last_result: Result[ModelsListResponse] | None, +) -> str: + last_error = ( + f"; last /v1/models poll did not succeed: {last_result}" + if last_result is not None and not isinstance(last_result, Success) + else "" + ) + return ( + f"model {model_name!r} was created but never became servable on the data " + f"plane within {timeout}s of /model/new (control/data-plane propagation or " + f"STORE_MODEL_IN_DB reload issue){last_error}" + ) + @dataclass(frozen=True, slots=True) class ProxyClient: transport: Transport poll_timeout: float = 120.0 poll_interval: float = 5.0 + model_servable_timeout: float = MODEL_SERVABLE_TIMEOUT + model_servable_interval: float = MODEL_SERVABLE_INTERVAL + model_servable_request_timeout: float = MODEL_SERVABLE_REQUEST_TIMEOUT # ---- keys / customers (satisfies lifecycle.ResourceClient) ---------- @@ -167,7 +250,12 @@ class ProxyClient: this returns can race the reload and 400 with "Invalid model name passed". We therefore poll the data-plane /v1/models until the model appears before handing back, so callers can invoke it immediately. In the monolithic case - it is already present on the first poll, so this adds one request.""" + it is already present on the first poll, so this adds one request. + + The wait is bounded by `model_servable_timeout` rather than the much longer + `poll_timeout` used for batched read-backs, so a stuck reload fails in under + a minute instead of two. Happy path still returns as soon as /v1/models lists + the model (usually the first poll).""" model_id = unwrap( self.transport.post( "/model/new", @@ -185,32 +273,34 @@ class ProxyClient: def _await_model_servable(self, model_name: str) -> None: """Block until the data plane lists `model_name`, or fail loudly if it does - not within poll_timeout (a real propagation/config problem, surfaced here - instead of as a downstream "Invalid model name passed").""" - deadline = time.monotonic() + self.poll_timeout - last_result: Result[ModelsListResponse] | None = None - while time.monotonic() < deadline: - last_result = self.transport.get( + not within model_servable_timeout (a real propagation/config problem, + surfaced here instead of as a downstream "Invalid model name passed").""" + outcome = await_servable( + lambda poll_timeout: self.transport.get( "/v1/models", headers=self.transport.master, params=NoBody(), response_type=ModelsListResponse, - ) - if isinstance(last_result, Success) and any( - entry.id == model_name for entry in last_result.data.data - ): + timeout=poll_timeout, + ), + model_name=model_name, + timeout=self.model_servable_timeout, + interval=self.model_servable_interval, + request_timeout=self.model_servable_request_timeout, + now=time.monotonic, + sleep=time.sleep, + ) + match outcome: + case Servable(): return - time.sleep(self.poll_interval) - last_error = ( - f"; last /v1/models poll did not succeed: {last_result}" - if last_result is not None and not isinstance(last_result, Success) - else "" - ) - raise AssertionError( - f"model {model_name!r} was created but never became servable on the data " - f"plane within {self.poll_timeout}s of /model/new (control/data-plane " - f"propagation or STORE_MODEL_IN_DB reload issue){last_error}" - ) + case NotServable(last_result=last_result): + raise AssertionError( + servable_timeout_message( + model_name=model_name, + timeout=self.model_servable_timeout, + last_result=last_result, + ) + ) def update_model(self, model_id: str, litellm_params: LiteLLMParamsBody) -> None: """Merge `litellm_params` over the deployment `model_id`'s stored params via diff --git a/tests/e2e/test_proxy_client_model_servable.py b/tests/e2e/test_proxy_client_model_servable.py new file mode 100644 index 00000000000..0cc63f882e2 --- /dev/null +++ b/tests/e2e/test_proxy_client_model_servable.py @@ -0,0 +1,190 @@ +"""Harness coverage for the bounded wait after /model/new (no live proxy). + +Model propagation is polled to a deadline so a stuck control/data-plane reload fails +fast instead of stalling every test that creates a model. The clock and sleep are +injected, so these assert the deadline arithmetic without waiting. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +from e2e_http import NetworkError, Result, Success +from models import ModelListEntry, ModelsListResponse +from proxy_client import ( + MODEL_SERVABLE_REQUEST_TIMEOUT, + MODEL_SERVABLE_TIMEOUT, + NotServable, + Servable, + await_servable, + servable_timeout_message, +) + + +def _listing(*model_names: str) -> Result[ModelsListResponse]: + return Success( + status_code=200, + data=ModelsListResponse(data=tuple(ModelListEntry(id=name) for name in model_names)), + ) + + +@dataclass(slots=True) +class FakeClock: + """A clock that only advances when the code under test sleeps or a slow poll runs.""" + + seconds: float = 0.0 + slept: list[float] = field(default_factory=list) # mutable-ok: records calls for assertions + + def now(self) -> float: + return self.seconds + + def sleep(self, duration: float) -> None: + self.slept.append(duration) + self.seconds += duration + + +@dataclass(slots=True) +class FakeModelList: + """Returns each queued /v1/models read in turn, repeating the last forever.""" + + responses: tuple[Result[ModelsListResponse], ...] + calls: int = 0 + timeouts: list[float] = field(default_factory=list) # mutable-ok: records call timeouts + + def __call__(self, request_timeout: float) -> Result[ModelsListResponse]: + self.timeouts.append(request_timeout) + response = self.responses[min(self.calls, len(self.responses) - 1)] + self.calls += 1 + return response + + +def test_returns_servable_on_first_listing_without_sleeping() -> None: + clock = FakeClock() + list_models = FakeModelList(responses=(_listing("my-model"),)) + + outcome = await_servable( + list_models, + model_name="my-model", + timeout=40.0, + interval=2.0, + request_timeout=5.0, + now=clock.now, + sleep=clock.sleep, + ) + + assert outcome == Servable() + assert list_models.calls == 1 + assert list_models.timeouts == [5.0] + assert clock.slept == [] + + +def test_polls_until_the_model_appears() -> None: + clock = FakeClock() + list_models = FakeModelList(responses=(_listing("other"), _listing("other"), _listing("other", "my-model"))) + + outcome = await_servable( + list_models, + model_name="my-model", + timeout=40.0, + interval=2.0, + request_timeout=5.0, + now=clock.now, + sleep=clock.sleep, + ) + + assert outcome == Servable() + assert list_models.calls == 3 + assert clock.seconds == 4.0 + + +def test_gives_up_at_the_deadline_rather_than_polling_forever() -> None: + clock = FakeClock() + list_models = FakeModelList(responses=(_listing("other"),)) + + outcome = await_servable( + list_models, + model_name="my-model", + timeout=10.0, + interval=2.0, + request_timeout=5.0, + now=clock.now, + sleep=clock.sleep, + ) + + assert isinstance(outcome, NotServable) + assert clock.seconds == 8.0 + assert list_models.calls == 5 + assert list_models.timeouts == [5.0, 5.0, 5.0, 4.0, 2.0] + + +def test_does_not_wait_past_the_overall_budget() -> None: + clock = FakeClock() + + outcome = await_servable( + FakeModelList(responses=(_listing("other"),)), + model_name="my-model", + timeout=MODEL_SERVABLE_TIMEOUT, + interval=2.0, + request_timeout=MODEL_SERVABLE_REQUEST_TIMEOUT, + now=clock.now, + sleep=clock.sleep, + ) + + assert isinstance(outcome, NotServable) + assert clock.seconds <= MODEL_SERVABLE_TIMEOUT + + +def test_clamps_request_timeout_to_remaining_deadline() -> None: + """A slow final poll must not receive the full request cap when less budget remains. + + Without the clamp, remaining=3 and cap=5 lets the transport block for 5s and the + overall wait overruns model_servable_timeout by up to ~cap seconds. + """ + clock = FakeClock() + timeouts: list[float] = [] + + def list_models(request_timeout: float) -> Result[ModelsListResponse]: + timeouts.append(request_timeout) + clock.seconds += request_timeout + return _listing("other") + + outcome = await_servable( + list_models, + model_name="my-model", + timeout=10.0, + interval=2.0, + request_timeout=5.0, + now=clock.now, + sleep=clock.sleep, + ) + + assert isinstance(outcome, NotServable) + assert timeouts[0] == 5.0 + assert any(timeout < 5.0 for timeout in timeouts) + assert timeouts[-1] == 3.0 + assert clock.seconds <= 10.0 + + +def test_reports_a_failed_read_distinctly_from_a_missing_model() -> None: + clock = FakeClock() + unreachable: Result[ModelsListResponse] = NetworkError(message="connection refused") + + outcome = await_servable( + FakeModelList(responses=(unreachable,)), + model_name="my-model", + timeout=1.0, + interval=0.5, + request_timeout=5.0, + now=clock.now, + sleep=clock.sleep, + ) + + assert outcome == NotServable(last_result=unreachable) + message = servable_timeout_message(model_name="my-model", timeout=1.0, last_result=unreachable) + assert "connection refused" in message + + listed_without_model = _listing("other") + propagation_message = servable_timeout_message( + model_name="my-model", timeout=1.0, last_result=listed_without_model + ) + assert "did not succeed" not in propagation_message diff --git a/tests/e2e/transport.py b/tests/e2e/transport.py index a6adf83ed1f..27b11befc8e 100644 --- a/tests/e2e/transport.py +++ b/tests/e2e/transport.py @@ -58,6 +58,7 @@ class Transport(Protocol): headers: BaseModel, params: BaseModel, response_type: type[R], + timeout: float | None = None, ) -> Result[R]: ... def delete[R: BaseModel]( @@ -136,13 +137,16 @@ class HttpTransport: headers: BaseModel, params: BaseModel, response_type: type[R], + timeout: float | None = None, ) -> Result[R]: + """`timeout` overrides the transport-wide request_timeout for this call, for + pollers whose own deadline is shorter than it.""" return e2e_http.get( self._url(path), headers=headers, params=params, response_type=response_type, - timeout=self.request_timeout, + timeout=self.request_timeout if timeout is None else timeout, ) def delete[R: BaseModel]( @@ -336,9 +340,14 @@ class SplitTransport: headers: BaseModel, params: BaseModel, response_type: type[R], + timeout: float | None = None, ) -> Result[R]: return self._route(path).get( - path, headers=headers, params=params, response_type=response_type + path, + headers=headers, + params=params, + response_type=response_type, + timeout=timeout, ) def delete[R: BaseModel]( From 5aa66ea33e6c194f66d360e22292f22eddac38e7 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 28 Jul 2026 17:09:20 -0700 Subject: [PATCH 2/6] fix(e2e): wait one default DB reload interval of continuous listing create_model returned after the first /v1/models hit that listed the model, so chat could still land on a cold gateway worker (numWorkers>1 / peer pod) and 400 Invalid model name. Require continuous listing for the product default add_deployment interval (30s) after first sight so every worker has synced from the DB; first listing still bounded at 40s (cherry picked from commit 7d1ee2ff861b970f6de3f6759ff015947af9d2a1) --- tests/e2e/proxy_client.py | 85 ++++++++---- tests/e2e/test_proxy_client_model_servable.py | 125 ++++++++++++------ 2 files changed, 148 insertions(+), 62 deletions(-) diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 87693175e62..dfbb90ac08e 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -75,14 +75,17 @@ from transport import HttpTransport, SplitTransport, Transport RowsPredicate = Callable[[list[SpendLogRow]], bool] -# After /model/new, poll the data plane until the model is listed (or fail). -# Shorter than poll_timeout (spend/log read-backs ~120s); longer than a single -# request. 40s is the harness middle ground: happy path returns on the first -# poll, a stuck reload fails in under a minute instead of two. +# After /model/new, the control-plane writer reloads itself immediately, but every +# other gateway worker (and peer pod) only picks the model up on its add_deployment +# job. That job runs every proxy_config_reload_interval_seconds (product default 30). +# A single /v1/models hit can land on a hot worker while the next /chat hits a cold +# one ("Invalid model name"). Wait for first listing within MODEL_SERVABLE_TIMEOUT, +# then require continuous listing for MODEL_SERVABLE_DB_SYNC_SECONDS (the default +# reload interval) so every worker has had a chance to sync from the DB. MODEL_SERVABLE_TIMEOUT = 40.0 +MODEL_SERVABLE_DB_SYNC_SECONDS = 30.0 MODEL_SERVABLE_INTERVAL = 2.0 -# Cap each /v1/models poll so one slow request cannot outlast the budget. -# Clamped further to remaining deadline inside await_servable. +# Cap each /v1/models poll so one slow request cannot outlast the remaining budget. MODEL_SERVABLE_REQUEST_TIMEOUT = 5.0 @@ -112,29 +115,55 @@ def await_servable( timeout: float, interval: float, request_timeout: float, + db_sync_seconds: float, now: Callable[[], float], sleep: Callable[[float], None], ) -> ServableOutcome: - """Poll `list_models` until the data plane lists `model_name` or `timeout` elapses. + """Poll until `model_name` is listed long enough for every worker to DB-sync. - `list_models` receives the per-poll request timeout, clamped to the remaining - deadline so a slow final poll cannot overrun the overall budget. Clock and sleep - are injected so this is exercised without wall-clock waits. Always polls at least - once when the loop starts with a positive budget.""" - deadline = now() + timeout + First listing must happen within `timeout`. After that, the model must stay + listed continuously for `db_sync_seconds` (any miss resets the continuous + window). `db_sync_seconds=0` returns on the first listing. Each poll's request + timeout is clamped to the remaining budget. Clock and sleep are injected.""" + started = now() + first_seen_at: float | None = None last_result: Result[ModelsListResponse] | None = None while True: - remaining = deadline - now() + t = now() + if first_seen_at is None: + deadline = started + timeout + else: + deadline = first_seen_at + db_sync_seconds + remaining = deadline - t if remaining <= 0 and last_result is not None: + if first_seen_at is not None and db_sync_seconds <= 0: + return Servable() + if first_seen_at is not None and t - first_seen_at >= db_sync_seconds: + return Servable() return NotServable(last_result=last_result) poll_timeout = min(request_timeout, remaining) if remaining > 0 else request_timeout last_result = list_models(poll_timeout) - if isinstance(last_result, Success) and any( + listed = isinstance(last_result, Success) and any( entry.id == model_name for entry in last_result.data.data - ): + ) + t = now() + if not listed: + first_seen_at = None + elif first_seen_at is None: + first_seen_at = t + if db_sync_seconds <= 0: + return Servable() + elif t - first_seen_at >= db_sync_seconds: return Servable() - if now() + interval >= deadline: - return NotServable(last_result=last_result) + if first_seen_at is None: + if now() + interval >= started + timeout: + return NotServable(last_result=last_result) + elif now() + interval >= first_seen_at + db_sync_seconds: + # Final stretch: sleep only the remainder of the continuous window. + remainder = first_seen_at + db_sync_seconds - now() + if remainder > 0: + sleep(remainder) + continue sleep(interval) @@ -142,6 +171,7 @@ def servable_timeout_message( *, model_name: str, timeout: float, + db_sync_seconds: float, last_result: Result[ModelsListResponse] | None, ) -> str: last_error = ( @@ -151,7 +181,8 @@ def servable_timeout_message( ) return ( f"model {model_name!r} was created but never became servable on the data " - f"plane within {timeout}s of /model/new (control/data-plane propagation or " + f"plane within {timeout}s of first listing (plus {db_sync_seconds}s continuous " + f"DB sync) after /model/new (control/data-plane propagation or " f"STORE_MODEL_IN_DB reload issue){last_error}" ) @@ -162,6 +193,7 @@ class ProxyClient: poll_timeout: float = 120.0 poll_interval: float = 5.0 model_servable_timeout: float = MODEL_SERVABLE_TIMEOUT + model_servable_db_sync_seconds: float = MODEL_SERVABLE_DB_SYNC_SECONDS model_servable_interval: float = MODEL_SERVABLE_INTERVAL model_servable_request_timeout: float = MODEL_SERVABLE_REQUEST_TIMEOUT @@ -252,10 +284,10 @@ class ProxyClient: handing back, so callers can invoke it immediately. In the monolithic case it is already present on the first poll, so this adds one request. - The wait is bounded by `model_servable_timeout` rather than the much longer - `poll_timeout` used for batched read-backs, so a stuck reload fails in under - a minute instead of two. Happy path still returns as soon as /v1/models lists - the model (usually the first poll).""" + First listing must arrive within `model_servable_timeout` (not the longer + spend `poll_timeout`). The model must then stay listed for + `model_servable_db_sync_seconds` (product default DB reload interval) so every + gateway worker has run add_deployment before callers use the model.""" model_id = unwrap( self.transport.post( "/model/new", @@ -272,9 +304,10 @@ class ProxyClient: return model_id def _await_model_servable(self, model_name: str) -> None: - """Block until the data plane lists `model_name`, or fail loudly if it does - not within model_servable_timeout (a real propagation/config problem, - surfaced here instead of as a downstream "Invalid model name passed").""" + """Block until the data plane lists `model_name` long enough for DB sync. + + Fails if first listing misses model_servable_timeout, or if continuous listing + for model_servable_db_sync_seconds never holds (multi-worker / peer reload).""" outcome = await_servable( lambda poll_timeout: self.transport.get( "/v1/models", @@ -287,6 +320,7 @@ class ProxyClient: timeout=self.model_servable_timeout, interval=self.model_servable_interval, request_timeout=self.model_servable_request_timeout, + db_sync_seconds=self.model_servable_db_sync_seconds, now=time.monotonic, sleep=time.sleep, ) @@ -298,6 +332,7 @@ class ProxyClient: servable_timeout_message( model_name=model_name, timeout=self.model_servable_timeout, + db_sync_seconds=self.model_servable_db_sync_seconds, last_result=last_result, ) ) diff --git a/tests/e2e/test_proxy_client_model_servable.py b/tests/e2e/test_proxy_client_model_servable.py index 0cc63f882e2..c9edd148c2b 100644 --- a/tests/e2e/test_proxy_client_model_servable.py +++ b/tests/e2e/test_proxy_client_model_servable.py @@ -1,8 +1,9 @@ """Harness coverage for the bounded wait after /model/new (no live proxy). -Model propagation is polled to a deadline so a stuck control/data-plane reload fails -fast instead of stalling every test that creates a model. The clock and sleep are -injected, so these assert the deadline arithmetic without waiting. +create_model must wait for the product default DB reload interval of continuous +listing so multi-worker gateways finish add_deployment before callers use the +model. Clock and sleep are injected so these assert the deadline arithmetic +without wall-clock waits. """ from __future__ import annotations @@ -12,6 +13,7 @@ from dataclasses import dataclass, field from e2e_http import NetworkError, Result, Success from models import ModelListEntry, ModelsListResponse from proxy_client import ( + MODEL_SERVABLE_DB_SYNC_SECONDS, MODEL_SERVABLE_REQUEST_TIMEOUT, MODEL_SERVABLE_TIMEOUT, NotServable, @@ -30,8 +32,6 @@ def _listing(*model_names: str) -> Result[ModelsListResponse]: @dataclass(slots=True) class FakeClock: - """A clock that only advances when the code under test sleeps or a slow poll runs.""" - seconds: float = 0.0 slept: list[float] = field(default_factory=list) # mutable-ok: records calls for assertions @@ -45,8 +45,6 @@ class FakeClock: @dataclass(slots=True) class FakeModelList: - """Returns each queued /v1/models read in turn, repeating the last forever.""" - responses: tuple[Result[ModelsListResponse], ...] calls: int = 0 timeouts: list[float] = field(default_factory=list) # mutable-ok: records call timeouts @@ -58,7 +56,7 @@ class FakeModelList: return response -def test_returns_servable_on_first_listing_without_sleeping() -> None: +def test_returns_on_first_listing_when_db_sync_is_zero() -> None: clock = FakeClock() list_models = FakeModelList(responses=(_listing("my-model"),)) @@ -68,17 +66,64 @@ def test_returns_servable_on_first_listing_without_sleeping() -> None: timeout=40.0, interval=2.0, request_timeout=5.0, + db_sync_seconds=0.0, now=clock.now, sleep=clock.sleep, ) assert outcome == Servable() assert list_models.calls == 1 - assert list_models.timeouts == [5.0] assert clock.slept == [] -def test_polls_until_the_model_appears() -> None: +def test_requires_continuous_listing_for_default_db_sync_interval() -> None: + clock = FakeClock() + list_models = FakeModelList(responses=(_listing("my-model"),)) + + outcome = await_servable( + list_models, + model_name="my-model", + timeout=40.0, + interval=2.0, + request_timeout=5.0, + db_sync_seconds=MODEL_SERVABLE_DB_SYNC_SECONDS, + now=clock.now, + sleep=clock.sleep, + ) + + assert outcome == Servable() + assert clock.seconds >= MODEL_SERVABLE_DB_SYNC_SECONDS + assert list_models.calls >= 2 + + +def test_resets_db_sync_window_when_a_poll_misses() -> None: + clock = FakeClock() + list_models = FakeModelList( + responses=( + _listing("my-model"), + _listing("my-model"), + _listing("other"), + _listing("my-model"), + ) + ) + + outcome = await_servable( + list_models, + model_name="my-model", + timeout=40.0, + interval=2.0, + request_timeout=5.0, + db_sync_seconds=6.0, + now=clock.now, + sleep=clock.sleep, + ) + + assert outcome == Servable() + assert list_models.calls >= 4 + assert clock.seconds >= 6.0 + + +def test_polls_until_the_model_first_appears() -> None: clock = FakeClock() list_models = FakeModelList(responses=(_listing("other"), _listing("other"), _listing("other", "my-model"))) @@ -88,6 +133,7 @@ def test_polls_until_the_model_appears() -> None: timeout=40.0, interval=2.0, request_timeout=5.0, + db_sync_seconds=0.0, now=clock.now, sleep=clock.sleep, ) @@ -97,7 +143,7 @@ def test_polls_until_the_model_appears() -> None: assert clock.seconds == 4.0 -def test_gives_up_at_the_deadline_rather_than_polling_forever() -> None: +def test_gives_up_if_first_listing_never_arrives() -> None: clock = FakeClock() list_models = FakeModelList(responses=(_listing("other"),)) @@ -107,6 +153,7 @@ def test_gives_up_at_the_deadline_rather_than_polling_forever() -> None: timeout=10.0, interval=2.0, request_timeout=5.0, + db_sync_seconds=30.0, now=clock.now, sleep=clock.sleep, ) @@ -114,32 +161,9 @@ def test_gives_up_at_the_deadline_rather_than_polling_forever() -> None: assert isinstance(outcome, NotServable) assert clock.seconds == 8.0 assert list_models.calls == 5 - assert list_models.timeouts == [5.0, 5.0, 5.0, 4.0, 2.0] - - -def test_does_not_wait_past_the_overall_budget() -> None: - clock = FakeClock() - - outcome = await_servable( - FakeModelList(responses=(_listing("other"),)), - model_name="my-model", - timeout=MODEL_SERVABLE_TIMEOUT, - interval=2.0, - request_timeout=MODEL_SERVABLE_REQUEST_TIMEOUT, - now=clock.now, - sleep=clock.sleep, - ) - - assert isinstance(outcome, NotServable) - assert clock.seconds <= MODEL_SERVABLE_TIMEOUT def test_clamps_request_timeout_to_remaining_deadline() -> None: - """A slow final poll must not receive the full request cap when less budget remains. - - Without the clamp, remaining=3 and cap=5 lets the transport block for 5s and the - overall wait overruns model_servable_timeout by up to ~cap seconds. - """ clock = FakeClock() timeouts: list[float] = [] @@ -154,6 +178,7 @@ def test_clamps_request_timeout_to_remaining_deadline() -> None: timeout=10.0, interval=2.0, request_timeout=5.0, + db_sync_seconds=30.0, now=clock.now, sleep=clock.sleep, ) @@ -161,10 +186,27 @@ def test_clamps_request_timeout_to_remaining_deadline() -> None: assert isinstance(outcome, NotServable) assert timeouts[0] == 5.0 assert any(timeout < 5.0 for timeout in timeouts) - assert timeouts[-1] == 3.0 assert clock.seconds <= 10.0 +def test_does_not_wait_past_first_listing_budget_when_missing() -> None: + clock = FakeClock() + + outcome = await_servable( + FakeModelList(responses=(_listing("other"),)), + model_name="my-model", + timeout=MODEL_SERVABLE_TIMEOUT, + interval=2.0, + request_timeout=MODEL_SERVABLE_REQUEST_TIMEOUT, + db_sync_seconds=MODEL_SERVABLE_DB_SYNC_SECONDS, + now=clock.now, + sleep=clock.sleep, + ) + + assert isinstance(outcome, NotServable) + assert clock.seconds <= MODEL_SERVABLE_TIMEOUT + + def test_reports_a_failed_read_distinctly_from_a_missing_model() -> None: clock = FakeClock() unreachable: Result[ModelsListResponse] = NetworkError(message="connection refused") @@ -175,16 +217,25 @@ def test_reports_a_failed_read_distinctly_from_a_missing_model() -> None: timeout=1.0, interval=0.5, request_timeout=5.0, + db_sync_seconds=0.0, now=clock.now, sleep=clock.sleep, ) assert outcome == NotServable(last_result=unreachable) - message = servable_timeout_message(model_name="my-model", timeout=1.0, last_result=unreachable) + message = servable_timeout_message( + model_name="my-model", + timeout=1.0, + db_sync_seconds=0.0, + last_result=unreachable, + ) assert "connection refused" in message listed_without_model = _listing("other") propagation_message = servable_timeout_message( - model_name="my-model", timeout=1.0, last_result=listed_without_model + model_name="my-model", + timeout=1.0, + db_sync_seconds=30.0, + last_result=listed_without_model, ) assert "did not succeed" not in propagation_message From 5953a66eab8212beac24d3265cd118b7f4a51176 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 28 Jul 2026 17:10:04 -0700 Subject: [PATCH 3/6] test(e2e): drop proxy_client model-servable unit tests Keep the create_model DB-sync wait in the harness; the pure-function unit file is not needed for this PR (cherry picked from commit 89204651d1a4537c6f21550c5ab85448ae0923f8) --- tests/e2e/test_proxy_client_model_servable.py | 241 ------------------ 1 file changed, 241 deletions(-) delete mode 100644 tests/e2e/test_proxy_client_model_servable.py diff --git a/tests/e2e/test_proxy_client_model_servable.py b/tests/e2e/test_proxy_client_model_servable.py deleted file mode 100644 index c9edd148c2b..00000000000 --- a/tests/e2e/test_proxy_client_model_servable.py +++ /dev/null @@ -1,241 +0,0 @@ -"""Harness coverage for the bounded wait after /model/new (no live proxy). - -create_model must wait for the product default DB reload interval of continuous -listing so multi-worker gateways finish add_deployment before callers use the -model. Clock and sleep are injected so these assert the deadline arithmetic -without wall-clock waits. -""" - -from __future__ import annotations - -from dataclasses import dataclass, field - -from e2e_http import NetworkError, Result, Success -from models import ModelListEntry, ModelsListResponse -from proxy_client import ( - MODEL_SERVABLE_DB_SYNC_SECONDS, - MODEL_SERVABLE_REQUEST_TIMEOUT, - MODEL_SERVABLE_TIMEOUT, - NotServable, - Servable, - await_servable, - servable_timeout_message, -) - - -def _listing(*model_names: str) -> Result[ModelsListResponse]: - return Success( - status_code=200, - data=ModelsListResponse(data=tuple(ModelListEntry(id=name) for name in model_names)), - ) - - -@dataclass(slots=True) -class FakeClock: - seconds: float = 0.0 - slept: list[float] = field(default_factory=list) # mutable-ok: records calls for assertions - - def now(self) -> float: - return self.seconds - - def sleep(self, duration: float) -> None: - self.slept.append(duration) - self.seconds += duration - - -@dataclass(slots=True) -class FakeModelList: - responses: tuple[Result[ModelsListResponse], ...] - calls: int = 0 - timeouts: list[float] = field(default_factory=list) # mutable-ok: records call timeouts - - def __call__(self, request_timeout: float) -> Result[ModelsListResponse]: - self.timeouts.append(request_timeout) - response = self.responses[min(self.calls, len(self.responses) - 1)] - self.calls += 1 - return response - - -def test_returns_on_first_listing_when_db_sync_is_zero() -> None: - clock = FakeClock() - list_models = FakeModelList(responses=(_listing("my-model"),)) - - outcome = await_servable( - list_models, - model_name="my-model", - timeout=40.0, - interval=2.0, - request_timeout=5.0, - db_sync_seconds=0.0, - now=clock.now, - sleep=clock.sleep, - ) - - assert outcome == Servable() - assert list_models.calls == 1 - assert clock.slept == [] - - -def test_requires_continuous_listing_for_default_db_sync_interval() -> None: - clock = FakeClock() - list_models = FakeModelList(responses=(_listing("my-model"),)) - - outcome = await_servable( - list_models, - model_name="my-model", - timeout=40.0, - interval=2.0, - request_timeout=5.0, - db_sync_seconds=MODEL_SERVABLE_DB_SYNC_SECONDS, - now=clock.now, - sleep=clock.sleep, - ) - - assert outcome == Servable() - assert clock.seconds >= MODEL_SERVABLE_DB_SYNC_SECONDS - assert list_models.calls >= 2 - - -def test_resets_db_sync_window_when_a_poll_misses() -> None: - clock = FakeClock() - list_models = FakeModelList( - responses=( - _listing("my-model"), - _listing("my-model"), - _listing("other"), - _listing("my-model"), - ) - ) - - outcome = await_servable( - list_models, - model_name="my-model", - timeout=40.0, - interval=2.0, - request_timeout=5.0, - db_sync_seconds=6.0, - now=clock.now, - sleep=clock.sleep, - ) - - assert outcome == Servable() - assert list_models.calls >= 4 - assert clock.seconds >= 6.0 - - -def test_polls_until_the_model_first_appears() -> None: - clock = FakeClock() - list_models = FakeModelList(responses=(_listing("other"), _listing("other"), _listing("other", "my-model"))) - - outcome = await_servable( - list_models, - model_name="my-model", - timeout=40.0, - interval=2.0, - request_timeout=5.0, - db_sync_seconds=0.0, - now=clock.now, - sleep=clock.sleep, - ) - - assert outcome == Servable() - assert list_models.calls == 3 - assert clock.seconds == 4.0 - - -def test_gives_up_if_first_listing_never_arrives() -> None: - clock = FakeClock() - list_models = FakeModelList(responses=(_listing("other"),)) - - outcome = await_servable( - list_models, - model_name="my-model", - timeout=10.0, - interval=2.0, - request_timeout=5.0, - db_sync_seconds=30.0, - now=clock.now, - sleep=clock.sleep, - ) - - assert isinstance(outcome, NotServable) - assert clock.seconds == 8.0 - assert list_models.calls == 5 - - -def test_clamps_request_timeout_to_remaining_deadline() -> None: - clock = FakeClock() - timeouts: list[float] = [] - - def list_models(request_timeout: float) -> Result[ModelsListResponse]: - timeouts.append(request_timeout) - clock.seconds += request_timeout - return _listing("other") - - outcome = await_servable( - list_models, - model_name="my-model", - timeout=10.0, - interval=2.0, - request_timeout=5.0, - db_sync_seconds=30.0, - now=clock.now, - sleep=clock.sleep, - ) - - assert isinstance(outcome, NotServable) - assert timeouts[0] == 5.0 - assert any(timeout < 5.0 for timeout in timeouts) - assert clock.seconds <= 10.0 - - -def test_does_not_wait_past_first_listing_budget_when_missing() -> None: - clock = FakeClock() - - outcome = await_servable( - FakeModelList(responses=(_listing("other"),)), - model_name="my-model", - timeout=MODEL_SERVABLE_TIMEOUT, - interval=2.0, - request_timeout=MODEL_SERVABLE_REQUEST_TIMEOUT, - db_sync_seconds=MODEL_SERVABLE_DB_SYNC_SECONDS, - now=clock.now, - sleep=clock.sleep, - ) - - assert isinstance(outcome, NotServable) - assert clock.seconds <= MODEL_SERVABLE_TIMEOUT - - -def test_reports_a_failed_read_distinctly_from_a_missing_model() -> None: - clock = FakeClock() - unreachable: Result[ModelsListResponse] = NetworkError(message="connection refused") - - outcome = await_servable( - FakeModelList(responses=(unreachable,)), - model_name="my-model", - timeout=1.0, - interval=0.5, - request_timeout=5.0, - db_sync_seconds=0.0, - now=clock.now, - sleep=clock.sleep, - ) - - assert outcome == NotServable(last_result=unreachable) - message = servable_timeout_message( - model_name="my-model", - timeout=1.0, - db_sync_seconds=0.0, - last_result=unreachable, - ) - assert "connection refused" in message - - listed_without_model = _listing("other") - propagation_message = servable_timeout_message( - model_name="my-model", - timeout=1.0, - db_sync_seconds=30.0, - last_result=listed_without_model, - ) - assert "did not succeed" not in propagation_message From 38d03fd341bf5b7030b2ebbc45c2555a7f1ee1e6 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 28 Jul 2026 17:16:30 -0700 Subject: [PATCH 4/6] fix(e2e): never skip the final deadline-clamped model-servable poll When less than one full poll interval remained in the first-listing budget, the pre-sleep check returned NotServable without another /v1/models call. Sleep only min(interval, time left) so a model that becomes listable in the last seconds of the timeout still gets a clamped final poll (cherry picked from commit 8439195922c913d118cb146409c75cc081d23e6e) --- tests/e2e/proxy_client.py | 43 ++++++++++++++++++++------------------- 1 file changed, 22 insertions(+), 21 deletions(-) diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index dfbb90ac08e..36dada5770a 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -124,24 +124,28 @@ def await_servable( First listing must happen within `timeout`. After that, the model must stay listed continuously for `db_sync_seconds` (any miss resets the continuous window). `db_sync_seconds=0` returns on the first listing. Each poll's request - timeout is clamped to the remaining budget. Clock and sleep are injected.""" + timeout is clamped to the remaining budget. Sleeps only min(interval, time left) + so a final deadline-clamped poll is never skipped just because a full interval + does not fit. Clock and sleep are injected.""" started = now() first_seen_at: float | None = None last_result: Result[ModelsListResponse] | None = None while True: t = now() - if first_seen_at is None: - deadline = started + timeout - else: - deadline = first_seen_at + db_sync_seconds - remaining = deadline - t - if remaining <= 0 and last_result is not None: - if first_seen_at is not None and db_sync_seconds <= 0: - return Servable() - if first_seen_at is not None and t - first_seen_at >= db_sync_seconds: + phase_deadline = ( + started + timeout if first_seen_at is None else first_seen_at + db_sync_seconds + ) + remaining = phase_deadline - t + if remaining <= 0: + if ( + last_result is not None + and first_seen_at is not None + and (db_sync_seconds <= 0 or t - first_seen_at >= db_sync_seconds) + ): return Servable() return NotServable(last_result=last_result) - poll_timeout = min(request_timeout, remaining) if remaining > 0 else request_timeout + + poll_timeout = min(request_timeout, remaining) last_result = list_models(poll_timeout) listed = isinstance(last_result, Success) and any( entry.id == model_name for entry in last_result.data.data @@ -155,16 +159,13 @@ def await_servable( return Servable() elif t - first_seen_at >= db_sync_seconds: return Servable() - if first_seen_at is None: - if now() + interval >= started + timeout: - return NotServable(last_result=last_result) - elif now() + interval >= first_seen_at + db_sync_seconds: - # Final stretch: sleep only the remainder of the continuous window. - remainder = first_seen_at + db_sync_seconds - now() - if remainder > 0: - sleep(remainder) - continue - sleep(interval) + + phase_deadline = ( + started + timeout if first_seen_at is None else first_seen_at + db_sync_seconds + ) + wait = min(interval, phase_deadline - now()) + if wait > 0: + sleep(wait) def servable_timeout_message( From 87be33f9354f399100c7ededdfba31b0e831d225 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 28 Jul 2026 17:35:13 -0700 Subject: [PATCH 5/6] fix(e2e): reject first listing that returns after the 40s deadline A poll may start with remaining budget and still return after started+timeout if the transport overruns its clamp. Recheck the first-listing deadline after the response so a late listing does not open the continuous DB-sync phase (cherry picked from commit 7ff2bcbf1498ee82f7dbe0b4330c1ab48927ed01) --- tests/e2e/proxy_client.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 36dada5770a..b3fc8538322 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -154,6 +154,8 @@ def await_servable( if not listed: first_seen_at = None elif first_seen_at is None: + if t > started + timeout: + return NotServable(last_result=last_result) first_seen_at = t if db_sync_seconds <= 0: return Servable() From 82fa66908bb72748efc574a8a1a8750155a85c14 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 28 Jul 2026 22:15:21 -0700 Subject: [PATCH 6/6] test(e2e): poll MCP tools across multi-worker lag (#35047) * fix(mcp): resolve call_tool by registry without requiring tool map Multi-worker reloads put MCP servers in the registry from the DB but do not re-run tools/list on every process. Gating call_tool on tool_name_to_mcp_server_name_mapping made cold workers 500 with Tool not found after another worker had already listed the tool. Treat a registry match on server id/name/alias as enough; upstream rejects unknown tools * test(e2e): poll MCP register, tools/list, and tools/call across multi-worker lag Stage multi-worker gateways only load MCP servers and tool maps on the process that handled the request. Poll until the server is listed, the tool appears on tools/list, and tools/call is not a cold-worker 500 so key-access and Datadog MCP e2e stop racing the LB * Revert "fix(mcp): resolve call_tool by registry without requiring tool map" This reverts commit 8b56e51e39b876d13d1112efa4130554ddf5f173. * test(e2e): tighten MCP multi-worker lag classifier Only retry tools/call on gateway shapes Tool not found and server_not_found, not any 500 that mentions tool/server not found, so upstream failures are not retried until the poll deadline * test(e2e): drop unit file for MCP lag classifier The live await_call_tool polls already cover multi-worker lag; a separate string-match unit module is not worth keeping (cherry picked from 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, 111 insertions(+), 24 deletions(-) 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}"