diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py
index 857aa2b4fef..7f29d646ffe 100644
--- a/litellm/litellm_core_utils/streaming_handler.py
+++ b/litellm/litellm_core_utils/streaming_handler.py
@@ -349,11 +349,14 @@ class CustomStreamWrapper:
make_call: Callable | None = None,
_response_headers: dict | httpx.Headers | None = None,
count_prompt_tokens: Callable[[], int] | None = None,
+ is_mock: bool = False,
):
self.model = model
self.make_call = make_call
self.count_prompt_tokens = count_prompt_tokens
self.custom_llm_provider = custom_llm_provider
+ self.is_mock = is_mock
+ self.uses_openai_chat_stream = False
self.logging_obj: LiteLLMLoggingObject = logging_obj
self.completion_stream = completion_stream
self.sent_first_chunk = False
@@ -376,6 +379,7 @@ class CustomStreamWrapper:
self._upstream_provider: str | None = None
self.received_finish_reason: str | None = None
self.intermittent_finish_reason: str | None = None # finish reasons that show up mid-stream
+ self.received_any_finish_reason = False # a later chunk can reset intermittent_finish_reason
self.special_tokens = [
"<|assistant|>",
"<|system|>",
@@ -1597,6 +1601,7 @@ class CustomStreamWrapper:
if response_obj["is_finished"]:
self.received_finish_reason = response_obj["finish_reason"]
else: # openai / azure chat model
+ self.uses_openai_chat_stream = True
if isinstance(chunk, ChatCompletionChunk):
self._openai_handler_receives_chat_chunks = True
if any(str(choice.finish_reason) == "abort" for choice in chunk.choices or ()):
@@ -1617,6 +1622,7 @@ class CustomStreamWrapper:
return _ProviderChunkEarlyReturn(None)
completion_obj["content"] = response_obj["text"]
self.intermittent_finish_reason = response_obj.get("finish_reason", None)
+ self.received_any_finish_reason |= response_obj["finish_reason"] is not None
if response_obj["is_finished"]:
if response_obj["finish_reason"] == "error":
raise Exception(
@@ -1996,6 +2002,22 @@ class CustomStreamWrapper:
model_response.choices[0].delta.content = close_tag
return model_response
+ def _should_fail_on_clean_eof(self) -> bool:
+ return (
+ not self.is_mock
+ 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
+ )
+
+ @staticmethod
+ def _stream_eof_error() -> httpx.RemoteProtocolError:
+ return httpx.RemoteProtocolError(
+ "Stream ended before a finish_reason was received",
+ request=httpx.Request(method="POST", url="https://docs.litellm.ai/docs"),
+ )
+
def _final_chunk_for_logging(self, processed_chunk: "ModelResponseStream") -> "ModelResponseStream | None":
if self.received_finish_reason is not None or self.intermittent_finish_reason is not None:
return processed_chunk
@@ -2181,6 +2203,13 @@ class CustomStreamWrapper:
return response
except StopIteration:
+ if self._should_fail_on_clean_eof():
+ error: Final = self._stream_eof_error()
+ eof_traceback: Final = traceback.format_exc()
+ self._record_partial_usage_for_failure()
+ threading.Thread(target=self.logging_obj.failure_handler, args=(error, eof_traceback)).start()
+ self.close()
+ self._handle_stream_fallback_error(error)
if self.sent_last_chunk is True:
complete_streaming_response: Final = self._build_complete_streaming_response()
response = self._new_usage_chunk()
@@ -2384,6 +2413,9 @@ class CustomStreamWrapper:
continue
return usage_free_chunk
except (StopAsyncIteration, StopIteration):
+ if self._should_fail_on_clean_eof():
+ await self.aclose()
+ self._log_stream_failure_and_raise(self._stream_eof_error())
return await self._finalize_completed_stream(cache_hit=cache_hit)
except httpx.TimeoutException as e: # if httpx read timeout error occues
traceback_exception = traceback.format_exc()
diff --git a/litellm/llms/bedrock/chat/agentcore/transformation.py b/litellm/llms/bedrock/chat/agentcore/transformation.py
index 18212b7390f..2982ae6cd08 100644
--- a/litellm/llms/bedrock/chat/agentcore/transformation.py
+++ b/litellm/llms/bedrock/chat/agentcore/transformation.py
@@ -544,6 +544,17 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
return AgentCoreParsedResponse(content=content, usage=usage_data, final_message=final_message)
+ @staticmethod
+ def _message_stop_chunk(model: str) -> ModelResponseStream:
+ chunk: Final = ModelResponseStream(
+ id=f"chatcmpl-{uuid.uuid4()}",
+ created=0,
+ model=model,
+ object="chat.completion.chunk",
+ )
+ chunk.choices = [StreamingChoices(finish_reason="stop", index=0, delta=Delta())]
+ return chunk
+
def _stream_agentcore_response_sync(
self,
response: httpx.Response,
@@ -553,6 +564,8 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
Internal sync generator that parses SSE and yields ModelResponse chunks.
"""
buffer = ""
+ sent_finish = False # rebind-ok: flips when a finish chunk goes out
+ message_stopped = False # rebind-ok: flips on messageStop, the only end some runtimes send
for text_chunk in response.iter_text():
buffer += text_chunk
@@ -624,7 +637,9 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
total_tokens=usage_data.get("totalTokens", 0),
),
)
+ sent_finish = True
yield chunk
+ message_stopped |= "messageStop" in event_payload
# Process final message
if "message" in data_obj and isinstance(data_obj["message"], dict):
@@ -641,12 +656,16 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
delta=Delta(),
)
]
+ sent_finish = True
yield chunk
except json.JSONDecodeError:
verbose_logger.debug("Skipping non-JSON SSE line: %s", line[:100])
continue
+ if message_stopped and not sent_finish:
+ yield self._message_stop_chunk(model)
+
def get_sync_custom_stream_wrapper(
self,
model: str,
@@ -775,6 +794,8 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
Internal async generator that parses SSE and yields ModelResponse chunks.
"""
buffer = ""
+ sent_finish = False # rebind-ok: flips when a finish chunk goes out
+ message_stopped = False # rebind-ok: flips on messageStop, the only end some runtimes send
async for text_chunk in response.aiter_text():
buffer += text_chunk
@@ -846,7 +867,9 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
total_tokens=usage_data.get("totalTokens", 0),
),
)
+ sent_finish = True
yield chunk
+ message_stopped |= "messageStop" in event_payload
# Process final message
if "message" in data_obj and isinstance(data_obj["message"], dict):
@@ -863,12 +886,16 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
delta=Delta(),
)
]
+ sent_finish = True
yield chunk
except json.JSONDecodeError:
verbose_logger.debug("Skipping non-JSON SSE line: %s", line[:100])
continue
+ if message_stopped and not sent_finish:
+ yield self._message_stop_chunk(model)
+
async def get_async_custom_stream_wrapper(
self,
model: str,
diff --git a/litellm/llms/langgraph/chat/sse_iterator.py b/litellm/llms/langgraph/chat/sse_iterator.py
index 6fe355ad98c..380c7893c77 100644
--- a/litellm/llms/langgraph/chat/sse_iterator.py
+++ b/litellm/llms/langgraph/chat/sse_iterator.py
@@ -41,8 +41,9 @@ class LangGraphSSEStreamIterator:
return self
def __aiter__(self):
- """Initialize async iteration."""
- self.async_line_iterator = aiter_sse_lines(self.response.aiter_text())
+ """Initialize async iteration once; CustomStreamWrapper re-enters `async for` on every chunk."""
+ if self.async_line_iterator is None:
+ self.async_line_iterator = aiter_sse_lines(self.response.aiter_text())
return self
def _parse_sse_line(self, line: str) -> ModelResponseStream | None:
diff --git a/litellm/main.py b/litellm/main.py
index 3a97fc7d5e6..78767b03eb7 100644
--- a/litellm/main.py
+++ b/litellm/main.py
@@ -1237,6 +1237,7 @@ def mock_completion(
custom_llm_provider="openai",
logging_obj=logging,
count_prompt_tokens=count_prompt_tokens,
+ is_mock=True,
)
return CustomStreamWrapper(
completion_stream=mock_completion_streaming_obj(
@@ -1246,6 +1247,7 @@ def mock_completion(
custom_llm_provider="openai",
logging_obj=logging,
count_prompt_tokens=count_prompt_tokens,
+ is_mock=True,
)
if isinstance(mock_response, litellm.MockException):
raise mock_response
diff --git a/tests/integration/streaming/test_decision_model_buffering.py b/tests/integration/streaming/test_decision_model_buffering.py
index a2594fd2b9b..d6fea067d69 100644
--- a/tests/integration/streaming/test_decision_model_buffering.py
+++ b/tests/integration/streaming/test_decision_model_buffering.py
@@ -229,11 +229,17 @@ def _frames(identity: str, endpoint: str) -> tuple[str, ...]:
'object': 'chat.completion.chunk',
'created': 1,
'model': 'scripted-chat',
- 'choices': [{'index': 0, 'delta': {'content': text}, 'finish_reason': None}],
+ 'choices': [
+ {
+ 'index': 0,
+ 'delta': {'content': text},
+ 'finish_reason': 'stop' if index == len(pieces) - 1 else None,
+ }
+ ],
}
)
}"
- for text in pieces
+ for index, text in enumerate(pieces)
),
"data: [DONE]",
)
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
new file mode 100644
index 00000000000..f3b427d91c4
--- /dev/null
+++ b/tests/integration/streaming/test_stream_clean_eof_without_finish_reason.py
@@ -0,0 +1,564 @@
+import asyncio
+import itertools
+import json
+import uuid
+from collections import Counter
+from collections.abc import Callable, Iterator, Mapping
+from concurrent.futures import ThreadPoolExecutor
+from contextlib import ExitStack
+from pathlib import Path
+from typing import Final, Literal
+
+import anthropic
+import httpx
+import openai
+import pytest
+from integration._support.anthropic_sse import CONTENT_BLOCK_START, CONTENT_TAIL, message_start, text_delta
+from integration._support.client import Gateway, Scenario, eventually, string_value
+from integration._support.database import read_rows, scratch_database
+from integration._support.database_relay import database_relay, plaintext_database_url
+from integration._support.openai_wire import answering_model_discovery, posted_targets
+from integration._support.process import owned_proxy, owned_proxy_process
+from integration._support.responses_stream import data_frames, spend_row, sse_events
+from integration._support.wire import Reply, Request, Wire, wire_server
+from integration.providers._count_tokens_system_lift import (
+ anthropic_client,
+ async_anthropic_client,
+ async_openai_client,
+ openai_client,
+)
+from pydantic import JsonValue
+
+Ending = Literal["ending-complete", "ending-clean-eof", "ending-done-without-finish", "ending-done-only"]
+Family = Literal["openai", "anthropic", "hosted_vllm", "gemini", "langgraph", "agentcore"]
+
+_ENDINGS: Final[tuple[Ending, ...]] = (
+ "ending-complete",
+ "ending-clean-eof",
+ "ending-done-without-finish",
+ "ending-done-only",
+)
+_PARTIAL: Final = ('{"findings":[{"title":', '"unfinished')
+_BACKEND: Final = "gpt-5.6"
+_CLAUDE: Final = "claude-sonnet-5-5"
+_GEMINI: Final = "gemini-3.8-flash"
+_AGENT_RUNTIME: Final = "arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/integration-agent"
+_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}
+_TARGETS: Final[Mapping[Family, str]] = {
+ "openai": "/chat/completions",
+ "hosted_vllm": "/chat/completions",
+ "anthropic": "/v1/messages",
+ "gemini": f"/models/{_GEMINI}:streamGenerateContent",
+ "langgraph": "/runs/stream",
+ "agentcore": "/runtimes/",
+}
+_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",
+ "agentcore": f"bedrock/agentcore/{_AGENT_RUNTIME}",
+}
+
+
+def _data(frame: JsonValue) -> bytes:
+ return b"data: " + json.dumps(frame).encode() + b"\n\n"
+
+
+def _chat_frames(identity: str, ending: Ending) -> tuple[bytes, ...]:
+ chunk: Final = {"id": identity, "object": "chat.completion.chunk", "created": 1, "model": _BACKEND}
+ content: Final = (
+ _data({**chunk, "choices": [{"index": 0, "delta": {"role": "assistant", "content": _PARTIAL[0]}}]}),
+ _data({**chunk, "choices": [{"index": 0, "delta": {"content": _PARTIAL[1]}, "finish_reason": None}]}),
+ )
+ match ending:
+ case "ending-complete":
+ finish: Final = _data({**chunk, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]})
+ usage: Final = _data({**chunk, "choices": [], "usage": _USAGE})
+ return (*content, finish + usage + b"data: [DONE]\n\n")
+ case "ending-clean-eof":
+ 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, ...]:
+ content: Final = (message_start(identity, _CLAUDE), CONTENT_BLOCK_START, *map(text_delta, _PARTIAL))
+ 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 _event(name: str, frame: JsonValue) -> bytes:
+ return f"event: {name}\n".encode() + _data(frame)
+
+
+def _langgraph_frames() -> tuple[bytes, ...]:
+ """A LangGraph server's run stream in its default mode, its run metadata first and then the state after each step"""
+ question: Final = {"type": "human", "content": "Answer in two steps."}
+ answers: Final = tuple({"type": "ai", "content": part} for part in _PARTIAL)
+ return (
+ _event("metadata", {"run_id": "run-1", "attempt": 1}),
+ *(_event("values", {"messages": [question, *answers[:step]]}) for step in range(len(answers) + 1)),
+ )
+
+
+def _agentcore_frames(ending: Ending) -> tuple[bytes, ...]:
+ """A runtime that forwards Bedrock Converse events without the metadata event, so its turn ends on messageStop"""
+ start: Final = _data({"event": {"messageStart": {"role": "assistant"}}})
+ content: Final = tuple(
+ _data({"event": {"contentBlockDelta": {"delta": {"text": part}, "contentBlockIndex": 0}}}) for part in _PARTIAL
+ )
+ end: Final = (
+ _data({"event": {"contentBlockStop": {"contentBlockIndex": 0}}}),
+ _data({"event": {"messageStop": {"stopReason": "end_turn"}}}),
+ )
+ return (start, *content, *end) if ending == "ending-complete" else (start, *content)
+
+
+def _frames(family: Family, identity: str, ending: Ending) -> tuple[bytes, ...]:
+ match family:
+ case "agentcore":
+ return _agentcore_frames(ending)
+ 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:
+ endings: Final = [ending for ending in _ENDINGS if ending.encode() in request.body]
+ assert len(endings) == 1, request.body
+ return endings[0]
+
+
+def _upstream(family: Family, ending: Ending | None = None) -> Callable[[Request], Reply]:
+ def respond(request: Request) -> Reply:
+ assert request.method == "POST" and _TARGETS[family] in request.target, request.target
+ frames: Final = _frames(family, f"upstream-{uuid.uuid4().hex}", ending or _ending(request))
+ return Reply(content_type="text/event-stream", chunks=frames)
+
+ return answering_model_discovery(respond)
+
+
+def _deploy(scenario: Scenario, family: Family, wire: Wire) -> str:
+ if family == "agentcore":
+ return scenario.model(
+ model=_DEPLOYMENTS[family],
+ aws_bedrock_runtime_endpoint=wire.url,
+ aws_access_key_id="AKIAINTEGRATIONAGENT",
+ aws_secret_access_key="synthetic-agentcore-secret",
+ aws_region_name="us-west-2",
+ )
+ return scenario.model(model=_DEPLOYMENTS[family], api_base=wire.url, api_key=_PROVIDER_KEY)
+
+
+def _prompt(ending: Ending) -> str:
+ return f"Return JSON {ending} {uuid.uuid4().hex}"
+
+
+def _post(gateway: Gateway, path: str, body: Mapping[str, JsonValue]) -> tuple[str, str]:
+ response: Final = gateway.request("POST", path, {**body, "stream": True, "num_retries": 0})
+ assert response.status_code == 200, response.text
+ return response.text, response.headers[_CALL_ID_HEADER]
+
+
+def _chat(gateway: Gateway, model: str, ending: Ending) -> tuple[str, str]:
+ return _post(
+ gateway,
+ "/v1/chat/completions",
+ {
+ "model": model,
+ "messages": [{"role": "user", "content": _prompt(ending)}],
+ "stream_options": {"include_usage": True},
+ },
+ )
+
+
+def _choices(text: str) -> Iterator[Mapping[str, JsonValue]]:
+ for frame in data_frames(text):
+ yield from (choice for choice in frame.get("choices") or [] if isinstance(choice, dict))
+
+
+def _finish_reasons(text: str) -> tuple[str, ...]:
+ return tuple(string_value(choice["finish_reason"]) for choice in _choices(text) if choice.get("finish_reason"))
+
+
+def _chat_text(text: str) -> str:
+ deltas: Final = (choice.get("delta") for choice in _choices(text))
+ return "".join(string_value(delta.get("content") or "") for delta in deltas if isinstance(delta, dict))
+
+
+def _chat_errors(text: str) -> tuple[Mapping[str, JsonValue], ...]:
+ return tuple(frame for frame in data_frames(text) if "error" in frame)
+
+
+def _event_types(text: str) -> tuple[str, ...]:
+ return tuple(event for event, _ in sse_events(text))
+
+
+@pytest.mark.parametrize("family", ["openai", "anthropic", "agentcore"])
+def test_chat_stream_that_ends_cleanly_without_a_finish_reason_reaches_the_client_as_an_error(
+ gateway: Gateway, family: Family
+) -> None:
+ with wire_server(_upstream(family)) as wire, gateway.scenario() as scenario:
+ model: Final = _deploy(scenario, family, wire)
+ text, call_id = _chat(gateway, model, "ending-clean-eof")
+ assert _chat_text(text) == "".join(_PARTIAL), text
+ assert _finish_reasons(text) == (), text
+ assert len(_chat_errors(text)) == 1, text
+ assert spend_row(call_id)["status"] == "failure"
+
+
+def test_chat_stream_that_sends_done_without_a_finish_reason_reaches_the_client_as_an_error(gateway: Gateway) -> None:
+ 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-without-finish")
+ assert _chat_text(text) == "".join(_PARTIAL), text
+ assert _finish_reasons(text) == (), text
+ assert len(_chat_errors(text)) == 1, text
+ assert spend_row(call_id)["status"] == "failure"
+
+
+def test_chat_stream_that_sends_only_done_still_completes_empty(gateway: Gateway) -> None:
+ """An upstream can answer with only [DONE]: Gemini's OpenAI-compatible endpoint did on 2026-10-20 when
+ reasoning used 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, _ = _chat(gateway, model, "ending-done-only")
+ assert _chat_text(text) == "", text
+ assert _finish_reasons(text) == ("stop",), text
+ assert _chat_errors(text) == (), text
+ assert text.rstrip().endswith("data: [DONE]"), text
+
+
+@pytest.mark.parametrize("family", ["openai", "anthropic", "gemini", "langgraph", "agentcore"])
+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 = _deploy(scenario, family, wire)
+ text, call_id = _chat(gateway, model, "ending-complete")
+ assert _chat_text(text) == "".join(_PARTIAL), text
+ assert _finish_reasons(text) == ("stop",), text
+ assert _chat_errors(text) == (), text
+ assert text.rstrip().endswith("data: [DONE]"), text
+ 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
+) -> None:
+ with wire_server(_upstream("hosted_vllm")) as wire, gateway.scenario() as scenario:
+ model: Final = scenario.model(model=_DEPLOYMENTS["hosted_vllm"], api_base=wire.url, api_key=_PROVIDER_KEY)
+ text, call_id = _post(
+ gateway,
+ "/v1/messages",
+ {"model": model, "max_tokens": 64, "messages": [{"role": "user", "content": _prompt(ending)}]},
+ )
+ types: Final = _event_types(text)
+ stop_reasons: Final = tuple(
+ data["delta"]["stop_reason"]
+ for event, data in sse_events(text)
+ if event == "message_delta" and isinstance(data.get("delta"), dict)
+ )
+ assert types[0] == "message_start", types
+ assert ("message_stop" in types) is terminal, types
+ assert ("error" in types) is not terminal, types
+ assert stop_reasons == (("end_turn",) if terminal else ()), types
+ assert spend_row(call_id)["status"] == ("success" if terminal else "failure")
+
+
+@pytest.mark.parametrize(("ending", "terminal"), [("ending-complete", True), ("ending-clean-eof", False)])
+def test_responses_over_a_chat_bridged_anthropic_deployment_completes_only_a_finished_stream(
+ gateway: Gateway, ending: Ending, terminal: bool
+) -> None:
+ with wire_server(_upstream("anthropic")) as wire, gateway.scenario() as scenario:
+ model: Final = scenario.model(model=_DEPLOYMENTS["anthropic"], api_base=wire.url, api_key=_PROVIDER_KEY)
+ text, call_id = _post(gateway, "/v1/responses", {"model": model, "input": _prompt(ending)})
+ types: Final = tuple(string_value(frame.get("type") or "") for frame in data_frames(text))
+ assert types[:1] == ("response.created",), text
+ assert ("response.completed" in types) is terminal, text
+ assert (types[-1] in {"error", "response.failed"}) is not terminal, text
+ assert spend_row(call_id)["status"] == ("success" if terminal else "failure")
+
+
+def _drain_openai_sdk(gateway: Gateway, model: str, sync_mode: bool) -> None:
+ messages: Final = [{"role": "user", "content": _prompt("ending-clean-eof")}]
+ if sync_mode:
+ for _ in openai_client(gateway).chat.completions.create(model=model, messages=messages, stream=True):
+ pass
+ return
+
+ async def drain() -> None:
+ stream: Final = await async_openai_client(gateway).chat.completions.create(
+ model=model, messages=messages, stream=True
+ )
+ async for _ in stream:
+ pass
+
+ asyncio.run(drain())
+
+
+def _drain_anthropic_sdk(gateway: Gateway, model: str, sync_mode: bool) -> None:
+ messages: Final = [{"role": "user", "content": _prompt("ending-clean-eof")}]
+ if sync_mode:
+ for _ in anthropic_client(gateway).messages.create(model=model, max_tokens=64, messages=messages, stream=True):
+ pass
+ return
+
+ async def drain() -> None:
+ stream: Final = await async_anthropic_client(gateway).messages.create(
+ model=model, max_tokens=64, messages=messages, stream=True
+ )
+ async for _ in stream:
+ pass
+
+ asyncio.run(drain())
+
+
+@pytest.mark.parametrize("sync_mode", [True, False])
+def test_openai_sdk_stream_cut_before_a_finish_reason_raises(gateway: Gateway, sync_mode: bool) -> None:
+ with wire_server(_upstream("openai")) as wire, gateway.scenario() as scenario:
+ model: Final = _deploy(scenario, "openai", wire)
+ with pytest.raises(openai.APIError, match="Stream ended before a finish_reason was received"):
+ _drain_openai_sdk(gateway, model, sync_mode)
+
+
+@pytest.mark.parametrize("sync_mode", [True, False])
+def test_anthropic_sdk_messages_stream_cut_before_a_finish_reason_raises(gateway: Gateway, sync_mode: bool) -> None:
+ with wire_server(_upstream("hosted_vllm")) as wire, gateway.scenario() as scenario:
+ model: Final = _deploy(scenario, "hosted_vllm", wire)
+ with pytest.raises(anthropic.APIError, match="Stream ended before a finish_reason was received"):
+ _drain_anthropic_sdk(gateway, model, sync_mode)
+
+
+def test_responses_stream_cut_on_its_deployment_finishes_on_the_request_fallback(gateway: Gateway) -> None:
+ with (
+ wire_server(_upstream("anthropic")) as primary,
+ wire_server(_upstream("anthropic", "ending-complete")) as backup,
+ gateway.scenario() as scenario,
+ ):
+ model: Final = _deploy(scenario, "anthropic", primary)
+ fallback: Final = _deploy(scenario, "anthropic", backup)
+ text, _ = _post(
+ gateway, "/v1/responses", {"model": model, "input": _prompt("ending-clean-eof"), "fallbacks": [fallback]}
+ )
+ types: Final = tuple(string_value(frame.get("type") or "") for frame in data_frames(text))
+ assert "response.completed" in types and "response.failed" not in types, text
+ assert (len(posted_targets(primary)), len(posted_targets(backup))) == (1, 1), text
+
+
+def test_chat_stream_cut_after_its_content_is_not_replayed_on_the_request_fallback(gateway: Gateway) -> None:
+ with (
+ wire_server(_upstream("openai")) as primary,
+ wire_server(_upstream("openai", "ending-complete")) as backup,
+ gateway.scenario() as scenario,
+ ):
+ model: Final = _deploy(scenario, "openai", primary)
+ fallback: Final = _deploy(scenario, "openai", backup)
+ text, _ = _post(
+ gateway,
+ "/v1/chat/completions",
+ {
+ "model": model,
+ "messages": [{"role": "user", "content": _prompt("ending-clean-eof")}],
+ "fallbacks": [fallback],
+ },
+ )
+ assert _chat_text(text) == "".join(_PARTIAL), text
+ assert len(_chat_errors(text)) == 1, text
+ assert posted_targets(backup) == (), text
+
+
+Route = Literal["/v1/chat/completions", "/v1/messages", "/v1/responses"]
+
+
+def _route_body(route: Route, model: str, ending: Ending) -> Mapping[str, JsonValue]:
+ match route:
+ case "/v1/chat/completions":
+ return {"model": model, "messages": [{"role": "user", "content": _prompt(ending)}]}
+ case "/v1/messages":
+ return {"model": model, "max_tokens": 64, "messages": [{"role": "user", "content": _prompt(ending)}]}
+ case "/v1/responses":
+ return {"model": model, "input": _prompt(ending)}
+
+
+def _finished(route: Route, text: str) -> bool:
+ match route:
+ case "/v1/chat/completions":
+ return _finish_reasons(text) == ("stop",) and _chat_errors(text) == ()
+ case "/v1/messages":
+ return "message_stop" in _event_types(text) and "error" not in _event_types(text)
+ case "/v1/responses":
+ return "response.completed" in tuple(string_value(frame.get("type") or "") for frame in data_frames(text))
+
+
+def _stream_outcome(gateway: Gateway, model: str, route: Route, ending: Ending) -> tuple[bool, str]:
+ text, call_id = _post(gateway, route, _route_body(route, model, ending))
+ return _finished(route, text), call_id
+
+
+_BRIDGED_FAMILY: Final[Mapping[Route, Family]] = {
+ "/v1/chat/completions": "openai",
+ "/v1/messages": "hosted_vllm",
+ "/v1/responses": "anthropic",
+}
+
+
+def test_concurrent_burst_ends_each_stream_by_whether_its_provider_finished_it(gateway: Gateway) -> None:
+ endings: Final[tuple[Ending, ...]] = ("ending-complete", "ending-clean-eof")
+ cases: Final = tuple(itertools.product(_BRIDGED_FAMILY, endings, range(5)))
+ with ExitStack() as wires, gateway.scenario() as scenario:
+ models: Final = {
+ route: _deploy(scenario, family, wires.enter_context(wire_server(_upstream(family))))
+ for route, family in _BRIDGED_FAMILY.items()
+ }
+ with ThreadPoolExecutor(max_workers=12) as pool:
+ outcomes: Final = tuple(
+ pool.map(lambda case: _stream_outcome(gateway, models[case[0]], case[0], case[1]), cases)
+ )
+ assert tuple(finished for finished, _ in outcomes) == tuple(
+ ending == "ending-complete" for _, ending, _ in cases
+ ), outcomes
+ assert tuple(spend_row(call_id)["status"] for _, call_id in outcomes) == tuple(
+ "success" if ending == "ending-complete" else "failure" for _, ending, _ in cases
+ )
+
+
+_SPEND_LOG_INSERT: Final = b'INSERT INTO "public"."LiteLLM_SpendLogs"'
+
+
+def _spend_statuses(database_url: str | None) -> Counter[tuple[str, str]]:
+ rows: Final = read_rows('SELECT litellm_call_id, status FROM "LiteLLM_SpendLogs"', (), database_url=database_url)
+ return Counter((string_value(row["litellm_call_id"]), string_value(row["status"])) for row in rows)
+
+
+def test_spend_log_outage_mid_burst_keeps_each_stream_ending_and_lands_every_row_once(
+ gateway: Gateway, tmp_path: Path
+) -> None:
+ endings: Final[tuple[Ending, ...]] = ("ending-complete", "ending-clean-eof")
+ cases: Final = tuple(itertools.product(_BRIDGED_FAMILY, endings))
+ with (
+ scratch_database() as database_url,
+ database_relay(database_url, _SPEND_LOG_INSERT) as (relay, relayed_url),
+ owned_proxy(gateway, tmp_path, {"DATABASE_URL": plaintext_database_url(relayed_url)}, workers=2) as candidate,
+ ExitStack() as wires,
+ candidate.scenario() as scenario,
+ ThreadPoolExecutor(max_workers=len(cases)) as pool,
+ ):
+ models: Final = {
+ route: _deploy(scenario, family, wires.enter_context(wire_server(_upstream(family))))
+ for route, family in _BRIDGED_FAMILY.items()
+ }
+
+ def burst() -> tuple[tuple[bool, str], ...]:
+ return tuple(pool.map(lambda case: _stream_outcome(candidate, models[case[0]], case[0], case[1]), cases))
+
+ relay.arm()
+ before: Final = burst()
+ assert relay.tripped.wait(90), "no spend log write reached the database after the first burst"
+ down: Final = not relay.reconnected.is_set()
+ during: Final = burst()
+ assert relay.reconnected.wait(90), "the proxy never reconnected to the database after the outage"
+ assert down and relay.refused >= 1, relay.refused
+ finished: Final = tuple(ending == "ending-complete" for _, ending in cases)
+ assert tuple(done for done, _ in before + during) == finished * 2, (before, during)
+ expected: Final = Counter(
+ (call_id, "success" if ending == "ending-complete" else "failure")
+ for (_, ending), (_, call_id) in zip(cases * 2, before + during, strict=True)
+ )
+ landed: Final = eventually(
+ lambda: _spend_statuses(database_url),
+ lambda rows: {call_id for call_id, _ in expected} <= {call_id for call_id, _ in rows},
+ seconds=120,
+ )
+ assert landed == expected
+
+
+def _attempt(gateway: Gateway, model: str, route: Route, ending: Ending) -> tuple[bool, str] | None:
+ """One stream's outcome, or None when the worker serving it was killed before the proxy finished the reply"""
+ try:
+ return _stream_outcome(gateway, model, route, ending)
+ except httpx.TransportError:
+ return None
+
+
+def test_worker_killed_mid_burst_leaves_the_proxy_ending_each_stream_by_whether_its_provider_finished_it(
+ gateway: Gateway, tmp_path: Path
+) -> None:
+ endings: Final[tuple[Ending, ...]] = ("ending-complete", "ending-clean-eof")
+ cases: Final = tuple(itertools.product(_BRIDGED_FAMILY, endings, range(4)))
+ probes: Final = tuple(itertools.product(_BRIDGED_FAMILY, endings))
+ with (
+ owned_proxy_process(gateway, tmp_path, {}, workers=2) as owned,
+ ExitStack() as wires,
+ owned.gateway.scenario() as scenario,
+ ):
+ candidate: Final = owned.gateway
+ models: Final = {
+ route: _deploy(scenario, family, wires.enter_context(wire_server(_upstream(family))))
+ for route, family in _BRIDGED_FAMILY.items()
+ }
+ workers: Final = eventually(owned.worker_pids, lambda pids: len(pids) == 2, seconds=60)
+ with owned.replacing_worker(min(workers)) as kill, ThreadPoolExecutor(max_workers=len(cases)) as pool:
+ futures: Final = [
+ pool.submit(_attempt, candidate, models[route], route, ending) for route, ending, _ in cases
+ ]
+ kill()
+ attempts: Final = tuple(future.result() for future in futures)
+ answered: Final = tuple(
+ (outcome[0], ending == "ending-complete")
+ for (_, ending, _), outcome in zip(cases, attempts, strict=True)
+ if outcome is not None
+ )
+ assert all(finished == expected for finished, expected in answered), attempts
+ after: Final = tuple(_stream_outcome(candidate, models[route], route, ending) for route, ending in probes)
+ assert tuple(finished for finished, _ in after) == tuple(ending == "ending-complete" for _, ending in probes)
+ statuses: Final = eventually(
+ lambda: _spend_statuses(candidate.database_url),
+ lambda rows: {call_id for _, call_id in after} <= {call_id for call_id, _ in rows},
+ seconds=120,
+ )
+ assert tuple(
+ statuses[(call_id, "success" if ending == "ending-complete" else "failure")]
+ for (_, ending), (_, call_id) in zip(probes, after, strict=True)
+ ) == (1,) * len(probes)
diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py
index 365ed3452cf..5c0c1642e01 100644
--- a/tests/unit/litellm_core_utils/test_streaming_handler.py
+++ b/tests/unit/litellm_core_utils/test_streaming_handler.py
@@ -32,8 +32,12 @@ from litellm.litellm_core_utils.streaming_handler import (
_ProviderChunkEarlyReturn,
_ProviderChunkParsed,
)
+from litellm.llms.anthropic.chat.handler import ModelResponseIterator as AnthropicModelResponseIterator
from litellm.llms.custom_httpx.http_handler import HTTPHandler
from litellm.llms.openai.chat.gpt_transformation import OpenAIChatCompletionStreamingHandler
+from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
+ ModelResponseIterator as GeminiModelResponseIterator,
+)
from litellm.responses.mcp.chat_completions_handler import _add_mcp_metadata_to_response
from litellm.types.utils import (
ChoiceLogprobs,
@@ -572,11 +576,6 @@ _GET_WEATHER_CALL: Final = {
"abcd",
0,
),
- (
- (_stream_chunk({"reasoning_content": "Still thinking"}),),
- "Still thinking",
- 0,
- ),
(
(_stream_chunk({"content": "hi"}), _stream_chunk({}, "stop")),
"hi",
@@ -588,7 +587,6 @@ _GET_WEATHER_CALL: Final = {
"reasoning_then_function_call",
"reasoning_cut_by_length",
"reasoning_after_answer",
- "stream_ends_mid_reasoning",
"answer_without_reasoning",
],
)
@@ -622,6 +620,26 @@ async def test_merged_reasoning_stream_keeps_think_tags_balanced(
)
+@pytest.mark.asyncio
+@pytest.mark.parametrize("sync_mode", [True, False])
+async def test_merged_reasoning_stream_cut_mid_reasoning_fails_instead_of_closing_the_think_block(
+ logging_obj: Logging, sync_mode: bool
+) -> None:
+ response: Final = CustomStreamWrapper(
+ completion_stream=ModelResponseListIterator(
+ model_responses=[_stream_chunk({"reasoning_content": "Still thinking"})]
+ ),
+ model="deepseek-reasoner",
+ logging_obj=logging_obj,
+ )
+ response.merge_reasoning_content_in_choices = True
+
+ with pytest.raises(MidStreamFallbackError) as excinfo:
+ await _collect(response, sync_mode)
+
+ assert excinfo.value.generated_content == "Still thinking"
+
+
def test_strip_sse_data_from_chunk():
"""Test the static method that strips 'data: ' prefix from SSE chunks"""
# Test with string inputs
@@ -2270,6 +2288,242 @@ async def test_mid_stream_error_keeps_credentials_in_the_api_base_out(routed_api
assert "deployment-secret" not in debug_info
+def _incomplete_openai_chat_chunk() -> ModelResponseStream:
+ return ModelResponseStream(
+ id="chatcmpl-incomplete-eof",
+ created=1,
+ model="gpt-5.6",
+ object="chat.completion.chunk",
+ choices=[
+ StreamingChoices(
+ index=0,
+ delta=Delta(content='{"findings":[{"title":"unfinished', role="assistant"),
+ finish_reason=None,
+ )
+ ],
+ )
+
+
+@pytest.mark.parametrize("custom_llm_provider", ["openai", "azure", "hosted_vllm"])
+@pytest.mark.asyncio
+async def test_async_openai_compatible_chat_clean_eof_without_finish_reason_raises_midstream_fallback(
+ logging_obj: Logging, custom_llm_provider: str
+) -> None:
+ """An OpenAI-compatible stream without a provider finish reason must fail at clean EOF."""
+
+ logging_obj.model_call_details["custom_llm_provider"] = custom_llm_provider
+ response: Final = CustomStreamWrapper(
+ completion_stream=_as_async(iter([_incomplete_openai_chat_chunk()])),
+ model="gpt-5.6",
+ logging_obj=logging_obj,
+ custom_llm_provider=custom_llm_provider,
+ stream_options={"include_usage": True},
+ )
+
+ first_chunk: Final = await response.__anext__()
+ with pytest.raises(MidStreamFallbackError) as excinfo:
+ await response.__anext__()
+
+ assert first_chunk.choices[0].delta.content == '{"findings":[{"title":"unfinished'
+ assert excinfo.value.generated_content == '{"findings":[{"title":"unfinished'
+ assert excinfo.value.is_pre_first_chunk is False
+ assert response.received_finish_reason is None
+
+
+@pytest.mark.parametrize("custom_llm_provider", ["openai", "azure", "hosted_vllm"])
+def test_sync_openai_compatible_chat_clean_eof_without_finish_reason_raises_midstream_fallback(
+ logging_obj: Logging, custom_llm_provider: str
+) -> None:
+ """The synchronous OpenAI-compatible stream must fail at clean EOF too."""
+
+ logging_obj.model_call_details["custom_llm_provider"] = custom_llm_provider
+ response: Final = CustomStreamWrapper(
+ completion_stream=iter([_incomplete_openai_chat_chunk()]),
+ model="gpt-5.6",
+ logging_obj=logging_obj,
+ custom_llm_provider=custom_llm_provider,
+ stream_options={"include_usage": True},
+ )
+
+ first_chunk: Final = next(response)
+ with pytest.raises(MidStreamFallbackError) as excinfo:
+ next(response)
+
+ assert first_chunk.choices[0].delta.content == '{"findings":[{"title":"unfinished'
+ assert excinfo.value.generated_content == '{"findings":[{"title":"unfinished'
+ assert excinfo.value.is_pre_first_chunk is False
+ 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
+) -> None:
+ """Gemini's OpenAI-compatible endpoint answers with only [DONE] when reasoning uses up max_tokens."""
+
+ wrapper: Final = CustomStreamWrapper(
+ completion_stream=iter(()) if sync_mode else _as_async(iter(())),
+ 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",
+ "message": {
+ "id": "msg_incomplete_eof",
+ "type": "message",
+ "role": "assistant",
+ "model": "claude-sonnet-4-5",
+ "content": [],
+ "stop_reason": None,
+ "usage": {"input_tokens": 9, "output_tokens": 1},
+ },
+ },
+ {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
+ {
+ "type": "content_block_delta",
+ "index": 0,
+ "delta": {"type": "text_delta", "text": '{"findings":[{"title":"unfinished'},
+ },
+)
+_ANTHROPIC_TERMINAL_EVENTS: Final[tuple[dict[str, object], ...]] = (
+ {"type": "content_block_stop", "index": 0},
+ {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 7}},
+ {"type": "message_stop"},
+)
+
+
+def _anthropic_sse_lines(events: Sequence[Mapping[str, object]]) -> list[str]:
+ return [line for event in events for line in (f"event: {event['type']}", f"data: {json.dumps(event)}", "")]
+
+
+def _anthropic_stream(
+ events: Sequence[Mapping[str, object]], sync_mode: bool, logging_obj: Logging
+) -> CustomStreamWrapper:
+ lines: Final = _anthropic_sse_lines(events)
+
+ return CustomStreamWrapper(
+ completion_stream=AnthropicModelResponseIterator(
+ streaming_response=iter(lines) if sync_mode else _as_async(iter(lines)), sync_stream=sync_mode
+ ),
+ model="claude-sonnet-4-5",
+ custom_llm_provider="anthropic",
+ logging_obj=logging_obj,
+ )
+
+
+@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(
+ sync_mode: bool, logging_obj: Logging
+) -> None:
+ wrapper: Final = _anthropic_stream(_ANTHROPIC_PARTIAL_EVENTS, sync_mode, logging_obj)
+
+ with pytest.raises(MidStreamFallbackError) as excinfo:
+ await _collect(wrapper, sync_mode)
+
+ assert excinfo.value.generated_content == '{"findings":[{"title":"unfinished'
+ assert excinfo.value.is_pre_first_chunk is False
+ assert isinstance(excinfo.value.original_exception, litellm.APIConnectionError)
+ assert "AnthropicException - Stream ended before a finish_reason was received" in str(excinfo.value)
+ assert wrapper.received_finish_reason is None
+
+
+@pytest.mark.parametrize("sync_mode", [True, False])
+@pytest.mark.asyncio
+async def test_anthropic_stream_with_message_stop_ends_with_its_finish_reason(
+ sync_mode: bool, logging_obj: Logging
+) -> None:
+ wrapper: Final = _anthropic_stream(
+ (*_ANTHROPIC_PARTIAL_EVENTS, *_ANTHROPIC_TERMINAL_EVENTS), sync_mode, logging_obj
+ )
+
+ 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"]
+ assert "".join(c.choices[0].delta.content or "" for c in received if c.choices) == (
+ '{"findings":[{"title":"unfinished'
+ )
+
+
+@pytest.mark.parametrize("sync_mode", [True, False])
+@pytest.mark.asyncio
+async def test_generic_chunk_stream_without_a_finish_reason_keeps_the_default_stop(
+ sync_mode: bool, logging_obj: Logging
+) -> None:
+ """Providers that stream GenericStreamingChunks and never mark the last one finished still end with stop,
+ but the logged response built from ``chunks`` never gains that synthetic stop."""
+ chunks: Final = [
+ {"text": text, "tool_use": None, "is_finished": False, "finish_reason": "", "usage": None, "index": 0}
+ for text in ("partial", " reply")
+ ]
+ wrapper: Final = CustomStreamWrapper(
+ completion_stream=ModelResponseListIterator(model_responses=chunks),
+ model="sagemaker/tgi-endpoint",
+ custom_llm_provider="sagemaker",
+ logging_obj=logging_obj,
+ )
+
+ 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"]
+ assert "".join(c.choices[0].delta.content or "" for c in received if c.choices) == "partial reply"
+ assert all(not (c.choices and c.choices[0].finish_reason) for c in wrapper.chunks)
+
+
+@pytest.mark.parametrize("sync_mode", [True, False])
+@pytest.mark.asyncio
+async def test_gemini_stream_with_a_usage_chunk_after_its_finish_reason_still_completes(
+ sync_mode: bool, logging_obj: Logging
+) -> None:
+ """Gemini marks every chunk unfinished, so its finishReason only shows up mid-stream before the usage chunk."""
+ lines: Final = [
+ f"data: {json.dumps(event)}"
+ for event in (
+ {"candidates": [{"content": {"parts": [{"text": "Hello"}], "role": "model"}, "index": 0}]},
+ {
+ "candidates": [
+ {"content": {"parts": [{"text": " there"}], "role": "model"}, "finishReason": "STOP", "index": 0}
+ ]
+ },
+ {"usageMetadata": {"promptTokenCount": 3, "candidatesTokenCount": 2, "totalTokenCount": 5}},
+ )
+ ]
+
+ logging_obj.update_environment_variables(
+ model="gemini-2.5-flash", user=None, optional_params={}, litellm_params={"litellm_call_id": "gemini-tail"}
+ )
+ iterator: Final = GeminiModelResponseIterator(
+ streaming_response=iter(lines) if sync_mode else _as_async(iter(lines)), sync_stream=sync_mode, logging_obj=logging_obj
+ )
+ wrapper: Final = CustomStreamWrapper(
+ completion_stream=iter(iterator) if sync_mode else iterator,
+ model="gemini-2.5-flash",
+ custom_llm_provider="gemini",
+ logging_obj=logging_obj,
+ )
+
+ received: Final = await _collect(wrapper, sync_mode)
+
+ assert "".join(c.choices[0].delta.content or "" for c in received if c.choices) == "Hello there"
+ assert [c.choices[0].finish_reason for c in received if c.choices and c.choices[0].finish_reason] == ["stop"]
+
def test_streaming_handler_with_created_time_propagation(
initialized_custom_stream_wrapper: CustomStreamWrapper, logging_obj: Logging
@@ -3638,11 +3892,8 @@ def test_usage_chunk_after_finish_reason_updates_hidden_params(logging_obj):
@pytest.mark.parametrize("sync_mode", [True, False])
-@pytest.mark.parametrize("ends_with_finish_chunk", [True, False])
@pytest.mark.asyncio
-async def test_stream_without_provider_usage_falls_back_to_the_token_estimate(
- sync_mode: bool, ends_with_finish_chunk: bool
-) -> None:
+async def test_stream_without_provider_usage_falls_back_to_the_token_estimate(sync_mode: bool) -> None:
model: Final = "gpt-4o"
messages: Final = [{"role": "user", "content": "Write a long sentence about a fox. " * 30}]
text: Final = "The quick brown fox jumps over the lazy dog. " * 20
@@ -3659,9 +3910,7 @@ async def test_stream_without_provider_usage_falls_back_to_the_token_estimate(
choices=[StreamingChoices(index=0, delta=Delta(content=""), finish_reason="stop")],
)
wrapper: Final = CustomStreamWrapper(
- completion_stream=ModelResponseListIterator(
- model_responses=[content_chunk, finish_chunk] if ends_with_finish_chunk else [content_chunk]
- ),
+ completion_stream=ModelResponseListIterator(model_responses=[content_chunk, finish_chunk]),
model=model,
custom_llm_provider="openai",
logging_obj=Logging(
@@ -3700,6 +3949,55 @@ async def test_stream_without_provider_usage_falls_back_to_the_token_estimate(
assert litellm.completion_cost(completion_response=assembled) == pytest.approx(spend)
+@pytest.mark.parametrize("sync_mode", [True, False])
+@pytest.mark.asyncio
+async def test_stream_cut_before_a_finish_chunk_bills_the_token_estimate_on_its_failure(sync_mode: bool) -> None:
+ model: Final = "gpt-4o"
+ messages: Final = [{"role": "user", "content": "Write a long sentence about a fox. " * 30}]
+ text: Final = "The quick brown fox jumps over the lazy dog. " * 20
+ logging_obj: Final = Logging(
+ model=model,
+ messages=messages,
+ stream=True,
+ call_type="completion",
+ start_time=datetime.datetime(2025, 3, 3, 21, 38, 10),
+ litellm_call_id="cut-no-usage-call",
+ function_id="cut-no-usage-fn",
+ )
+ logging_obj.update_environment_variables(
+ model=model, user=None, optional_params={}, litellm_params={"litellm_call_id": "cut-no-usage-call"}
+ )
+ wrapper: Final = CustomStreamWrapper(
+ completion_stream=ModelResponseListIterator(
+ model_responses=[
+ ModelResponseStream(
+ id="chatcmpl-cut-no-usage",
+ created=1741037890,
+ model=model,
+ choices=[StreamingChoices(index=0, delta=Delta(role="assistant", content=text))],
+ )
+ ]
+ ),
+ model=model,
+ custom_llm_provider="openai",
+ logging_obj=logging_obj,
+ stream_options=None,
+ )
+
+ with pytest.raises(MidStreamFallbackError):
+ await _collect(wrapper, sync_mode)
+
+ expected: Final = (
+ litellm.token_counter(model=model, messages=messages),
+ litellm.token_counter(model=model, text=text, count_response_tokens=True),
+ )
+ usage: Final = logging_obj.model_call_details["combined_usage_object"]
+ assert (usage.prompt_tokens, usage.completion_tokens) == expected
+ rates: Final = litellm.model_cost[model]
+ spend: Final = expected[0] * rates["input_cost_per_token"] + expected[1] * rates["output_cost_per_token"]
+ assert logging_obj.model_call_details["response_cost"] == pytest.approx(spend)
+
+
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
async def test_stream_with_usage_after_the_finish_chunk_bills_only_the_reported_usage(sync_mode: bool) -> None:
@@ -8582,10 +8880,8 @@ async def test_no_synthetic_finish_reason_logged_when_provider_sent_none(sync_mo
custom_llm_provider="bedrock",
logging_obj=logging_obj,
)
- if sync_mode:
- list(response)
- else:
- [c async for c in response]
+ with pytest.raises(MidStreamFallbackError):
+ await _collect(response, sync_mode)
assert response.received_finish_reason is None
assert all(not (c.choices and c.choices[0].finish_reason) for c in response.chunks)
diff --git a/tests/unit/llms/bedrock/chat/agentcore/test_agentcore_transformation.py b/tests/unit/llms/bedrock/chat/agentcore/test_agentcore_transformation.py
index 413061a732d..acfd954215d 100644
--- a/tests/unit/llms/bedrock/chat/agentcore/test_agentcore_transformation.py
+++ b/tests/unit/llms/bedrock/chat/agentcore/test_agentcore_transformation.py
@@ -19,8 +19,9 @@ from typing import Final
from unittest.mock import MagicMock, Mock, patch
import litellm
+from litellm.exceptions import MidStreamFallbackError
from litellm.llms.bedrock.chat.agentcore.transformation import AmazonAgentCoreConfig
-from litellm.llms.custom_httpx.http_handler import HTTPHandler
+from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
class TestAgentCoreAcceptHeader:
@@ -1271,6 +1272,87 @@ def test_agentcore_runtime_response_becomes_assistant_message(
assert response.choices[0].finish_reason == "stop"
+_AGENTCORE_RUNTIME_URL: Final = (
+ "https://bedrock-agentcore.us-west-2.amazonaws.com/runtimes/"
+ "arn%3Aaws%3Abedrock-agentcore%3Aus-west-2%3A888602223428%3Aruntime%2Fhosted_agent_13sf6-cALnp38iZD/invocations"
+)
+_AGENTCORE_DELTAS: Final = (
+ 'data: {"event":{"messageStart":{"role":"assistant"}}}\n\n'
+ 'data: {"event":{"contentBlockDelta":{"delta":{"text":"Machine learning lets "},"contentBlockIndex":0}}}\n\n'
+ 'data: {"event":{"contentBlockDelta":{"delta":{"text":"computers learn from data."},"contentBlockIndex":0}}}\n\n'
+)
+
+
+async def _stream_agentcore(body: str, sync_mode: bool) -> list[tuple[str | None, str | None]]:
+ """Stream `body` from the runtime as SSE and return each chunk's (content, finish_reason)."""
+ kwargs: Final = {
+ "model": "bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/"
+ "hosted_agent_13sf6-cALnp38iZD",
+ "messages": [{"role": "user", "content": "Explain machine learning in simple terms"}],
+ "aws_access_key_id": "AKIAAGENTCOREKEY",
+ "aws_secret_access_key": "agentcore-secret",
+ "stream": True,
+ }
+ if sync_mode:
+ return [
+ (chunk.choices[0].delta.content, chunk.choices[0].finish_reason)
+ for chunk in litellm.completion(**kwargs, client=HTTPHandler())
+ ]
+ return [
+ (chunk.choices[0].delta.content, chunk.choices[0].finish_reason)
+ async for chunk in await litellm.acompletion(
+ **kwargs, client=AsyncHTTPHandler(transport=httpx.AsyncHTTPTransport())
+ )
+ ]
+
+
+@pytest.mark.parametrize(
+ "ending",
+ [
+ pytest.param(
+ 'data: {"event":{"contentBlockStop":{"contentBlockIndex":0}}}\n\n'
+ 'data: {"event":{"messageStop":{"stopReason":"end_turn"}}}\n\n',
+ id="message_stop",
+ ),
+ pytest.param(
+ 'data: {"event":{"messageStop":{"stopReason":"end_turn"}}}\n\n'
+ 'data: {"event":{"metadata":{"usage":{"inputTokens":9,"outputTokens":7,"totalTokens":16}}}}\n\n'
+ 'data: {"message":{"role":"assistant",'
+ '"content":[{"text":"Machine learning lets computers learn from data."}]}}\n\n',
+ id="strands_full",
+ ),
+ ],
+)
+@pytest.mark.parametrize("sync_mode", [True, False])
+@pytest.mark.respx(assert_all_called=True)
+async def test_agentcore_stream_that_signals_its_end_completes(
+ ending: str, sync_mode: bool, respx_mock: respx.Router
+) -> None:
+ respx_mock.post(_AGENTCORE_RUNTIME_URL).mock(
+ return_value=httpx.Response(
+ 200, content=(_AGENTCORE_DELTAS + ending).encode(), headers={"content-type": "text/event-stream"}
+ )
+ )
+
+ chunks: Final = await _stream_agentcore(_AGENTCORE_DELTAS + ending, sync_mode)
+
+ assert "".join(content for content, _ in chunks if content) == "Machine learning lets computers learn from data."
+ assert [finish for _, finish in chunks if finish] == ["stop"]
+
+
+@pytest.mark.parametrize("sync_mode", [True, False])
+@pytest.mark.respx(assert_all_called=True)
+async def test_agentcore_stream_cut_before_it_signals_its_end_raises(sync_mode: bool, respx_mock: respx.Router) -> None:
+ respx_mock.post(_AGENTCORE_RUNTIME_URL).mock(
+ return_value=httpx.Response(
+ 200, content=_AGENTCORE_DELTAS.encode(), headers={"content-type": "text/event-stream"}
+ )
+ )
+
+ with pytest.raises(MidStreamFallbackError, match="Stream ended before a finish_reason was received"):
+ await _stream_agentcore(_AGENTCORE_DELTAS, sync_mode)
+
+
def test_get_complete_url_sends_a_non_string_runtime_endpoint_to_that_endpoint():
vpc_endpoint: Final = "https://vpce-1.bedrock-agentcore.us-west-2.vpce.amazonaws.com"
diff --git a/tests/unit/llms/langgraph/chat/test_langgraph_transformation.py b/tests/unit/llms/langgraph/chat/test_langgraph_transformation.py
index 238f7dd2f26..256876b7eea 100644
--- a/tests/unit/llms/langgraph/chat/test_langgraph_transformation.py
+++ b/tests/unit/llms/langgraph/chat/test_langgraph_transformation.py
@@ -22,12 +22,14 @@ import asyncio
import importlib
import json
import os
+from collections.abc import AsyncIterator
+from typing import Final
import httpx
import pytest
import litellm
-from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
+from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client
from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
@@ -268,3 +270,28 @@ async def test_langgraph_streaming_acompletion_without_a_client_posts_through_th
"assistant_id": "agent",
"input": {"messages": [{"role": "human", "content": "What is 25 * 4?"}]},
}
+
+
+@pytest.mark.asyncio
+async def test_langgraph_async_stream_delivers_every_step_of_the_run_before_its_stop() -> None:
+ async def events() -> AsyncIterator[bytes]:
+ question: Final = {"type": "human", "content": "What is 25 * 4?"}
+ answers: Final = ({"type": "ai", "content": "25 * 4"}, {"type": "ai", "content": " is 100"})
+ yield b'event: metadata\ndata: {"run_id": "run-1", "attempt": 1}\n\n'
+ for step in range(len(answers) + 1):
+ yield b"event: values\ndata: " + json.dumps({"messages": [question, *answers[:step]]}).encode() + b"\n\n"
+
+ client: Final = AsyncHTTPHandler(
+ transport=httpx.MockTransport(lambda request: httpx.Response(200, content=events()))
+ )
+
+ response: Final = await litellm.acompletion(
+ model="langgraph/agent",
+ messages=[{"role": "user", "content": "What is 25 * 4?"}],
+ api_base="http://langgraph.test",
+ stream=True,
+ client=client,
+ )
+ streamed: Final = [(chunk.choices[0].delta.content, chunk.choices[0].finish_reason) async for chunk in response]
+
+ assert streamed == [("25 * 4", None), (" is 100", None), (None, "stop")]
diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py
index 9bfa78ab288..93da00ff064 100644
--- a/tests/unit/test_router/test_router.py
+++ b/tests/unit/test_router/test_router.py
@@ -22656,6 +22656,9 @@ _SSE_CHUNKS: Final[tuple[bytes, ...]] = tuple(
b'data: {"id":"c","object":"chat.completion.chunk","created":1,"model":"gpt-5.6",'
b'"choices":[{"index":0,"delta":{"content":"x"},"finish_reason":null}]}\n\n'
for _ in range(5)
+) + (
+ b'data: {"id":"c","object":"chat.completion.chunk","created":1,"model":"gpt-5.6",'
+ b'"choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}\n\n',
)