diff --git a/tests/e2e/e2e_http.py b/tests/e2e/e2e_http.py index d4978601b20..58817d399a8 100644 --- a/tests/e2e/e2e_http.py +++ b/tests/e2e/e2e_http.py @@ -676,6 +676,38 @@ def send( return streaming_outcome(resp, stream, sent_at=sent_at) +class AbandonedRequest(BaseModel): + """A non-streaming request the client walked away from: the socket was closed + ``after`` seconds in, before the proxy had answered, so the proxy saw a client + disconnect with the upstream call still in flight.""" + + kind: Literal["abandoned"] = "abandoned" + after: float + + +def abandon( + url: URL, *, headers: BaseModel, json: BaseModel, after: float, connect_timeout: float = 10.0 +) -> AbandonedRequest | StreamingResponse: + """POST and hang up ``after`` seconds if no response head has arrived by then, + closing the connection so the proxy observes the disconnect. Returns the + response instead when the proxy answered first, so a test can tell a real + disconnect from a generation that finished too fast to be cancelled.""" + sent_at: Final = time.monotonic() + session: Final = requests.Session() + try: + resp = session.post( + str(url), + headers=_headers(headers), + json=wire_body(json), + timeout=(connect_timeout, after), + ) + except requests.exceptions.ReadTimeout: + return AbandonedRequest(after=after) + finally: + session.close() + return streaming_outcome(resp, False, sent_at=sent_at) + + def stream(url: URL, *, headers: BaseModel, json: BaseModel, timeout: float = 60.0) -> StreamingResponse: """Streaming (SSE) call: consumes the stream counting events, and captures the x-litellm-call-id + content-type headers. Body is elided.""" diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 11a58742058..a0f1fa47c67 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -1047,8 +1047,8 @@ class ModelInfoBody(BaseModel): mode: ModelMode | None = None access_groups: list[str] | None = None team_id: str | None = None - allowed_fails_policy: dict[str, int] | None = None allowed_fails: int | None = None + allowed_fails_policy: dict[str, int] | None = None class ModelNewBody(BaseModel): diff --git a/tests/e2e/router/test_reliability_cancel_on_disconnect_e2e.py b/tests/e2e/router/test_reliability_cancel_on_disconnect_e2e.py index 21fd2d70603..7d46a980034 100644 --- a/tests/e2e/router/test_reliability_cancel_on_disconnect_e2e.py +++ b/tests/e2e/router/test_reliability_cancel_on_disconnect_e2e.py @@ -34,9 +34,9 @@ import time import pytest from complexity_router_client import ComplexityRouterClient from e2e_config import unique_marker -from e2e_http import NetworkError, StreamingResponse +from e2e_http import AbandonedRequest, StreamingResponse from lifecycle import ResourceManager -from models import ChatMessage, ChatResponse, ReliabilityChatBody, RouterSettingsOverride +from models import ChatMessage, ReliabilityChatBody, RouterSettingsOverride from reliability_support import ( chat_override, create_azure_benched_on_first_failure_deployment, @@ -65,8 +65,8 @@ def _say_hi(client: ComplexityRouterClient, key: str, group: str) -> StreamingRe def _hang_up_mid_answer(client: ComplexityRouterClient, key: str, group: str) -> None: """Send a request whose answer takes far longer than the client waits, so the - read timeout closes the socket while the provider is still generating.""" - outcome = client.proxy.transport.post( + client closes the socket while the provider is still generating.""" + outcome = client.proxy.transport.abandon( "/chat/completions", headers=client.proxy.transport.bearer(key), json=ReliabilityChatBody( @@ -80,16 +80,15 @@ def _hang_up_mid_answer(client: ComplexityRouterClient, key: str, group: str) -> max_tokens=LONG_ANSWER_MAX_TOKENS, router_settings_override=RouterSettingsOverride(num_retries=0), ), - response_type=ChatResponse, - timeout=CLIENT_HANGS_UP_AFTER_SECONDS, + after=CLIENT_HANGS_UP_AFTER_SECONDS, ) match outcome: - case NetworkError(message=message) if "Read timed out" in message: + case AbandonedRequest(): return - case _: + case StreamingResponse(status_code=status_code, body=body): pytest.fail( f"the client should have hung up {CLIENT_HANGS_UP_AFTER_SECONDS:.0f}s into a long answer with the " - f"call still in flight, but the proxy answered first: {outcome!r}" + f"call still in flight, but the proxy answered first with {status_code}: {body[:300]}" ) diff --git a/tests/e2e/transport.py b/tests/e2e/transport.py index 0022c0c4355..a3eec815441 100644 --- a/tests/e2e/transport.py +++ b/tests/e2e/transport.py @@ -13,6 +13,7 @@ from typing import Protocol import e2e_http from e2e_http import ( URL, + AbandonedRequest, AuthHeaders, BinaryStream, NetworkError, @@ -58,6 +59,10 @@ class Transport(Protocol): stream: bool = False, ) -> StreamingResponse: ... + def abandon( + self, path: str, *, headers: BaseModel, json: BaseModel, after: float + ) -> AbandonedRequest | StreamingResponse: ... + def get[R: BaseModel]( self, path: str, @@ -243,6 +248,11 @@ class HttpTransport: timeout=self.request_timeout, ) + def abandon( + self, path: str, *, headers: BaseModel, json: BaseModel, after: float + ) -> AbandonedRequest | StreamingResponse: + return e2e_http.abandon(self._url(path), headers=headers, json=json, after=after) + def probe(self, path: str, *, params: BaseModel, headers: BaseModel | None = None) -> ProbeResult: return e2e_http.probe( self._url(path), @@ -420,6 +430,11 @@ class SplitTransport: ) -> StreamingResponse: return self._route(path).send(path, headers=headers, json=json, params=params, stream=stream) + def abandon( + self, path: str, *, headers: BaseModel, json: BaseModel, after: float + ) -> AbandonedRequest | StreamingResponse: + return self._route(path).abandon(path, headers=headers, json=json, after=after) + def probe(self, path: str, *, params: BaseModel, headers: BaseModel | None = None) -> ProbeResult: return self._route(path).probe(path, params=params, headers=headers)