From 09857e30c0c8ec8af8a9f71c2c9c8f9d22d8ebda Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Fri, 30 Oct 2026 01:35:30 -0700
Subject: [PATCH] fix(streaming): let a finished generic chunk end every
 streamed choice

An OpenAI-compatible provider that streams n=2 without per-choice finish_reason
ends with `[DONE]`, which arrives as one finished generic chunk on choice 0. A
drop after it errored because choice 1 never finished, while SDK 2 ended the
stream cleanly. A finished generic chunk now ends every choice streamed so far
---
 .../litellm_core_utils/streaming_handler.py   | 14 ++++++-
 .../test_streaming_handler.py                 | 40 +++++++++++++++++++
 2 files changed, 52 insertions(+), 2 deletions(-)

diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py
index 6e608062cb4..cc2343a62af 100644
--- a/litellm/litellm_core_utils/streaming_handler.py
+++ b/litellm/litellm_core_utils/streaming_handler.py
@@ -108,6 +108,14 @@ def _choice_progress(chunk: object) -> tuple[tuple[int, bool], ...]:
             return ()
 
 
+def _ends_every_choice(chunk: object) -> bool:
+    match chunk:
+        case {"is_finished": True}:
+            return True
+        case _:
+            return False
+
+
 def _stamp_usage_cost(
     usage_chunk: ModelResponse | TextCompletionResponse | ModelResponseStream,
     complete_response: ModelResponse | TextCompletionResponse | ModelResponseStream,
@@ -2277,8 +2285,10 @@ class CustomStreamWrapper:
     def _record_choice_progress(self, chunk: object) -> None:
         progress: Final = _choice_progress(chunk)
         self._streamed_choice_indices = self._streamed_choice_indices.union(index for index, _ in progress)
-        self._finished_choice_indices = self._finished_choice_indices.union(
-            index for index, finished in progress if finished
+        self._finished_choice_indices = (
+            self._streamed_choice_indices
+            if _ends_every_choice(chunk)
+            else self._finished_choice_indices.union(index for index, finished in progress if finished)
         )
 
     def _every_choice_finished(self) -> bool:
diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py
index e7445b0873e..f379104d97c 100644
--- a/tests/unit/litellm_core_utils/test_streaming_handler.py
+++ b/tests/unit/litellm_core_utils/test_streaming_handler.py
@@ -32,6 +32,7 @@ from litellm.litellm_core_utils.streaming_handler import (
     _ProviderChunkParsed,
 )
 from litellm.llms.custom_httpx.http_handler import HTTPHandler
+from litellm.llms.fireworks_ai.chat.transformation import FireworksAIChatCompletionStreamingHandler
 from litellm.llms.ollama.completion.transformation import OllamaTextCompletionResponseIterator
 from litellm.llms.openai.chat.gpt_transformation import OpenAIChatCompletionStreamingHandler
 from litellm.responses.mcp.chat_completions_handler import _add_mcp_metadata_to_response
@@ -5452,6 +5453,45 @@ async def test_ollama_generate_stream_dropped_after_its_done_line_ends_cleanly(l
     assert [chunk.choices[0].finish_reason for chunk in collected if chunk.choices[0].finish_reason] == ["stop"]
 
 
+
+def _fireworks_sse_line_for_choice(index: int, content: str) -> str:
+    chunk: Final = {
+        "id": "chatcmpl-two-choices",
+        "object": "chat.completion.chunk",
+        "created": 1,
+        "model": "accounts/fireworks/models/deepseek-v4p1-flash",
+        "choices": [{"index": index, "delta": {"role": "assistant", "content": content}, "finish_reason": None}],
+    }
+    return f"data: {json.dumps(chunk)}"
+
+
+def _fireworks_n2_lines_without_finish_reasons_through_done_then(failure: Exception) -> Iterator[str]:
+    yield _fireworks_sse_line_for_choice(0, "Hi")
+    yield _fireworks_sse_line_for_choice(1, "Hel")
+    yield "data: [DONE]"
+    raise failure
+
+
+@pytest.mark.asyncio
+@_SYNC_AND_ASYNC
+async def test_n2_stream_dropped_after_done_without_choice_finish_reasons_ends_cleanly(
+    logging_obj: Logging, sync_mode: bool
+):
+    logging_obj.update_environment_variables(litellm_params={}, optional_params={"n": 2})
+    lines: Final = _fireworks_n2_lines_without_finish_reasons_through_done_then(_httpx2_read_error())
+    wrapper: Final = CustomStreamWrapper(
+        completion_stream=FireworksAIChatCompletionStreamingHandler(
+            streaming_response=lines if sync_mode else _as_async(lines), sync_stream=sync_mode
+        ),
+        model="accounts/fireworks/models/deepseek-v4p1-flash",
+        custom_llm_provider="fireworks_ai",
+        logging_obj=logging_obj,
+    )
+
+    collected: Final = await _collect(wrapper, sync_mode)
+
+    assert "".join(chunk.choices[0].delta.content or "" for chunk in collected if chunk.choices) == "HiHel"
+
 def _generic_chunks_with_a_second_content_block_then(failure: Exception) -> Iterator[GenericStreamingChunk]:
     yield GenericStreamingChunk(text="Hi", is_finished=False, finish_reason="", usage=None, index=0)
     yield GenericStreamingChunk(text="!", is_finished=False, finish_reason="", usage=None, index=1)