From caba34983f9bf40d5ef73d88760d6793e7105548 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 31 Oct 2026 19:58:39 -0700 Subject: [PATCH] fix(proxy): read polled chunks as SSE events and keep draining past a mid-chunk [DONE] A background poll split a chunk into data lines, so one event with two data lines parsed as two payloads, and a frame with an integer past the digit limit raised a ValueError the parse guard did not catch, which failed the poll and dropped its spend row. The poll now joins each event's data lines the way an SSE reader does and skips any frame it cannot parse A [DONE] behind another frame in one chunk ended the poll and left the stream unbilled. Only a chunk that opens with [DONE] ends the poll now, as on main, and the poll stops reading a chunk at its first terminal event, so two terminal events in one chunk poll the first one as main did --- .../common_utils/responses_stream_errors.py | 5 +++ .../response_polling/background_streaming.py | 18 ++++----- .../test_response_polling_handler.py | 39 +++++++++++++++++++ 3 files changed, 51 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/common_utils/responses_stream_errors.py b/litellm/proxy/common_utils/responses_stream_errors.py index ca948f65bd2..8dccd938a90 100644 --- a/litellm/proxy/common_utils/responses_stream_errors.py +++ b/litellm/proxy/common_utils/responses_stream_errors.py @@ -484,6 +484,11 @@ def _terminal_blocks(frame: str) -> tuple[tuple[tuple[str, ...], _StreamEvent], return tuple((lines, event) for lines, event in read if event.type in _TERMINAL_EVENT_TYPES) +def frame_payloads(frame: str) -> tuple[tuple[str, ...], bool]: + payloads: Final = tuple(_block_data(lines).strip() for lines in _blocks(frame) if _field_values(lines, "data")) + return tuple(takewhile(lambda payload: payload != "[DONE]", payloads)), payloads[:1] == ("[DONE]",) + + def frame_error_data(frame: str) -> str | None: terminal: Final = _terminal_blocks(frame) if not terminal: diff --git a/litellm/proxy/response_polling/background_streaming.py b/litellm/proxy/response_polling/background_streaming.py index 9d211fe4e71..73190ad8df8 100644 --- a/litellm/proxy/response_polling/background_streaming.py +++ b/litellm/proxy/response_polling/background_streaming.py @@ -12,7 +12,6 @@ import asyncio import json from collections.abc import AsyncIterable, Callable, Mapping, Sequence from http import HTTPStatus -from itertools import takewhile from time import perf_counter from typing import TYPE_CHECKING, Final, Protocol, TypeAlias, runtime_checkable @@ -21,12 +20,12 @@ from starlette.types import Message from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger -from litellm.llms.custom_httpx.sse_lines import split_sse_lines from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth from litellm.proxy.common_request_processing import RESPONSES_SSE_FRAMING, ProxyBaseLLMRequestProcessing from litellm.proxy.common_utils.responses_stream_errors import ( event_error_object, frame_error_data, + frame_payloads, responses_error_fields, sse_error_object, ) @@ -91,11 +90,6 @@ class _StreamsBody(Protocol): def body_iterator(self) -> AsyncIterable[str | bytes | memoryview]: ... -def _payloads_before_done(frame: str) -> tuple[tuple[str, ...], bool]: - payloads: Final = tuple(line[6:].strip() for line in split_sse_lines(frame) if line.startswith("data: ")) - return tuple(takewhile(lambda payload: payload != "[DONE]", payloads)), "[DONE]" in payloads - - def _terminal_response(event: _StreamEvent) -> _TerminalResponse: match event.get("response"): case Mapping() as response: @@ -289,7 +283,7 @@ async def background_streaming_task( ): frame_failure = responses_error_fields(_event_error(_StreamEventParser.parse(error_data))) - payloads, stream_done = _payloads_before_done(chunk) if isinstance(chunk, str) else ((), False) + payloads, stream_done = frame_payloads(chunk) if isinstance(chunk, str) else ((), False) for chunk_data in payloads: try: event: _StreamEvent = _StreamEventParser.parse(chunk_data) @@ -312,7 +306,7 @@ async def background_streaming_task( # Process different event types based on OpenAI streaming spec if event_type == "response.output_item.added": # New output item added - item = event.get("item", {}) + item = event.get("item", _OutputItem()) item_id = item.get("id") if item_id: output_items[item_id] = item @@ -379,7 +373,7 @@ async def background_streaming_task( elif event_type == "response.output_item.done": # Output item completed - use final item data - item = event.get("item", {}) + item = event.get("item", _OutputItem()) item_id = item.get("id") if item_id: output_items[item_id] = item @@ -443,10 +437,12 @@ async def background_streaming_task( continue state_dirty = True - except (json.JSONDecodeError, AttributeError, TypeError) as e: + except (ValueError, AttributeError, TypeError) as e: verbose_proxy_logger.warning("Failed to parse streaming chunk: %s", e) else: await flush_state_if_needed() + if event_type in _event_to_status: + break if stream_done: break diff --git a/tests/unit/proxy/response_polling/test_response_polling_handler.py b/tests/unit/proxy/response_polling/test_response_polling_handler.py index 43184f95700..752dfc5f124 100644 --- a/tests/unit/proxy/response_polling/test_response_polling_handler.py +++ b/tests/unit/proxy/response_polling/test_response_polling_handler.py @@ -1701,6 +1701,41 @@ class TestBackgroundStreamingTerminalEvents: None, {"input_tokens": 3}, ), + ( + ( + 'data: {"type": "response.output_text.delta", "delta": "hi"}\n\n' + 'data: {"type": "hook.huge", "n": ' + "9" * 4400 + "}\n\n", + 'data: {"type": "response.completed", "response": {"status": "completed", "usage": {"input_tokens": 3}}}\n\n', + ), + "completed", + None, + {"input_tokens": 3}, + ), + ( + ( + 'data: {"type": "response.output_text.delta", "delta": "hi"}\n\ndata: [DONE]\n\n', + 'data: {"type": "response.completed", "response": {"status": "completed", "usage": {"input_tokens": 3}}}\n\n', + ), + "completed", + None, + {"input_tokens": 3}, + ), + ( + ( + 'data: {"type": "response.completed", "response": {"status": "completed", "usage": {"input_tokens": 3}}}\n\n' + 'data: {"type": "response.failed", "response": {"status": "failed", ' + '"error": {"code": "server_error", "message": "late"}}}\n\n', + ), + "completed", + None, + {"input_tokens": 3}, + ), + ( + ('data: event: response.incomplete\ndata: {"type": "response.incomplete", "response": {"status": "incomplete"}}\n\n',), + "completed", + None, + None, + ), ], ids=[ "error_frame_after_a_delta", @@ -1714,6 +1749,10 @@ class TestBackgroundStreamingTerminalEvents: "failed_event_whose_response_is_null", "error_frame_and_incomplete_in_one_chunk", "error_frame_then_completed_behind_a_delta_in_one_chunk", + "frame_past_the_int_digit_limit_behind_a_delta_then_completed", + "done_behind_a_delta_in_one_chunk_then_completed", + "completed_then_failed_in_one_chunk", + "event_line_sent_as_a_data_line_of_one_event", ], ) async def test_a_poll_stores_what_its_stream_ended_with(