fix(proxy): report a bridged stream failure once, after the client stream's own errors, and skip disconnect billing for failed or background Responses streams

litellm/proxy/common_request_processing.py

@@ -4,6 +4,7 @@ import json
 import logging
 import math
 from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Coroutine, Mapping, Sequence
+from dataclasses import dataclass
 from datetime import datetime
 from functools import lru_cache
 from types import MappingProxyType
@@ -1962,6 +1963,43 @@ def _content_text(fields: Mapping[str, object]) -> str:
     return str(fields.get("content", ""))
 
 
+async def _no_stream_failure_to_release() -> None:
+    return None
+
+
+@dataclass(slots=True)
+class _ClientStreamFailureReport:
+    logging_obj: LiteLLMLoggingObj
+    report: Callable[[Exception], Awaitable[None]]
+    held: Exception | None = None
+
+    @classmethod
+    def hold_while_streaming(cls, logging_obj: LiteLLMLoggingObj | None) -> Callable[[], Awaitable[None]]:
+        if logging_obj is None or logging_obj.on_detached_stream_failure is None:
+            return _no_stream_failure_to_release
+        return cls(logging_obj, logging_obj.on_detached_stream_failure)._start_holding()
+
+    def _start_holding(self) -> Callable[[], Awaitable[None]]:
+        self.logging_obj.on_detached_stream_failure = self.hold
+        return self.release
+
+    async def hold(self, failure: Exception) -> None:
+        if self.held is None:
+            self.held = failure
+
+    async def release(self) -> None:
+        self.logging_obj.on_detached_stream_failure = self.report
+        if self.held is None:
+            return
+        try:
+            with anyio.CancelScope(shield=True):
+                await self.report(self.held)
+        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)
+            )
+
+
 class ProxyBaseLLMRequestProcessing:
     def __init__(self, data: dict, served_model: str | None = None):
         self.data = data
@@ -4404,6 +4442,9 @@ class ProxyBaseLLMRequestProcessing:
             response=response,
             request_data=request_data,
         )
+        release_stream_failure: Final = _ClientStreamFailureReport.hold_while_streaming(
+            stream_logging_obj if isinstance(stream_logging_obj, LiteLLMLoggingObj) else None
+        )
         try:
             str_so_far = ""
             async for chunk in guarded_stream:
@@ -4508,6 +4549,7 @@ 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,

litellm/responses/streaming_iterator.py

@@ -1067,12 +1067,14 @@ class BaseResponsesAPIStreamingIterator:
         the same shape the success handler assembles for a natural end of stream.
         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
