Merge pull request #42295 from BerriAI/litellm_fix_azure_cancellederror_cooldown

fix(azure): propagate asyncio.CancelledError instead of raising a 500

litellm/llms/azure/azure.py

@@ -481,7 +481,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
                 additional_args={"complete_input_dict": data},
                 original_response=str(e),
             )
-            raise AzureOpenAIError(status_code=500, message=str(e))
+            raise
         except Exception as e:
             message: Final = getattr(e, "message", str(e))
             body: Final = getattr(e, "body", None)

tests/e2e/coverage_registry/reliability.yaml

@@ -12,6 +12,7 @@
 - {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.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.cooldown.client_disconnect.stays_healthy, module: reliability, tier: P0, behavior: cooldown, variant: client_disconnect, assertions: [stays_healthy], exercised_on: [chat_completions], source: "llms/azure/azure.py:484", fail_before_fix: proven, rationale: "A client hanging up mid-request under cancel_on_disconnect never benches the Azure deployment it was talking to: the cancellation used to surface as a fake 500 that tripped the cooldown and sent every caller behind it to billed fallbacks (GitHub issues #35329 and #42222)"}
 - {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"}
 - {id: reliability.routing.cost_based.picks_lowest_cost, module: reliability, tier: P1, behavior: routing, variant: cost_based, assertions: [picks_lowest_cost], exercised_on: [chat_completions], source: "router_strategy/lowest_cost.py", rationale: "Spend-aware routing"}

tests/e2e/e2e_http.py

@@ -676,6 +676,35 @@ def send(
     return streaming_outcome(resp, stream, sent_at=sent_at)
 
 
+class AbandonedRequest(BaseModel):
+    """A non-streaming request whose socket the client closed ``after`` seconds in,
+    before the proxy had answered."""
+
+    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 close the connection ``after`` seconds if no response head has arrived
+    by then; returns the response instead when the proxy answered first."""
+    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."""

tests/e2e/gateway/stage_mirror_ci_config.yml

@@ -10,6 +10,7 @@ general_settings:
   store_prompts_in_spend_logs: true
   database_connection_pool_limit: 10
   forward_client_headers_to_llm_api: false
+  cancel_on_disconnect: true
   maximum_spend_logs_retention_period: "60d"
   maximum_spend_logs_cleanup_cron: "0 1 * * *"
   proxy_budget_rescheduler_min_time: 15

tests/e2e/models.py

@@ -951,6 +951,23 @@ class RouterSettingsResponse(BaseModel):
     current_values: RouterCurrentValues
 
 
+class ConfigListParams(BaseModel):
+    config_type: Literal["general_settings"]
+
+
+class ConfigField(BaseModel):
+    """One row of GET /config/list: a general_settings field and the value the
+    proxy is running with, the two fields a test preconditions on."""
+
+    model_config = ConfigDict(extra="ignore")
+    field_name: str
+    field_value: JsonValue = None
+
+
+class ConfigFieldList(RootModel[tuple[ConfigField, ...]]):
+    """GET /config/list answers with a bare array of general_settings fields."""
+
+
 class CostMapEntry(BaseModel):
     model_config = ConfigDict(extra="ignore")
     litellm_provider: str | None = None
@@ -1066,6 +1083,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/proxy_client.py

@@ -47,6 +47,8 @@ from models import (
     AnthropicMessagesResponse,
     ChatBody,
     ChatResponse,
+    ConfigFieldList,
+    ConfigListParams,
     CostMap,
     CostMapEntry,
     CountTokensBody,
@@ -630,6 +632,19 @@ class ProxyClient:
             provider_live=provider_live,
         )
 
+    def general_setting_enabled(self, field_name: str) -> bool:
+        """Whether the proxy is running with the named general_settings flag on, for
+        a test whose behavior only exists under a config flag the stack has to carry."""
+        fields = unwrap(
+            self.transport.get(
+                "/config/list",
+                headers=self.transport.master,
+                params=ConfigListParams(config_type="general_settings"),
+                response_type=ConfigFieldList,
+            )
+        ).root
+        return any(entry.field_name == field_name and entry.field_value is True for entry in fields)
+
     def register_model(
         self, body: ModelNewBody, listed_for: str | None = None, *, provider_live: bool = False
     ) -> str:

tests/e2e/router/reliability_support.py

@@ -42,7 +42,7 @@ REAL_KEY = "os.environ/OPENAI_API_KEY"
 CACHING_MODEL = "anthropic/claude-haiku-4-5"
 CACHING_KEY = "os.environ/ANTHROPIC_API_KEY"
 
-CONTENT_FILTERED_MODEL = "azure/gpt-5.4-nano"
+AZURE_MODEL = "azure/gpt-5.4-nano"
 AZURE_KEY = "os.environ/AZURE_API_KEY"
 AZURE_BASE = "os.environ/AZURE_API_BASE"
 AZURE_API_VERSION = "2024-10-21"
@@ -53,6 +53,7 @@ CONTENT_POLICY_PROMPT = (
 )
 
 COOLDOWN_SECONDS = 30.0
+REPLICA_PROPAGATION_SECONDS = 15.0
 
 # The smallest-context chat model OpenAI still serves (16385 tokens). A prompt
 # past that limit comes back as a real `context_length_exceeded` 400, which is
@@ -111,7 +112,7 @@ def create_content_filtered_deployment(proxy: ProxyClient, name: str) -> str:
     return proxy.create_model(
         name,
         LiteLLMParamsBody(
-            model=CONTENT_FILTERED_MODEL,
+            model=AZURE_MODEL,
             api_key=AZURE_KEY,
             api_base=AZURE_BASE,
             api_version=AZURE_API_VERSION,
@@ -120,6 +121,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:
+    """The live Azure OpenAI deployment holding all of the group's shuffle weight,
+    benched on its first failure of any class, with the client's own retries off."""
+    return proxy.register_model(
+        ModelNewBody(
+            model_name=name,
+            litellm_params=LiteLLMParamsBody(
+                model=AZURE_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,139 @@
+"""Live e2e: a client hanging up mid-request under cancel_on_disconnect never
+benches the deployment it was talking to.
+
+The group is the cooldown suite's pair: the live Azure deployment holding all of
+the shuffle weight, benched on its first failure of any class with a cooldown that
+outlasts the test, plus a healthy backup at weight 0 the shuffle only reaches once
+the Azure deployment is benched. A cheap call first proves the Azure deployment
+answers the key and warms its auth path. The test then asks for an answer far
+longer than CLIENT_HANGS_UP_AFTER_SECONDS of generation, retries off, and hangs up
+that many seconds in: late enough that the proxy has handed the call to Azure (a
+hang-up before the provider call is in flight cancels nothing the router could
+bench, so the cell would pass vacuously). An answer that comes back inside the
+window proves nothing and benches nothing either, since a success never counts
+against the deployment, so the cell asks again up to HANG_UP_ATTEMPTS times and
+fails out loud naming the window only when every ask came back early. After the
+cooldown suite's replica propagation window, every one of the next calls has to
+come back 200 from the Azure deployment itself, named in x-litellm-model-id; a
+single answer from the backup means the hang-up was booked as a failure.
+
+The test reads `cancel_on_disconnect` back from the proxy first: without the flag
+the hang-up cancels nothing and the cell would pass vacuously.
+"""
+
+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, StreamingResponse
+from lifecycle import ResourceManager
+from models import ChatMessage, ReliabilityChatBody, RouterSettingsOverride
+from reliability_support import (
+    REPLICA_PROPAGATION_SECONDS,
+    chat_override,
+    create_azure_benched_on_first_failure_deployment,
+    create_zero_weight_backup_deployment,
+    model_id_of,
+)
+
+pytestmark = pytest.mark.e2e
+
+CLIENT_HANGS_UP_AFTER_SECONDS = 5.0
+HANG_UP_ATTEMPTS = 3
+LONG_ANSWER_MAX_TOKENS = 16384
+BENCH_OUTLASTS_TEST_SECONDS = 300.0
+CALLS_AFTER_HANGUP = 6
+
+
+def _say_hi(client: ComplexityRouterClient, key: str, group: str) -> StreamingResponse:
+    return chat_override(
+        client.proxy,
+        key,
+        group,
+        f"say hi {unique_marker()}",
+        override=RouterSettingsOverride(num_retries=0),
+    )
+
+
+def _ask_for_a_long_answer_then_hang_up(
+    client: ComplexityRouterClient, key: str, group: str
+) -> AbandonedRequest | StreamingResponse:
+    return client.proxy.transport.abandon(
+        "/chat/completions",
+        headers=client.proxy.transport.bearer(key),
+        json=ReliabilityChatBody(
+            model=group,
+            messages=[
+                ChatMessage(
+                    role="user",
+                    content=(
+                        "Write an essay on the history of the telegraph with one section per decade from the 1830s "
+                        f"to the 2020s, each section at least 300 words. {unique_marker()}"
+                    ),
+                )
+            ],
+            max_tokens=LONG_ANSWER_MAX_TOKENS,
+            router_settings_override=RouterSettingsOverride(num_retries=0),
+        ),
+        after=CLIENT_HANGS_UP_AFTER_SECONDS,
+    )
+
+
+def _hang_up_mid_answer(client: ComplexityRouterClient, key: str, group: str) -> None:
+    for attempt in range(1, HANG_UP_ATTEMPTS + 1):
+        match _ask_for_a_long_answer_then_hang_up(client, key, group):
+            case AbandonedRequest():
+                return
+            case StreamingResponse(status_code=200):
+                continue
+            case StreamingResponse(status_code=status_code, body=body):
+                pytest.fail(
+                    f"hang-up attempt {attempt} should have found the long answer still in flight after "
+                    f"{CLIENT_HANGS_UP_AFTER_SECONDS:.0f}s, but the proxy answered {status_code}: {body[:300]}"
+                )
+    pytest.fail(
+        f"the proxy answered all {HANG_UP_ATTEMPTS} long asks within {CLIENT_HANGS_UP_AFTER_SECONDS:.0f}s, so the "
+        "client never hung up with a call still in flight and the bench this cell guards against could not happen"
+    )
+
+
+class TestReliabilityCancelOnDisconnect:
+    @pytest.mark.covers("reliability.cooldown.client_disconnect.stays_healthy")
+    def test_client_hanging_up_never_benches_the_deployment(
+        self, client: ComplexityRouterClient, resources: ResourceManager, scoped_key: str
+    ) -> None:
+        assert client.proxy.general_setting_enabled("cancel_on_disconnect"), (
+            "this cell needs general_settings.cancel_on_disconnect: true in the proxy config; without it the "
+            "hang-up cancels nothing and the bench it guards against can never happen"
+        )
+
+        group = f"reliability-cooldown-disconnect-{unique_marker()}"
+        azure = create_azure_benched_on_first_failure_deployment(
+            client.proxy, group, cooldown_time=BENCH_OUTLASTS_TEST_SECONDS
+        )
+        resources.defer(lambda: client.proxy.delete_model(azure))
+        backup = create_zero_weight_backup_deployment(client.proxy, group)
+        resources.defer(lambda: client.proxy.delete_model(backup))
+
+        warm_up = _say_hi(client, scoped_key, group)
+        assert warm_up.status_code == 200 and model_id_of(warm_up) == azure, (
+            f"before any hang-up the Azure deployment {azure} should answer the group, got {warm_up.status_code} "
+            f"from {model_id_of(warm_up)!r}: {warm_up.body[:300]}"
+        )
+
+        _hang_up_mid_answer(client, scoped_key, group)
+        time.sleep(REPLICA_PROPAGATION_SECONDS)
+
+        for call in range(1, CALLS_AFTER_HANGUP + 1):
+            resp = _say_hi(client, scoped_key, group)
+            assert resp.status_code == 200, (
+                f"call {call} after the hang-up should have been a plain 200 from the group, got "
+                f"{resp.status_code}: {resp.body[:300]}"
+            )
+            assert model_id_of(resp) == azure, (
+                f"call {call} after the hang-up should have been served by the Azure deployment {azure}, the proxy "
+                f"named {model_id_of(resp)!r}: the cancelled call was booked as a failure and benched it"
+            )

tests/e2e/router/test_reliability_cooldowns_e2e.py

@@ -43,6 +43,7 @@ from lifecycle import ResourceManager
 from models import KeyGenerateBody, RouterSettingsOverride
 from reliability_support import (
     COOLDOWN_SECONDS,
+    REPLICA_PROPAGATION_SECONDS,
     chat_override,
     create_always_5xx_deployment,
     create_always_rate_limited_deployment,
@@ -57,7 +58,6 @@ from reliability_support import (
 pytestmark = pytest.mark.e2e
 
 RECOVERY_GRACE_SECONDS = 10
-REPLICA_PROPAGATION_SECONDS = 15.0
 PROPAGATION_POLL_SECONDS = 0.25
 BENCH_MARGIN_SECONDS = 4.0
 

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)
 

tests/test_litellm/llms/azure/test_azure.py

@@ -1,10 +1,13 @@
 """Tests for litellm/llms/azure/azure.py AzureChatCompletion handler behaviour."""
 
+import asyncio
 import time
 from typing import Final
 
-from openai import AzureOpenAI
+import pytest
+from openai import AsyncAzureOpenAI, AzureOpenAI
 
+import litellm
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 from litellm.llms.azure.azure import AzureChatCompletion
 
@@ -52,3 +55,25 @@ def test_sync_streaming_stamps_response_headers_on_the_logging_obj() -> None:
     )
 
     assert logging_obj.model_call_details["response_headers"] == {"x-ms-is-spilled-over": "true"}
+
+
+class _CancelledRawCompletions:
+    async def create(self, **kwargs):
+        raise asyncio.CancelledError()
+
+
+@pytest.mark.asyncio
+async def test_acompletion_propagates_cancelled_error() -> None:
+    client = AsyncAzureOpenAI(
+        api_key="fake-key",
+        api_version="2024-02-01",
+        azure_endpoint="https://fake-resource.openai.azure.com",
+    )
+    client.chat.completions.with_raw_response = _CancelledRawCompletions()
+
+    with pytest.raises(asyncio.CancelledError):
+        await litellm.acompletion(
+            model="azure/fake-deployment",
+            messages=[{"role": "user", "content": "hi"}],
+            client=client,
+        )