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 db08191e18c..19f971a5f37 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -2741,7 +2741,7 @@ _STUCK_PROVIDER_CHUNK_CAP: Final = 10 * litellm.REPEATED_STREAMING_CHUNK_LIMIT class _StuckProvider: def __init__(self) -> None: - self.closed = threading.Event() + self.closed_early = threading.Event() self.chunks_sent = 0 def _next_chunk(self) -> ModelResponseStream: @@ -2751,15 +2751,17 @@ class _StuckProvider: def sync_stream(self) -> Iterator[ModelResponseStream]: try: yield from (self._next_chunk() for _ in range(_STUCK_PROVIDER_CHUNK_CAP)) - finally: - self.closed.set() + 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() - finally: - self.closed.set() + except GeneratorExit: + self.closed_early.set() + raise class _FailureRecorder(CustomLogger): @@ -2786,7 +2788,7 @@ def _stuck_stream_wrapper( messages=[{"role": "user", "content": "hi"}], stream=True, call_type=call_type, - start_time=datetime.datetime.now(), + 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], @@ -2815,7 +2817,8 @@ async def test_async_stream_stuck_on_one_chunk_closes_the_provider_stream_and_lo with pytest.raises(Exception, match="repeating the same chunk"): async for _ in wrapper: pass - assert provider.closed.is_set() + 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): @@ -2836,7 +2839,8 @@ def test_sync_stream_stuck_on_one_chunk_closes_the_provider_stream_and_logs_one_ with pytest.raises(Exception, match="repeating the same chunk"): for _ in wrapper: pass - assert provider.closed.is_set() + 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)