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

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:

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
 

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(