fix(streaming): exact-match the SSE [DONE] marker in BaseModelResponseIterator (#47832)

* fix(streaming): exact-match the SSE [DONE] marker in BaseModelResponseIterator

* test(streaming): cover reasoning text and error frames that mention [DONE]

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>

litellm/llms/base_llm/base_model_iterator.py

@@ -3,6 +3,7 @@ from abc import abstractmethod
 from typing import TYPE_CHECKING, Final, cast
 
 import litellm
+from litellm.constants import STREAM_SSE_DONE_STRING
 
 if TYPE_CHECKING:
     import httpx
@@ -142,7 +143,8 @@ class BaseModelResponseIterator:
     def _handle_string_chunk(self, str_line: str) -> GenericStreamingChunk | ModelResponseStream:
         # chunk is a str at this point
         stripped_json_chunk: Final = BaseModelResponseIterator._string_to_dict_parser(str_line=str_line)
-        if "[DONE]" in str_line:
+        stripped_sse_chunk: Final = litellm.CustomStreamWrapper.strip_sse_data_from_chunk(str_line)
+        if stripped_sse_chunk is not None and stripped_sse_chunk.strip() == STREAM_SSE_DONE_STRING:
             return GenericStreamingChunk(
                 text="",
                 is_finished=True,

tests/unit/llms/base_llm/test_base_model_iterator.py

@@ -4,15 +4,19 @@ and non-string objects (e.g. Pydantic BaseModel events from the Responses API) p
 """
 
 import json
-from collections.abc import AsyncIterator
+from collections.abc import AsyncIterator, Mapping
+from itertools import chain
 from typing import Final
 
 import httpx
 import pytest
 from pydantic import BaseModel
 
+from litellm.litellm_core_utils.litellm_logging import Logging
+from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
 from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator, convert_model_response_to_streaming
 from litellm.llms.openai.chat.gpt_transformation import OpenAIChatCompletionStreamingHandler
+from litellm.llms.openai.common_utils import OpenAIError
 from litellm.types.utils import Choices, GenericStreamingChunk, Message, ModelResponse
 
 
@@ -336,3 +340,98 @@ def test_fake_stream_chunk_indexes_thinking_blocks_by_position_without_touching_
         {"type": "thinking", "thinking": "Step 2.", "signature": "sig-b", "index": 2},
     ]
     assert [block.get("index") for block in response.choices[0].message.thinking_blocks] == [None, None, None]
+
+
+def _sse_chunk(delta: Mapping[str, object], finish_reason: str | None = None) -> str:
+    return "data: " + json.dumps(
+        {
+            "id": "chatcmpl-done",
+            "object": "chat.completion.chunk",
+            "created": 1700000000,
+            "model": "mistral-small-latest",
+            "choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}],
+        }
+    )
+
+
+async def _aiter_lines(lines: tuple[str, ...]) -> AsyncIterator[str]:
+    for line in lines:
+        yield line
+
+
+@pytest.mark.parametrize("sync_mode", [True, False])
+@pytest.mark.asyncio
+async def test_done_inside_streamed_json_is_content_not_the_end_of_stream(sync_mode: bool) -> None:
+    lines: Final = (
+        _sse_chunk({"role": "assistant", "reasoning_content": "Mark it [DONE] once tests pass"}),
+        _sse_chunk({"content": "Checklist: "}),
+        _sse_chunk({"content": "- [DONE] write tests"}),
+        _sse_chunk({"content": " - [ ] ship it"}),
+        _sse_chunk(
+            {
+                "tool_calls": [
+                    {
+                        "index": 0,
+                        "id": "call_1",
+                        "type": "function",
+                        "function": {"name": "set_status", "arguments": ""},
+                    }
+                ]
+            }
+        ),
+        _sse_chunk({"tool_calls": [{"index": 0, "function": {"arguments": '{"status": '}}]}),
+        _sse_chunk({"tool_calls": [{"index": 0, "function": {"arguments": '"[DONE]"'}}]}),
+        _sse_chunk({"tool_calls": [{"index": 0, "function": {"arguments": ', "n": 1}'}}]}),
+        _sse_chunk({}, "tool_calls"),
+        "data: [DONE]",
+    )
+    wrapper: Final = CustomStreamWrapper(
+        completion_stream=OpenAIChatCompletionStreamingHandler(
+            streaming_response=iter(lines) if sync_mode else _aiter_lines(lines), sync_stream=sync_mode
+        ),
+        model="mistral/mistral-small-latest",
+        custom_llm_provider="mistral",
+        logging_obj=Logging(
+            model="mistral/mistral-small-latest",
+            messages=[{"role": "user", "content": "hi"}],
+            stream=True,
+            call_type="completion",
+            start_time=1_790_000_000.0,
+            litellm_call_id="call-id",
+            function_id="function-id",
+        ),
+    )
+
+    chunks: Final = list(wrapper) if sync_mode else [chunk async for chunk in wrapper]
+
+    deltas: Final = [chunk.choices[0].delta for chunk in chunks if chunk.choices]
+    tool_calls: Final = tuple(chain.from_iterable(delta.tool_calls or () for delta in deltas))
+    assert "".join(getattr(delta, "reasoning_content", None) or "" for delta in deltas) == (
+        "Mark it [DONE] once tests pass"
+    )
+    assert "".join(delta.content or "" for delta in deltas) == "Checklist: - [DONE] write tests - [ ] ship it"
+    assert json.loads("".join(tool_call.function.arguments or "" for tool_call in tool_calls)) == {
+        "status": "[DONE]",
+        "n": 1,
+    }
+    assert [chunk.choices[0].finish_reason for chunk in chunks if chunk.choices and chunk.choices[0].finish_reason] == [
+        "tool_calls"
+    ]
+
+
+@pytest.mark.parametrize("line", ["data: [DONE]", "data:[DONE]", "[DONE]", "data:  [DONE]", "data: [DONE]  "])
+def test_done_marker_ends_the_stream(line: str) -> None:
+    chunks: Final = list(BaseModelResponseIterator(streaming_response=iter((line,)), sync_stream=True))
+
+    assert chunks == [
+        GenericStreamingChunk(text="", is_finished=True, finish_reason="stop", usage=None, index=0, tool_use=None)
+    ]
+
+
+def test_error_frame_mentioning_done_raises_instead_of_ending_the_stream() -> None:
+    line: Final = 'data: {"error": {"message": "prompt too long [DONE]", "code": 400}}'
+
+    with pytest.raises(OpenAIError, match=r"prompt too long \[DONE\]") as error:
+        list(OpenAIChatCompletionStreamingHandler(streaming_response=iter((line,)), sync_stream=True))
+
+    assert error.value.status_code == 400