From 96792d09138deecd5bc9eb5a746e040ea6f49beb Mon Sep 17 00:00:00 2001 From: Arsh Goyal Date: Mon, 26 Oct 2026 23:33:37 +0530 Subject: [PATCH] test(streaming): pin that a stream stuck on one chunk closes its provider stream and logs one failure (#38753) * fix(streaming): stop re-raising and re-logging a failed stream on every next() A stream that raised keeps its provider iterator alive, so a consumer that iterates past the error re-detected, re-raised and re-logged the same failure on every call, flooding failure callbacks for one request. Latch the failure and re-raise it without logging again. * fix(streaming): drop the retained traceback before re-raising a failed stream * test(streaming): count failures through a real logging callback and require a clean stop on later reads * test(streaming): require the stuck stream to close early and release its socket The provider stream now records a GeneratorExit and caps its chunks, so draining it or detecting the loop late fails. Two wire cases check a sync and an async stuck stream release their TCP session. * test(streaming): avoid failure recorder name collision --------- Co-authored-by: Devin AI Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Co-authored-by: moe-berri --- ...est_sync_stream_connection_release_wire.py | 33 +++++ .../test_streaming_handler.py | 118 ++++++++++++++++++ 2 files changed, 151 insertions(+) diff --git a/tests/integration/sdk/test_sync_stream_connection_release_wire.py b/tests/integration/sdk/test_sync_stream_connection_release_wire.py index a8b49d285c3..21774df66d0 100644 --- a/tests/integration/sdk/test_sync_stream_connection_release_wire.py +++ b/tests/integration/sdk/test_sync_stream_connection_release_wire.py @@ -212,3 +212,36 @@ def test_sync_stream_that_raises_mid_stream_releases_the_connection() -> None: assert _open_sessions_settling_to(wire, at_most=0) == 0, ( "a stream that raised mid-way left its TCP session open" ) + + +_STUCK: Final = _Shape( + "hosted_vllm", + "hosted_vllm/stuck-model", + _OPENAI_OPENING, + (_openai_delta("the model is stuck on this", None),) * (10 * litellm.REPEATED_STREAMING_CHUNK_LIMIT), +) + + +def _still_generating_peer(request: Request) -> Reply: + return Reply(content_type="text/event-stream", chunks=(_STUCK.opening, *_STUCK.closing), pause_between_chunks=0.01) + + +def test_sync_stream_stuck_repeating_one_chunk_releases_the_connection() -> None: + with wire_server(_still_generating_peer, keep_alive=True) as wire: + stuck: Final = litellm.completion(**_request_kwargs(_STUCK, wire)) + with pytest.raises(Exception, match="repeating the same chunk"): + for _ in stuck: + pass + assert _open_sessions_settling_to(wire, at_most=0) == 0, "a stream stuck on one chunk left its TCP session open" + + +@pytest.mark.asyncio +async def test_async_stream_stuck_repeating_one_chunk_releases_the_connection() -> None: + with wire_server(_still_generating_peer, keep_alive=True) as wire: + stuck: Final = await litellm.acompletion(**_request_kwargs(_STUCK, wire)) + with pytest.raises(Exception, match="repeating the same chunk"): + async for _ in stuck: + pass + assert await _open_sessions_settling_to_while_the_loop_runs(wire, at_most=0) == 0, ( + "a stream stuck on one chunk left its TCP session open" + ) diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index c13cea51183..cc1fa1b8803 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -1,6 +1,7 @@ import asyncio import datetime import json +import queue import sys import threading import time @@ -2819,6 +2820,123 @@ def test_raise_on_model_repetition_tolerates_empty_choices( wrapper.raise_on_model_repetition() +_STUCK_CHUNK: Final = "the model is stuck on this" +_STUCK_PROVIDER_CHUNK_CAP: Final = 10 * litellm.REPEATED_STREAMING_CHUNK_LIMIT + + +class _StuckProvider: + def __init__(self) -> None: + self.closed_early = threading.Event() + self.chunks_sent = 0 + + def _next_chunk(self) -> ModelResponseStream: + self.chunks_sent += 1 + return _make_chunk(_STUCK_CHUNK) + + def sync_stream(self) -> Iterator[ModelResponseStream]: + try: + yield from (self._next_chunk() for _ in range(_STUCK_PROVIDER_CHUNK_CAP)) + except GeneratorExit: + self.closed_early.set() + raise + + async def async_stream(self) -> AsyncIterator[ModelResponseStream]: + try: + for _ in range(_STUCK_PROVIDER_CHUNK_CAP): + yield self._next_chunk() + except GeneratorExit: + self.closed_early.set() + raise + + +class _StuckStreamFailureRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.failures: Final = queue.SimpleQueue[str]() + + def log_failure_event(self, kwargs: object, response_obj: object, start_time: object, end_time: object) -> None: + self.failures.put(str(cast(Mapping[str, object], kwargs)["exception"])) + + async def async_log_failure_event( + self, kwargs: object, response_obj: object, start_time: object, end_time: object + ) -> None: + self.failures.put(str(cast(Mapping[str, object], kwargs)["exception"])) + + +def _stuck_stream_wrapper( + provider_stream: Iterator[ModelResponseStream] | AsyncIterator[ModelResponseStream], + call_type: Literal["completion", "acompletion"], + recorder: _StuckStreamFailureRecorder, +) -> CustomStreamWrapper: + logging_obj: Final = Logging( + model="test-model", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type=call_type, + start_time=datetime.datetime(2026, 1, 1), + litellm_call_id=f"stuck-stream-{call_type}", + function_id=f"stuck-stream-{call_type}", + dynamic_failure_callbacks=[recorder], + dynamic_async_failure_callbacks=[recorder], + ) + logging_obj.update_environment_variables( + model="test-model", + optional_params={}, + litellm_params=get_litellm_params(custom_llm_provider="openai", acompletion=call_type == "acompletion"), + ) + return CustomStreamWrapper( + completion_stream=provider_stream, + model="test-model", + logging_obj=logging_obj, + custom_llm_provider="openai", + ) + + +@pytest.mark.asyncio +async def test_async_stream_stuck_on_one_chunk_closes_the_provider_stream_and_logs_one_failure(): + """https://github.com/BerriAI/litellm/issues/19428 and https://github.com/BerriAI/litellm/issues/13786""" + provider: Final = _StuckProvider() + recorder: Final = _StuckStreamFailureRecorder() + wrapper: Final = _stuck_stream_wrapper(provider.async_stream(), "acompletion", recorder) + + with pytest.raises(Exception, match="repeating the same chunk"): + async for _ in wrapper: + pass + assert provider.closed_early.is_set() + assert provider.chunks_sent <= litellm.REPEATED_STREAMING_CHUNK_LIMIT + 1 + chunks_sent_at_failure: Final = provider.chunks_sent + + with pytest.raises(StopAsyncIteration): + await wrapper.__anext__() + await wrapper.wait_for_failure_logging(timeout=5) + + assert provider.chunks_sent == chunks_sent_at_failure + assert "repeating the same chunk" in recorder.failures.get_nowait() + assert recorder.failures.empty() + + +def test_sync_stream_stuck_on_one_chunk_closes_the_provider_stream_and_logs_one_failure(): + """Sync counterpart of the async test above""" + provider: Final = _StuckProvider() + recorder: Final = _StuckStreamFailureRecorder() + wrapper: Final = _stuck_stream_wrapper(provider.sync_stream(), "completion", recorder) + + with pytest.raises(Exception, match="repeating the same chunk"): + for _ in wrapper: + pass + assert provider.closed_early.is_set() + assert provider.chunks_sent <= litellm.REPEATED_STREAMING_CHUNK_LIMIT + 1 + chunks_sent_at_failure: Final = provider.chunks_sent + assert "repeating the same chunk" in recorder.failures.get(timeout=5) + + with pytest.raises(StopIteration): + next(wrapper) + + assert provider.chunks_sent == chunks_sent_at_failure + with pytest.raises(queue.Empty): + recorder.failures.get(timeout=0.5) + + def test_usage_chunk_after_finish_reason_updates_hidden_params(logging_obj): """ Test that provider-reported usage from a post-finish_reason chunk