-        or the stream already reached its terminal event and billed itself.
+        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.
         """
         created: Final = self._created_response
         generated: Final = self._generated_content + self._generated_tool_arguments
-        if self.completed_response is not None or created is None or not generated:
+        if created is None or not generated or self._billed_without_disconnect():
             return None
         usage: Final = _estimate_usage_safely(
             self.model or "",
@@ -1093,6 +1095,11 @@ class BaseResponsesAPIStreamingIterator:
         stamp_responses_usage_cost(partial, self.logging_obj)
         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
+        )
+
     def _note_yielded_event(self, event: ResponsesAPIStreamingResponse) -> None:
         self._yielded_first_chunk = True
         if event.type not in PRE_OUTPUT_LIFECYCLE_EVENT_TYPES:

tests/unit/proxy/test_common_request_processing.py

@@ -10815,27 +10815,93 @@ class TestDetachedStreamFailureHook:
         )
         return logging_obj, stream
 
+    @staticmethod
+    def _hook_reporting(upstream_failure: Exception, after_report=None):
+        def iterator_hook_for(logging_obj):
+            async def iterator_hook(**_kwargs):
+                yield b"event: message_start\n\n"
+                await logging_obj.on_detached_stream_failure(upstream_failure)
+                yield b"event: error\n\n"
+                yield b"event: message_stop\n\n"
+                if after_report is not None:
+                    raise after_report
+
+            return iterator_hook
+
+        return iterator_hook_for
+
     @pytest.mark.asyncio
-    async def test_failure_reported_by_the_stream_reaches_the_failure_hook_once_when_an_iterator_hook_raises_after_it(
-        self,
-    ):
+    async def test_iterator_hook_failure_after_a_bridged_failure_is_the_one_the_failure_hook_gets(self):
         recorder = _FailureHookRecorder()
         upstream_failure = RuntimeError("upstream Responses stream failed")
         hook_failure = RuntimeError("iterator hook rejected the stream")
 
-        def iterator_hook_for(logging_obj):
-            async def iterator_hook(**_kwargs):
-                yield b"event: error\n\n"
-                await logging_obj.on_detached_stream_failure(upstream_failure)
-                raise hook_failure
-
-            return iterator_hook
-
-        _, stream = self._armed_stream(recorder, iterator_hook_for)
+        _, stream = self._armed_stream(recorder, self._hook_reporting(upstream_failure, after_report=hook_failure))
 
         frames = [frame async for frame in stream]
 
-        assert frames == ["data: chunk\n\n", "data: error\n\n"]
+        assert frames == ["data: chunk\n\n"] * 3 + ["data: error\n\n"]
+        assert [call["original_exception"] for call in recorder.calls] == [hook_failure]
+
+    @pytest.mark.asyncio
+    async def test_failure_hook_replacing_the_iterator_hook_failure_aborts_the_stream(self):
+        replacement = HTTPException(status_code=418, detail="rejected by a custom failure hook")
+        recorder = _FailureHookRecorder(raises=replacement)
+        upstream_failure = RuntimeError("upstream Responses stream failed")
+        hook_failure = RuntimeError("iterator hook rejected the stream")
+
+        _, stream = self._armed_stream(recorder, self._hook_reporting(upstream_failure, after_report=hook_failure))
+
+        frames = [await stream.__anext__() for _ in range(3)]
+        with pytest.raises(HTTPException) as raised:
+            await stream.__anext__()
+
+        assert raised.value is replacement
+        assert frames == ["data: chunk\n\n"] * 3
+        assert [call["original_exception"] for call in recorder.calls] == [hook_failure]
+
+    @pytest.mark.asyncio
+    async def test_failure_reported_by_the_stream_reaches_the_failure_hook_once_after_the_stream_ends(self):
+        recorder = _FailureHookRecorder()
+        upstream_failure = RuntimeError("upstream Responses stream failed")
+
+        _, stream = self._armed_stream(recorder, self._hook_reporting(upstream_failure))
+
+        frames = []
+        async for frame in stream:
+            frames.append(frame)
+            assert recorder.calls == []
+
+        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_reported_by_the_stream_reaches_the_failure_hook_once_when_the_client_leaves_early(self):
+        recorder = _FailureHookRecorder()
+        upstream_failure = RuntimeError("upstream Responses stream failed")
+
+        _, stream = self._armed_stream(recorder, self._hook_reporting(upstream_failure))
+
+        await stream.__anext__()
+        await stream.__anext__()
+        await stream.aclose()
+
+        assert [call["original_exception"] for call in recorder.calls] == [upstream_failure]
+
+    @pytest.mark.asyncio
+    async def test_failure_reported_after_the_stream_ends_still_reaches_the_failure_hook(self):
+        recorder = _FailureHookRecorder()
+        upstream_failure = RuntimeError("upstream Responses stream failed after the client stream closed")
+
+        async def iterator_hook(**_kwargs):
+            yield b"event: message_start\n\n"
+
+        logging_obj, stream = self._armed_stream(recorder, lambda _logging_obj: iterator_hook)
+
+        frames = [frame async for frame in stream]
+        await logging_obj.on_detached_stream_failure(upstream_failure)
+
+        assert frames == ["data: chunk\n\n"]
         assert [call["original_exception"] for call in recorder.calls] == [upstream_failure]
 
     @pytest.mark.asyncio

tests/unit/responses/test_responses_streaming_iterator.py

@@ -440,17 +440,28 @@ def _completed_after(text: str) -> Mapping[str, object]:
     }
 
 
-def _live_iterator(logging_obj: LitellmLogging) -> ResponsesAPIStreamingIterator:
+def _live_iterator(
+    logging_obj: LitellmLogging, upstream: object = None, background: bool = False
+) -> ResponsesAPIStreamingIterator:
     return ResponsesAPIStreamingIterator(
-        response=httpx.Response(200),
+        response=cast(httpx.Response, upstream or httpx.Response(200)),
         model="gpt-5.4-nano",
         responses_api_provider_config=OpenAIResponsesAPIConfig(),
         logging_obj=logging_obj,
-        request_data={"model": "gpt-5.4-nano", "input": "tell me a story"},
+        request_data={"model": "gpt-5.4-nano", "input": "tell me a story", "background": background},
         custom_llm_provider="openai",
     )
 
 
+class _UpstreamThatTimesOutAfterRelayingText:
+    headers: dict[str, str] = {}
+
+    async def aiter_bytes(self) -> AsyncIterator[bytes]:
+        for event in (_CREATED_WITH_PRIORITY_TIER, _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")
+
+
 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)):
@@ -492,3 +503,22 @@ async def test_partial_response_on_disconnect_is_none_when_there_is_nothing_to_b
 
     assert iterator.partial_response_on_disconnect() is None
     await iterator._await_pending_logging()
+
+
+@pytest.mark.asyncio
+async def test_partial_response_on_disconnect_is_none_once_the_failure_log_billed_the_relayed_text():
+    iterator: Final = _live_iterator(_make_logging_obj(), upstream=_UpstreamThatTimesOutAfterRelayingText())
+
+    with pytest.raises(litellm.Timeout):
+        await _ChunkRecorder().drain(iterator)
+
+    assert iterator.partial_response_on_disconnect() is None
+    await iterator._await_pending_logging()
+
+
+def test_partial_response_on_disconnect_is_none_for_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