test(e2e): client disconnect must not bench the Azure deployment it cancelled

Live proxy with cancel_on_disconnect, two-deployment group, generic allowed_fails=0, red at the pre-fix handler and green with the bare raise

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

tests/e2e/coverage_registry/reliability.yaml

@@ -11,6 +11,7 @@
 - {id: reliability.cooldown.5xx.trips_then_recovers, module: reliability, tier: P0, behavior: cooldown, variant: "5xx", assertions: [trips_then_recovers], exercised_on: [chat_completions], source: "cooldown_handlers.py:40", rationale: "Deployment cools after repeated 5xx, recovers after cooldown_time"}
 - {id: reliability.cooldown.429.trips_then_recovers, module: reliability, tier: P0, behavior: cooldown, variant: "429", assertions: [trips_then_recovers], exercised_on: [chat_completions], source: "cooldown_handlers.py:69", rationale: "Cools on 429, avoids hammering exhausted provider"}
 - {id: reliability.cooldown.auth.trips_then_recovers, module: reliability, tier: P1, behavior: cooldown, variant: auth, assertions: [trips_then_recovers], exercised_on: [chat_completions], source: "cooldown_handlers.py:74", rationale: "Cools on 401 auth error"}
+- {id: reliability.cooldown.client_disconnect.stays_healthy, module: reliability, tier: P0, behavior: cooldown, variant: client_disconnect, assertions: [stays_healthy], exercised_on: [chat_completions], source: "litellm/llms/azure/azure.py", fail_before_fix: proven, rationale: "With cancel_on_disconnect on, a client hanging up mid-request cancels the upstream call; that cancellation must not be recorded as a deployment 500 that benches a healthy Azure deployment"}
 - {id: reliability.cooldown.timeout.trips_then_recovers, module: reliability, tier: P1, behavior: cooldown, variant: timeout, assertions: [trips_then_recovers], exercised_on: [chat_completions], source: "cooldown_handlers.py:77", rationale: "Cools on 408 timeout"}
 - {id: reliability.routing.simple_shuffle.picks_healthy_deployment, module: reliability, tier: P1, behavior: routing, variant: simple_shuffle, assertions: [picks_healthy_deployment], exercised_on: [chat_completions], source: "router_strategy/simple_shuffle.py", rationale: "Baseline weighted/uniform pick"}
 - {id: reliability.routing.latency_based.picks_lowest_latency, module: reliability, tier: P1, behavior: routing, variant: latency_based, assertions: [picks_lowest_latency], exercised_on: [chat_completions], source: "router_strategy/lowest_latency.py", rationale: "Routes to lowest-latency deployment"}

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."""
@@ -877,7 +909,10 @@ class PreparedForward:
 
 
 def prepare_forward(
-    method: str, url: str, headers: dict[str, str], body: bytes | None,
+    method: str,
+    url: str,
+    headers: dict[str, str],
+    body: bytes | None,
 ) -> PreparedForward | NetworkError:
     try:
         with requests.Session() as session:
@@ -896,7 +931,8 @@ def forward_prepared_stream(prepared: PreparedForward, timeout: float) -> Stream
     except requests.RequestException as exc:
         return NetworkError(message=str(exc))
     return StreamHead(
-        resp.status_code, {name.lower(): value for name, value in resp.headers.items()},
+        resp.status_code,
+        {name.lower(): value for name, value in resp.headers.items()},
         primed_steps(_stream_steps(resp)),
     )
 

tests/e2e/gateway/stage_mirror_ci_config.yml

@@ -1,5 +1,6 @@
 general_settings:
   proxy_batch_write_at: 5
+  cancel_on_disconnect: true
   enable_jwt_auth: true
   litellm_jwtauth:
     user_id_jwt_field: sub

tests/e2e/models.py

@@ -1030,6 +1030,7 @@ class ModelInfoBody(BaseModel):
     mode: ModelMode | None = None
     access_groups: list[str] | None = None
     team_id: str | None = None
+    allowed_fails: int | None = None
     allowed_fails_policy: dict[str, int] | None = None
 
 

tests/e2e/router/reliability_support.py

@@ -17,9 +17,6 @@ from __future__ import annotations
 
 from collections.abc import Sequence
 
-from pydantic import ValidationError
-
-from proxy_client import ProxyClient
 from e2e_config import CHEAP_OPENAI_MODEL, PROXY_BASE_URL, unique_marker
 from e2e_http import NetworkError, StreamHead, StreamingResponse
 from models import (
@@ -35,6 +32,8 @@ from models import (
     TextContentPart,
     Usage,
 )
+from proxy_client import ProxyClient
+from pydantic import ValidationError
 
 REAL_MODEL = "openai/gpt-5.5"
 REAL_KEY = "os.environ/OPENAI_API_KEY"
@@ -120,6 +119,26 @@ def create_content_filtered_deployment(proxy: ProxyClient, name: str) -> str:
     )
 
 
+def create_azure_benched_on_first_failure_deployment(proxy: ProxyClient, name: str, cooldown_time: float) -> str:
+    """A healthy real Azure deployment benched on its first failure of any kind, so a cancellation the proxy
+    wrongly records as a 500 shows up as the next call landing on the backup."""
+    return proxy.register_model(
+        ModelNewBody(
+            model_name=name,
+            litellm_params=LiteLLMParamsBody(
+                model=CONTENT_FILTERED_MODEL,
+                api_key=AZURE_KEY,
+                api_base=AZURE_BASE,
+                api_version=AZURE_API_VERSION,
+                max_retries=0,
+                weight=1,
+                cooldown_time=cooldown_time,
+            ),
+            model_info=ModelInfoBody(allowed_fails=0),
+        ),
+    )
+
+
 def create_caching_deployment(proxy: ProxyClient, name: str) -> str:
     """Register the Anthropic deployment whose prompt cache the affinity check pins to."""
     return proxy.create_model(name, LiteLLMParamsBody(model=CACHING_MODEL, api_key=CACHING_KEY, weight=1))

tests/e2e/router/test_reliability_cancel_on_disconnect_e2e.py

@@ -0,0 +1,99 @@
+"""Live e2e: a client hanging up on an in-flight request must not bench the healthy
+deployment that was serving it.
+
+The proxy runs with `cancel_on_disconnect: true`, so when the client closes the
+socket before the answer arrives the proxy cancels the upstream call. That
+cancellation is the client's doing, so it must never count as a failure of the
+deployment: a deployment that benches on its very first failure of any kind has
+to keep serving the next request, and a request served right after the hang-up
+has to come from that same deployment rather than its zero-weight backup.
+
+The disconnect is real: a non-streaming /chat/completions asking the real Azure
+OpenAI deployment for a long generation, with the client closing the connection
+ABANDON_AFTER_SECONDS in, well before any answer. If Azure ever answers within
+that window the test fails loudly rather than passing without a disconnect.
+"""
+
+from __future__ import annotations
+
+import time
+
+import pytest
+from complexity_router_client import ComplexityRouterClient
+from e2e_config import unique_marker
+from e2e_http import AbandonedRequest
+from lifecycle import ResourceManager
+from models import ChatMessage, ReliabilityChatBody, RouterSettingsOverride
+from reliability_support import (
+    chat_override,
+    create_azure_benched_on_first_failure_deployment,
+    create_zero_weight_backup_deployment,
+    model_id_of,
+)
+
+pytestmark = pytest.mark.e2e
+
+ABANDON_AFTER_SECONDS = 2.0
+LONG_GENERATION_MAX_TOKENS = 4000
+FOLLOW_UP_CALLS = 3
+FOLLOW_UP_SPACING_SECONDS = 1.0
+COOLDOWN_SECONDS = 60.0
+
+
+def _long_generation_prompt(marker: str) -> str:
+    return (
+        f"Write a detailed, multi-chapter short story of at least 3000 words about {marker}. "
+        "Do not stop early and do not summarize."
+    )
+
+
+class TestReliabilityCancelOnDisconnect:
+    @pytest.mark.covers("reliability.cooldown.client_disconnect.stays_healthy")
+    def test_client_disconnect_does_not_bench_healthy_deployment(
+        self, client: ComplexityRouterClient, resources: ResourceManager, scoped_key: str
+    ) -> None:
+        group = f"reliability-cancel-on-disconnect-{unique_marker()}"
+        azure_deployment = create_azure_benched_on_first_failure_deployment(
+            client.proxy, group, cooldown_time=COOLDOWN_SECONDS
+        )
+        resources.defer(lambda: client.proxy.delete_model(azure_deployment))
+        backup_deployment = create_zero_weight_backup_deployment(client.proxy, group)
+        resources.defer(lambda: client.proxy.delete_model(backup_deployment))
+
+        abandoned = client.proxy.transport.abandon(
+            "/chat/completions",
+            headers=client.proxy.transport.bearer(scoped_key),
+            json=ReliabilityChatBody(
+                model=group,
+                messages=[ChatMessage(role="user", content=_long_generation_prompt(unique_marker()))],
+                max_tokens=LONG_GENERATION_MAX_TOKENS,
+                stream=False,
+                router_settings_override=RouterSettingsOverride(num_retries=0),
+                cache={"no-cache": True},
+            ),
+            after=ABANDON_AFTER_SECONDS,
+        )
+        assert isinstance(abandoned, AbandonedRequest), (
+            f"the proxy answered within {ABANDON_AFTER_SECONDS}s so the client never disconnected mid-request, "
+            f"got {abandoned.status_code}: {abandoned.body[:300]}"
+        )
+
+        for attempt in range(1, FOLLOW_UP_CALLS + 1):
+            time.sleep(FOLLOW_UP_SPACING_SECONDS)
+            resp = chat_override(
+                client.proxy,
+                scoped_key,
+                group,
+                f"say hi {unique_marker()}",
+                override=RouterSettingsOverride(num_retries=0),
+            )
+            assert resp.status_code == 200, (
+                f"follow-up {attempt}/{FOLLOW_UP_CALLS} should still land on the Azure deployment the client hung up on; "
+                f"landing on the backup means the cancellation was recorded as a deployment failure and benched it, "
+                f"got {resp.status_code}: {resp.body[:300]}"
+            )
+            assert model_id_of(resp) == azure_deployment, (
+                f"follow-up {attempt}/{FOLLOW_UP_CALLS} should still land on the Azure deployment the client hung up on; "
+                f"landing on the backup means the cancellation was recorded as a deployment failure and benched it, "
+                f"got {model_id_of(resp)!r}"
+            )

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)