From a7832c290ab6e817799fbb2698dcf85a87fbccb8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 14:52:02 -0700 Subject: [PATCH] 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 | 42 +++++++++ litellm/responses/streaming_iterator.py | 13 ++- .../proxy/test_common_request_processing.py | 92 ++++++++++++++++--- .../test_responses_streaming_iterator.py | 36 +++++++- 4 files changed, 164 insertions(+), 19 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 5ff6f5c92fd..b57ff91f650 100644 --- a/litellm/proxy/common_request_processing.py +++ b/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, diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 6edb1fe10a0..8ec2d2563d8 100644 --- a/litellm/responses/streaming_iterator.py +++ b/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: diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 8b56c046235..4bca44dbe3c 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/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 diff --git a/tests/unit/responses/test_responses_streaming_iterator.py b/tests/unit/responses/test_responses_streaming_iterator.py index 4dbe795ab1b..9f64540d6ba 100644 --- a/tests/unit/responses/test_responses_streaming_iterator.py +++ b/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