fix(proxy): keep the bridged failure report and stream cleanup alive through a client cancel, bill failed estimates at the served tier, and count extra_body background

litellm/proxy/common_request_processing.py

@@ -129,7 +129,7 @@ from litellm.proxy.utils import (  # noqa: F401  # legacy module exports
     _check_and_merge_model_level_guardrails,  # pyright: ignore[reportPrivateUsage,reportUnusedImport]  # backwards-compatible package export
     check_and_merge_model_level_guardrails,
 )
-from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect
+from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect, spawn_logging_task
 from litellm.router import Router
 from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict
 from litellm.router_utils.auto_router_model_naming import AUTO_ROUTER_MODEL_PREFIX
@@ -1991,9 +1991,15 @@ class _ClientStreamFailureReport:
         self.logging_obj.on_detached_stream_failure = self.report
         if self.held is None:
             return
+        await asyncio.shield(
+            spawn_logging_task(
+                asyncio.get_running_loop(), self._report_held(self.held), task_name="stream failure report"
+            )
+        )
+
+    async def _report_held(self, failure: Exception) -> None:
         try:
-            with anyio.CancelScope(shield=True):
-                await self.report(self.held)
+            await self.report(failure)
         except Exception as hook_failure:  # noqa: BLE001  # a failing failure hook must not skip the stream cleanup
             verbose_proxy_logger.warning(
                 "Stream failure hook raised: %s(%s)", type(hook_failure).__name__, sanitize_for_log(hook_failure)
@@ -4549,16 +4555,18 @@ class ProxyBaseLLMRequestProcessing:
             seal: Final = "" if seal_open_frame is None else seal_open_frame(recent_tail)
             yield seal + error_frame if seal else error_frame
         finally:
-            await release_stream_failure()
-            await ProxyBaseLLMRequestProcessing.finalize_streaming_generator_cleanup(
-                request=request,
-                request_data=request_data,
-                response=response,
-                stream_completed=stream_completed,
-                client_disconnected=client_disconnected,
-                user_api_key_dict=user_api_key_dict,
-                proxy_logging_obj=proxy_logging_obj,
-            )
+            try:
+                await release_stream_failure()
+            finally:
+                await ProxyBaseLLMRequestProcessing.finalize_streaming_generator_cleanup(
+                    request=request,
+                    request_data=request_data,
+                    response=response,
+                    stream_completed=stream_completed,
+                    client_disconnected=client_disconnected,
+                    user_api_key_dict=user_api_key_dict,
+                    proxy_logging_obj=proxy_logging_obj,
+                )
 
     @staticmethod
     def async_sse_data_generator(

litellm/responses/streaming_iterator.py

@@ -41,7 +41,7 @@ from litellm.exceptions import MidStreamFallbackError, RateLimitError
 from litellm.litellm_core_utils.asyncify import run_async_function
 from litellm.litellm_core_utils.core_helpers import process_response_headers
 from litellm.litellm_core_utils.exception_mapping_utils import with_deployment_api_base
-from litellm.litellm_core_utils.internal_call_metadata import without_budget_reservation
+from litellm.litellm_core_utils.internal_call_metadata import requests_background, without_budget_reservation
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base
 from litellm.litellm_core_utils.llm_response_utils.get_headers import provider_request_id_headers
@@ -974,7 +974,7 @@ class BaseResponsesAPIStreamingIterator:
         generated_text: Final = self._generated_content + self._generated_tool_arguments
         if (
             not generated_text
-            or self.request_data.get("background") is True
+            or requests_background(self.request_data)
             or self.logging_obj.model_call_details.get("combined_usage_object") is not None
         ):
             return
@@ -990,7 +990,12 @@ class BaseResponsesAPIStreamingIterator:
         usage: Final = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(  # pyright: ignore[reportPrivateUsage]  # same shared transform the success path above uses
             estimate
         )
-        self.logging_obj.record_assembled_response_for_failure(ModelResponse(model=self.model, usage=usage))
+        served_tier: Final[object] = getattr(self._created_response, "service_tier", None)
+        self.logging_obj.record_assembled_response_for_failure(
+            ModelResponse(
+                model=self.model, usage=usage, service_tier=served_tier if isinstance(served_tier, str) else None
+            )
+        )
 
     def _handle_failure(self, exception: Exception):
         """
@@ -1068,9 +1073,8 @@ class BaseResponsesAPIStreamingIterator:
         Its response is the response.created payload (it carries the served
         service_tier) plus the text and tool arguments relayed so far, with counted
         usage since the provider never sent any. None when nothing was relayed yet,
-        the stream already reached its terminal event and billed itself, its failure
-        log already billed the relayed output, or it runs in the background, where
-        retrieving it bills the whole run.
+        the stream already reached its terminal event and billed itself, or its failure
+        log already billed the relayed output.
         """
         created: Final = self._created_response
         generated: Final = self._generated_content + self._generated_tool_arguments
@@ -1096,9 +1100,7 @@ class BaseResponsesAPIStreamingIterator:
         return ResponseIncompleteEvent(type=ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, response=partial)
 
     def _billed_without_disconnect(self) -> bool:
-        return (
-            self.completed_response is not None or self._failure_handled or self.request_data.get("background") is True
-        )
+        return self.completed_response is not None or self._failure_handled
 
     def _note_yielded_event(self, event: ResponsesAPIStreamingResponse) -> None:
         self._yielded_first_chunk = True

tests/unit/proxy/test_common_request_processing.py

@@ -10676,12 +10676,19 @@ class TestResolveLitellmCallId:
 class _FailureHookRecorder:
     """Stands in for ProxyLogging.post_call_failure_hook, recording what the detached-failure closure hands it."""
 
-    def __init__(self, raises: Optional[Exception] = None):
+    def __init__(self, raises: Optional[Exception] = None, gate: Optional[asyncio.Event] = None):
         self.calls = []
+        self.entered = asyncio.Event()
+        self.reported = asyncio.Event()
         self._raises = raises
+        self._gate = gate
 
     async def post_call_failure_hook(self, **kwargs):
+        self.entered.set()
+        if self._gate is not None:
+            await self._gate.wait()
         self.calls.append(kwargs)
+        self.reported.set()
         if self._raises is not None:
             raise self._raises
 
@@ -10782,7 +10789,7 @@ class TestDetachedStreamFailureHook:
         assert [call["original_exception"] for call in recorder.calls] == [failure]
 
     @staticmethod
-    def _armed_stream(recorder: _FailureHookRecorder, iterator_hook_for):
+    def _armed_stream(recorder: _FailureHookRecorder, iterator_hook_for, response: Optional[MagicMock] = None):
         from litellm.proxy._types import UserAPIKeyAuth as RealUserAPIKeyAuth
 
         logging_obj = LiteLLMLoggingObj(
@@ -10806,7 +10813,7 @@ class TestDetachedStreamFailureHook:
             proxy_logging_obj=proxy_logging_obj,
         )
         stream = ProxyBaseLLMRequestProcessing.async_streaming_data_generator(
-            response=MagicMock(),
+            response=MagicMock() if response is None else response,
             user_api_key_dict=user_api_key_dict,
             request_data=processing_obj.data,
             proxy_logging_obj=proxy_logging_obj,
@@ -10875,6 +10882,30 @@ class TestDetachedStreamFailureHook:
         assert frames == ["data: chunk\n\n"] * 3
         assert [call["original_exception"] for call in recorder.calls] == [upstream_failure]
 
+    @pytest.mark.asyncio
+    async def test_failure_report_and_stream_cleanup_survive_the_client_task_being_cancelled_mid_report(self):
+        gate = asyncio.Event()
+        recorder = _FailureHookRecorder(gate=gate)
+        upstream_failure = RuntimeError("upstream Responses stream failed")
+        response = MagicMock()
+        response.aclose = AsyncMock()
+
+        _, stream = self._armed_stream(recorder, self._hook_reporting(upstream_failure), response=response)
+
+        async def client_reads_the_whole_stream():
+            return [frame async for frame in stream]
+
+        client = asyncio.create_task(client_reads_the_whole_stream())
+        await asyncio.wait_for(recorder.entered.wait(), timeout=5)
+        client.cancel()
+        with pytest.raises(asyncio.CancelledError):
+            await client
+
+        response.aclose.assert_awaited_once()
+        gate.set()
+        await asyncio.wait_for(recorder.reported.wait(), timeout=5)
+        assert [call["original_exception"] for call in recorder.calls] == [upstream_failure]
+
     @pytest.mark.asyncio
     async def test_failure_reported_by_the_stream_reaches_the_failure_hook_once_when_the_client_leaves_early(self):
         recorder = _FailureHookRecorder()

tests/unit/responses/test_responses_streaming_iterator.py

@@ -270,7 +270,7 @@ async def _failure_payload_after_mid_stream_timeout(
     monkeypatch: pytest.MonkeyPatch,
     request_input: str | list[dict[str, object]],
     upstream: _UpstreamThatTimesOutAfterThreeDeltas | None = None,
-    background: bool = False,
+    background_fields: Mapping[str, object] | None = None,
 ) -> tuple[FailureRecorder, StandardLoggingPayload]:
     upstream_response: Final = upstream or _UpstreamThatTimesOutAfterThreeDeltas()
     recorder = FailureRecorder()
@@ -290,7 +290,7 @@ async def _failure_payload_after_mid_stream_timeout(
         responses_api_provider_config=OpenAIResponsesAPIConfig(),
         logging_obj=logging_obj,
         custom_llm_provider="openai",
-        request_data={"model": "gpt-5.4-nano", "input": request_input, "stream": True, "background": background},
+        request_data={"model": "gpt-5.4-nano", "input": request_input, "stream": True, **(background_fields or {})},
     )
     received = _ChunkRecorder()
 
@@ -336,11 +336,18 @@ async def test_mid_stream_read_timeout_bills_the_tool_call_arguments_already_str
     assert payload["response_cost"] > 0
 
 
+@pytest.mark.parametrize(
+    "background_fields",
+    [{"background": True}, {"extra_body": {"background": True}}],
+    ids=["top_level", "extra_body"],
+)
 @pytest.mark.asyncio
 async def test_mid_stream_failure_of_a_background_response_bills_nothing_since_retrieving_it_bills_the_whole_run(
-    monkeypatch: pytest.MonkeyPatch,
+    monkeypatch: pytest.MonkeyPatch, background_fields: Mapping[str, object]
 ) -> None:
-    _, payload = await _failure_payload_after_mid_stream_timeout(monkeypatch, _SEA_PROMPT, background=True)
+    _, payload = await _failure_payload_after_mid_stream_timeout(
+        monkeypatch, _SEA_PROMPT, background_fields=background_fields
+    )
 
     assert payload["status"] == "failure"
     assert payload["completion_tokens"] == 0
@@ -462,6 +469,55 @@ class _UpstreamThatTimesOutAfterRelayingText:
         raise httpx.ReadTimeout("Timeout on reading data from socket")
 
 
+class _UpstreamServedAtTierThatTimesOut:
+    headers: dict[str, str] = {}
+
+    def __init__(self, served_tier: str) -> None:
+        self.served_tier = served_tier
+
+    async def aiter_bytes(self) -> AsyncIterator[bytes]:
+        created: Final = {
+            **_CREATED_WITH_PRIORITY_TIER,
+            "response": {**_CREATED_RESPONSE_WITH_PRIORITY_TIER, "service_tier": self.served_tier},
+        }
+        for event in (created, _text_delta("Once upon a time", 1)):
+            yield f"data: {json.dumps(event)}\n\n".encode()
+        raise httpx.ReadTimeout("Timeout on reading data from socket")
+
+
+async def _mid_stream_failure_cost(
+    monkeypatch: pytest.MonkeyPatch, requested_tier: str | None, served_tier: str
+) -> float:
+    recorder: Final = FailureRecorder()
+    monkeypatch.setattr(litellm, "failure_callback", [recorder])
+    monkeypatch.setattr(litellm, "_async_failure_callback", [recorder])
+    logging_obj: Final = _make_logging_obj()
+    logging_obj.update_environment_variables(
+        model="gpt-5.4-nano",
+        optional_params={"stream": True} if requested_tier is None else {"stream": True, "service_tier": requested_tier},
+        litellm_params={"custom_llm_provider": "openai", "aresponses": True},
+    )
+    iterator: Final = _live_iterator(logging_obj, upstream=_UpstreamServedAtTierThatTimesOut(served_tier))
+
+    with pytest.raises(litellm.Timeout):
+        await _ChunkRecorder().drain(iterator)
+    await iterator._await_pending_logging()
+
+    (payload,) = recorder.failure_payloads
+    return payload["response_cost"]
+
+
+@pytest.mark.asyncio
+async def test_mid_stream_failure_bills_the_relayed_text_at_the_tier_the_provider_served(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    flex_requested_default_served: Final = await _mid_stream_failure_cost(monkeypatch, "flex", "default")
+    nothing_requested_default_served: Final = await _mid_stream_failure_cost(monkeypatch, None, "default")
+
+    assert nothing_requested_default_served > 0
+    assert flex_requested_default_served == nothing_requested_default_served
+
+
 def test_partial_response_on_disconnect_bills_the_relayed_text_at_the_served_service_tier():
     iterator: Final = _live_iterator(_make_logging_obj())
     for event in (_CREATED_WITH_PRIORITY_TIER, _text_delta("Once upon a time ", 1), _text_delta("a fox slept.", 2)):
@@ -516,9 +572,12 @@ async def test_partial_response_on_disconnect_is_none_once_the_failure_log_bille
     await iterator._await_pending_logging()
 
 
-def test_partial_response_on_disconnect_is_none_for_a_background_response():
+def test_partial_response_on_disconnect_reports_the_relayed_text_of_a_background_response():
     iterator: Final = _live_iterator(_make_logging_obj(), background=True)
     for event in (_CREATED_WITH_PRIORITY_TIER, _text_delta("Once upon a time", 1)):
         iterator._process_chunk(json.dumps(event))
 
-    assert iterator.partial_response_on_disconnect() is None
+    incomplete_event: Final = iterator.partial_response_on_disconnect()
+
+    assert incomplete_event is not None
+    assert incomplete_event.response.output[0].content[0].text == "Once upon a time"