From 8bdce62ec26225eff7a94e4ade20811203535254 Mon Sep 17 00:00:00 2001 From: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 13:52:12 -0700 Subject: [PATCH] fix(streaming): keep the synthetic stop for a stream that sent no chunks Gemini's OpenAI-compatible endpoint answers with only [DONE] when reasoning uses up max_tokens. Failing every openai/azure stream that sent nothing turned that 200 into a 500, while every other provider already kept the stop for an empty stream. Adds integration cells for Gemini's trailing usage chunk, multi-event LangGraph runs, a mocked deployment and the [DONE]-only answer. --- .../litellm_core_utils/streaming_handler.py | 3 +- ..._stream_clean_eof_without_finish_reason.py | 94 ++++++++++++++++--- .../test_streaming_handler.py | 37 ++++++-- 3 files changed, 114 insertions(+), 20 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 70bad6395ef..d867f926d87 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -71,7 +71,6 @@ _SYNC_ITER_EXHAUSTED: Final = object() _GCHUNK_FIELDS: Final[frozenset] = frozenset(GChunk.__annotations__) _CompleteStreamingResponse = ModelResponse | TextCompletionResponse | ModelResponseStream | None _USAGE_COST_HEADER_PROVIDERS: Final[frozenset[str]] = frozenset({LlmProviders.OPENROUTER.value}) -_OPENAI_AZURE_CHAT_PROVIDERS: Final[frozenset[str]] = frozenset({LlmProviders.OPENAI.value, LlmProviders.AZURE.value}) def _next_sync_or_exhausted(it: Iterator[object]) -> object: @@ -1908,7 +1907,7 @@ class CustomStreamWrapper: def _should_fail_on_clean_eof(self) -> bool: return ( not self.is_mock - and (self.custom_llm_provider in _OPENAI_AZURE_CHAT_PROVIDERS or self.uses_openai_chat_stream) + and self.uses_openai_chat_stream and self.received_finish_reason is None and self.intermittent_finish_reason is None and not self.received_any_finish_reason diff --git a/tests/integration/streaming/test_stream_clean_eof_without_finish_reason.py b/tests/integration/streaming/test_stream_clean_eof_without_finish_reason.py index f1dad7e8844..eea9f8e449e 100644 --- a/tests/integration/streaming/test_stream_clean_eof_without_finish_reason.py +++ b/tests/integration/streaming/test_stream_clean_eof_without_finish_reason.py @@ -11,13 +11,19 @@ from integration._support.responses_stream import data_frames, spend_row, sse_ev from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue -Ending = Literal["ending-complete", "ending-clean-eof", "ending-done-without-finish"] -Family = Literal["openai", "anthropic", "hosted_vllm"] +Ending = Literal["ending-complete", "ending-clean-eof", "ending-done-without-finish", "ending-done-only"] +Family = Literal["openai", "anthropic", "hosted_vllm", "gemini", "langgraph"] -_ENDINGS: Final[tuple[Ending, ...]] = ("ending-complete", "ending-clean-eof", "ending-done-without-finish") +_ENDINGS: Final[tuple[Ending, ...]] = ( + "ending-complete", + "ending-clean-eof", + "ending-done-without-finish", + "ending-done-only", +) _PARTIAL: Final = ('{"findings":[{"title":', '"unfinished') _BACKEND: Final = "gpt-4o-mini" _CLAUDE: Final = "claude-sonnet-4-5" +_GEMINI: Final = "gemini-2.5-flash" _PROVIDER_KEY: Final = "synthetic-provider-key" _CALL_ID_HEADER: Final = "x-litellm-call-id" _USAGE: Final = {"prompt_tokens": 9, "completion_tokens": 7, "total_tokens": 16} @@ -25,15 +31,19 @@ _TARGETS: Final[Mapping[Family, str]] = { "openai": "/chat/completions", "hosted_vllm": "/chat/completions", "anthropic": "/v1/messages", + "gemini": f"/models/{_GEMINI}:streamGenerateContent", + "langgraph": "/runs/stream", } _DEPLOYMENTS: Final[Mapping[Family, str]] = { "openai": f"openai/{_BACKEND}", "hosted_vllm": f"hosted_vllm/{_BACKEND}", "anthropic": f"anthropic/{_CLAUDE}", + "gemini": f"gemini/{_GEMINI}", + "langgraph": "langgraph/agent", } -def _data(frame: Mapping[str, JsonValue]) -> bytes: +def _data(frame: JsonValue) -> bytes: return b"data: " + json.dumps(frame).encode() + b"\n\n" @@ -52,6 +62,8 @@ def _chat_frames(identity: str, ending: Ending) -> tuple[bytes, ...]: return content case "ending-done-without-finish": return (*content, b"data: [DONE]\n\n") + case "ending-done-only": + return (b"data: [DONE]\n\n",) def _anthropic_frames(identity: str, ending: Ending) -> tuple[bytes, ...]: @@ -59,19 +71,55 @@ def _anthropic_frames(identity: str, ending: Ending) -> tuple[bytes, ...]: return (*content, CONTENT_TAIL) if ending == "ending-complete" else content +def _gemini_frames(ending: Ending) -> tuple[bytes, ...]: + """Gemini puts finishReason on the last text chunk and can follow it with a usage-only chunk.""" + finished: Final = ending == "ending-complete" + content: Final = tuple( + _data( + { + "candidates": [ + { + "content": {"parts": [{"text": part}], "role": "model"}, + "index": 0, + **({"finishReason": "STOP"} if finished and part == _PARTIAL[-1] else {}), + } + ] + } + ) + for part in _PARTIAL + ) + usage: Final = _data({"usageMetadata": {"promptTokenCount": 9, "candidatesTokenCount": 7, "totalTokenCount": 16}}) + return (*content, usage) if finished else content + + +def _langgraph_frames() -> tuple[bytes, ...]: + """LangGraph ends a run by closing the stream after its metadata event.""" + return ( + *(_data(["messages", [[{"type": "ai", "content": part}, {}]]]) for part in _PARTIAL), + _data(["metadata", {"run_id": "run-1"}]), + ) + + +def _frames(family: Family, identity: str, ending: Ending) -> tuple[bytes, ...]: + match family: + case "anthropic": + return _anthropic_frames(identity, ending) + case "gemini": + return _gemini_frames(ending) + case "langgraph": + return _langgraph_frames() + case "openai" | "hosted_vllm": + return _chat_frames(identity, ending) + + def _ending(request: Request) -> Ending: return next(ending for ending in _ENDINGS if ending.encode() in request.body) def _upstream(family: Family) -> Callable[[Request], Reply]: def respond(request: Request) -> Reply: - assert request.method == "POST" and request.target.endswith(_TARGETS[family]), request.target - identity: Final = f"upstream-{uuid.uuid4().hex}" - frames: Final = ( - _anthropic_frames(identity, _ending(request)) - if family == "anthropic" - else _chat_frames(identity, _ending(request)) - ) + assert request.method == "POST" and _TARGETS[family] in request.target, request.target + frames: Final = _frames(family, f"upstream-{uuid.uuid4().hex}", _ending(request)) return Reply(content_type="text/event-stream", chunks=frames) return answering_model_discovery(respond) @@ -150,7 +198,18 @@ def test_chat_stream_that_sends_done_without_a_finish_reason_reaches_the_client_ assert spend_row(call_id)["status"] == "failure" -@pytest.mark.parametrize("family", ["openai", "anthropic"]) +def test_chat_stream_that_sends_only_done_still_completes_empty(gateway: Gateway) -> None: + """Gemini's OpenAI-compatible endpoint answers with only [DONE] when reasoning uses up max_tokens.""" + with wire_server(_upstream("openai")) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_DEPLOYMENTS["openai"], api_base=wire.url, api_key=_PROVIDER_KEY) + text, call_id = _chat(gateway, model, "ending-done-only") + assert _chat_text(text) == "", text + assert _finish_reasons(text) == ("stop",), text + assert _chat_errors(text) == (), text + assert spend_row(call_id)["status"] == "success" + + +@pytest.mark.parametrize("family", ["openai", "anthropic", "gemini", "langgraph"]) def test_chat_stream_with_a_provider_finish_reason_still_completes(gateway: Gateway, family: Family) -> None: with wire_server(_upstream(family)) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=_DEPLOYMENTS[family], api_base=wire.url, api_key=_PROVIDER_KEY) @@ -162,6 +221,17 @@ def test_chat_stream_with_a_provider_finish_reason_still_completes(gateway: Gate assert spend_row(call_id)["status"] == "success" +def test_chat_stream_from_a_mock_response_deployment_still_completes(gateway: Gateway) -> None: + """mock_completion streams never carry a finish reason, so a mocked deployment must keep its synthetic stop.""" + with gateway.scenario() as scenario: + model: Final = scenario.model(model=_DEPLOYMENTS["openai"], mock_response="".join(_PARTIAL)) + text, call_id = _chat(gateway, model, "ending-clean-eof") + assert _chat_text(text) == "".join(_PARTIAL), text + assert _finish_reasons(text) == ("stop",), text + assert _chat_errors(text) == (), text + assert spend_row(call_id)["status"] == "success" + + @pytest.mark.parametrize(("ending", "terminal"), [("ending-complete", True), ("ending-clean-eof", False)]) def test_messages_over_a_chat_bridged_deployment_reports_end_turn_only_for_a_finished_stream( gateway: Gateway, ending: Ending, terminal: bool diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 2d1c2a581cf..0586af6e72b 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -1651,6 +1651,37 @@ def test_sync_openai_compatible_chat_clean_eof_without_finish_reason_raises_mids assert response.received_finish_reason is None +async def _collect(wrapper: CustomStreamWrapper, sync_mode: bool) -> list[ModelResponseStream]: + if sync_mode: + return list(wrapper) + return [chunk async for chunk in wrapper] + + +@pytest.mark.parametrize("custom_llm_provider", ["openai", "azure", "hosted_vllm"]) +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_stream_that_sends_no_chunks_keeps_the_default_stop( + logging_obj: Logging, custom_llm_provider: str, sync_mode: bool +): + """Gemini's OpenAI-compatible endpoint answers with only [DONE] when reasoning uses up max_tokens.""" + + async def no_chunks() -> AsyncIterator[ModelResponseStream]: + for chunk in (): + yield chunk + + wrapper: Final = CustomStreamWrapper( + completion_stream=iter(()) if sync_mode else no_chunks(), + model="gemini-2.5-flash", + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, + stream_options={"include_usage": True}, + ) + + received: Final = await _collect(wrapper, sync_mode) + + assert [c.choices[0].finish_reason for c in received if c.choices and c.choices[0].finish_reason] == ["stop"] + + _ANTHROPIC_PARTIAL_EVENTS: Final[tuple[dict[str, object], ...]] = ( { "type": "message_start", @@ -1701,12 +1732,6 @@ def _anthropic_stream( ) -async def _collect(wrapper: CustomStreamWrapper, sync_mode: bool) -> list[ModelResponseStream]: - if sync_mode: - return list(wrapper) - return [chunk async for chunk in wrapper] - - @pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.asyncio async def test_anthropic_stream_that_ends_before_message_delta_fails_with_the_eof_message(