diff --git a/litellm/constants.py b/litellm/constants.py index 5d0eab30b90..966a5b71b8c 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -736,6 +736,12 @@ ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE: Final = get_env_int( # upstream stream immediately, instead of continuing to drain for the real # terminal usage. ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS: Final = get_env_int("ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS", 100) +ANTHROPIC_MESSAGES_STREAM_CLIENT_STALL_TIMEOUT_SECONDS: Final = get_env_float( + "ANTHROPIC_MESSAGES_STREAM_CLIENT_STALL_TIMEOUT_SECONDS", 600.0 +) +ANTHROPIC_MESSAGES_DETACHED_STREAM_DRAIN_TIMEOUT_SECONDS: Final = get_env_float( + "ANTHROPIC_MESSAGES_DETACHED_STREAM_DRAIN_TIMEOUT_SECONDS", 600.0 +) LOGGING_WORKER_CONCURRENCY: Final = get_env_int("LOGGING_WORKER_CONCURRENCY", 100) # Must be above 0 LOGGING_WORKER_MAX_QUEUE_SIZE: Final = get_env_int("LOGGING_WORKER_MAX_QUEUE_SIZE", 50_000) LOGGING_WORKER_MAX_TIME_PER_COROUTINE: Final = get_env_float("LOGGING_WORKER_MAX_TIME_PER_COROUTINE", 20.0) diff --git a/litellm/llms/anthropic/messages/native/streaming_iterator.py b/litellm/llms/anthropic/messages/native/streaming_iterator.py index 5a7633698ec..b166b822b4f 100644 --- a/litellm/llms/anthropic/messages/native/streaming_iterator.py +++ b/litellm/llms/anthropic/messages/native/streaming_iterator.py @@ -1,17 +1,22 @@ import asyncio +import contextlib import json +import math import re from collections.abc import AsyncGenerator, AsyncIterator, Mapping, Sequence from contextvars import ContextVar from datetime import datetime from typing import TYPE_CHECKING, Any, Final, Protocol, runtime_checkable +import anyio import httpx from pydantic import TypeAdapter, ValidationError from typing_extensions import TypedDict from litellm.constants import ( + ANTHROPIC_MESSAGES_DETACHED_STREAM_DRAIN_TIMEOUT_SECONDS, ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS, + ANTHROPIC_MESSAGES_STREAM_CLIENT_STALL_TIMEOUT_SECONDS, ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE, ) from litellm.litellm_core_utils.core_helpers import process_response_headers @@ -43,6 +48,41 @@ _UPSTREAM_PUMP_TASKS: Final[set[asyncio.Task[None]]] = set() # mutable-ok: stdl _DETACHED_STREAM_DRAINS: Final[set[asyncio.Task[None]]] = set() # mutable-ok: bounded strong-ref set, detached drains +class _PumpExited: + pass + + +_PUMP_EXITED: Final = _PumpExited() + + +def _start_detached_drain_deadline(drain_deadline: anyio.CancelScope) -> None: + if drain_deadline.deadline == math.inf: + drain_deadline.deadline = ( # rebind-ok: anyio cancel scopes are armed by setting their deadline + anyio.current_time() + ANTHROPIC_MESSAGES_DETACHED_STREAM_DRAIN_TIMEOUT_SECONDS + ) + + +async def _next_relay_item( + queue: "asyncio.Queue[bytes | None | BaseException]", pump_task: "asyncio.Task[None]" +) -> bytes | None | BaseException | _PumpExited: + """The next queued item, or ``_PUMP_EXITED`` once the pump is gone and will never send another.""" + with contextlib.suppress(asyncio.QueueEmpty): + return queue.get_nowait() + if pump_task.done(): + return _PUMP_EXITED + getter: Final = asyncio.ensure_future(queue.get()) + try: + await asyncio.wait(frozenset((getter, pump_task)), return_when=asyncio.FIRST_COMPLETED) + finally: + if not getter.done(): + getter.cancel() + if getter.done() and not getter.cancelled(): + return getter.result() + with contextlib.suppress(asyncio.QueueEmpty): + return queue.get_nowait() + return _PUMP_EXITED + + def is_message_stop_chunk(chunk: object) -> bool: if isinstance(chunk, dict): return chunk.get("type") == "message_stop" @@ -204,7 +244,11 @@ def _try_claim_detached_drain_slot() -> bool: def _exception_left_unconsumed(queue: "asyncio.Queue[bytes | None | BaseException]", exc: BaseException) -> bool: + """Drop ``exc`` from ``queue`` if the client never consumed it, keeping the chunks queued ahead of it.""" remaining: Final = tuple(queue.get_nowait() for _ in range(queue.qsize())) + for item in remaining: + if item is not exc: + queue.put_nowait(item) return any(item is exc for item in remaining) @@ -606,8 +650,12 @@ class BaseAnthropicMessagesStreamingIterator: keeps a single ``collected_chunks`` copy for billing, and the number of such post-disconnect drains running at once is capped so client behavior can't create unbounded worker state; over the cap the pump bills what it - has rather than draining further. Detached-drain lifetime is otherwise - bounded by the upstream stream/read timeout. + has rather than draining further. A drain still reading + ``ANTHROPIC_MESSAGES_DETACHED_STREAM_DRAIN_TIMEOUT_SECONDS`` after the + client left, a stalled provider for one, bills what it has and closes the + provider stream too, so it cannot hold its buffered response forever. A + relay whose pump exited without an end-of-stream sentinel ends the stream + with an incomplete-stream error instead of waiting on an empty queue. An upstream failure (Bedrock read / decode / chunk-conversion error) that happens while the client is still connected is forwarded through @@ -622,23 +670,33 @@ class BaseAnthropicMessagesStreamingIterator: maxsize=ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE ) client_detached: Final = asyncio.Event() + drain_deadline: Final = anyio.CancelScope() - pump_task: Final = asyncio.create_task(self._pump_upstream_to_queue(completion_stream, queue, client_detached)) + pump_task: Final = asyncio.create_task( + self._pump_upstream_to_queue(completion_stream, queue, client_detached, drain_deadline) + ) _UPSTREAM_PUMP_TASKS.add(pump_task) pump_task.add_done_callback(_UPSTREAM_PUMP_TASKS.discard) reached_end = False # rebind-ok: flipped once the relay consumes the end-of-stream sentinel + relayed_terminal_event = False # rebind-ok: flipped once a message_stop or error event reaches the client try: while True: - item = await queue.get() + item = await _next_relay_item(queue, pump_task) + if isinstance(item, _PumpExited): + if not relayed_terminal_event: + yield _incomplete_stream_error_sse_event() + break if item is None: reached_end = True break if isinstance(item, BaseException): raise item + relayed_terminal_event = relayed_terminal_event or _is_terminal_stream_chunk(item) yield item finally: client_detached.set() + _start_detached_drain_deadline(drain_deadline) if not reached_end: self._dispatch_pending_deferred_logging() @@ -720,13 +778,18 @@ class BaseAnthropicMessagesStreamingIterator: put_task: Final = asyncio.ensure_future(queue.put(item)) detached_task: Final = asyncio.ensure_future(client_detached.wait()) try: - await asyncio.wait(frozenset((put_task, detached_task)), return_when=asyncio.FIRST_COMPLETED) + await asyncio.wait( + frozenset((put_task, detached_task)), + timeout=ANTHROPIC_MESSAGES_STREAM_CLIENT_STALL_TIMEOUT_SECONDS, + return_when=asyncio.FIRST_COMPLETED, + ) finally: if not detached_task.done(): detached_task.cancel() if put_task.done() and not put_task.cancelled(): return True put_task.cancel() + client_detached.set() return False async def _pump_upstream_to_queue( @@ -734,6 +797,7 @@ class BaseAnthropicMessagesStreamingIterator: completion_stream: AsyncIterator[bytes | GenericStreamingChunk | ModelResponseStream | Mapping[str, object]], queue: "asyncio.Queue[bytes | None | BaseException]", client_detached: "asyncio.Event", + drain_deadline: anyio.CancelScope, ) -> None: """Drain the whole upstream into ``queue`` (backpressured) and bill once. @@ -750,32 +814,41 @@ class BaseAnthropicMessagesStreamingIterator: collected_chunks: Final[list[bytes]] = [] # mutable-ok: SSE billing buffer appended to across the drain saw_terminal_event = False # rebind-ok: accumulates across the upstream loop draining_detached = False # rebind-ok: set once this pump claims a detached-drain slot + over_drain_cap = False # rebind-ok: set when no detached-drain slot is free try: - async for chunk in completion_stream: - if self.completion_start_time is None: - self.completion_start_time = datetime.now() - saw_terminal_event = saw_terminal_event or _is_terminal_stream_chunk(chunk) - encoded_chunk = self._convert_chunk_to_sse_format(chunk) - collected_chunks.append(encoded_chunk) - if not client_detached.is_set(): - await self._enqueue_for_client(queue, client_detached, encoded_chunk) - continue - if not draining_detached: - if not _try_claim_detached_drain_slot(): - verbose_proxy_logger.warning( - "async_sse_wrapper: detached-drain cap (%d) reached; billing %d partial " - "chunks and aborting the upstream stream to stop provider billing", - ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS, - len(collected_chunks), - ) - await self._bill_collected_chunks(collected_chunks, stream_teardown=True) - await self._abort_upstream(completion_stream) - return - draining_detached = True + with drain_deadline: + async for chunk in completion_stream: + if self.completion_start_time is None: + self.completion_start_time = datetime.now() + saw_terminal_event = saw_terminal_event or _is_terminal_stream_chunk(chunk) + encoded_chunk = self._convert_chunk_to_sse_format(chunk) + collected_chunks.append(encoded_chunk) + if not client_detached.is_set(): + if not await self._enqueue_for_client(queue, client_detached, encoded_chunk): + _start_detached_drain_deadline(drain_deadline) + continue + if not draining_detached: + if not _try_claim_detached_drain_slot(): + over_drain_cap = True + break + draining_detached = True except Exception as exc: # noqa: BLE001 # upstream errors are handled/forwarded by _handle_pump_upstream_error await self._handle_pump_upstream_error(queue, client_detached, collected_chunks, exc) return + if over_drain_cap or drain_deadline.cancelled_caught: + verbose_proxy_logger.warning( + "async_sse_wrapper: %s; billing %d partial chunks and aborting the upstream stream to stop " + "provider billing", + f"detached-drain cap ({ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS}) reached" + if over_drain_cap + else f"detached drain still running {ANTHROPIC_MESSAGES_DETACHED_STREAM_DRAIN_TIMEOUT_SECONDS}s " + "after the client left", + len(collected_chunks), + ) + await self._bill_collected_chunks(collected_chunks, stream_teardown=True) + await self._abort_upstream(completion_stream) + return if client_detached.is_set(): await self._bill_collected_chunks(collected_chunks, stream_teardown=True) return @@ -814,7 +887,10 @@ class BaseAnthropicMessagesStreamingIterator: exception=described, ) if not client_detached.is_set() and await self._enqueue_for_client(queue, client_detached, described): - await client_detached.wait() + with contextlib.suppress(asyncio.TimeoutError): + await asyncio.wait_for( + client_detached.wait(), timeout=ANTHROPIC_MESSAGES_STREAM_CLIENT_STALL_TIMEOUT_SECONDS + ) if not _exception_left_unconsumed(queue, described): return await self._fire_detached_failure_hook(described) diff --git a/tests/documentation_tests/test_env_keys.py b/tests/documentation_tests/test_env_keys.py index 2a7da58f39c..72386877704 100644 --- a/tests/documentation_tests/test_env_keys.py +++ b/tests/documentation_tests/test_env_keys.py @@ -38,6 +38,8 @@ EXCLUDED_ROLLOUT_FLAGS = { # These are advanced settings with sensible defaults that most users should not modify EXCLUDED_INTERNAL_TUNING_VARS = { "ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS", + "ANTHROPIC_MESSAGES_DETACHED_STREAM_DRAIN_TIMEOUT_SECONDS", + "ANTHROPIC_MESSAGES_STREAM_CLIENT_STALL_TIMEOUT_SECONDS", "ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", "DEFAULT_S3_MAX_ADAPTIVE_CONCURRENCY", } diff --git a/tests/unit/llms/anthropic/messages/native/test_streaming_iterator.py b/tests/unit/llms/anthropic/messages/native/test_streaming_iterator.py index 39f7a4a05cb..ded599512d0 100644 --- a/tests/unit/llms/anthropic/messages/native/test_streaming_iterator.py +++ b/tests/unit/llms/anthropic/messages/native/test_streaming_iterator.py @@ -1,5 +1,7 @@ import asyncio +import gc import json +import tracemalloc from collections.abc import AsyncIterator, Mapping, Sequence from datetime import datetime from typing import Final, Literal @@ -1921,3 +1923,303 @@ async def test_safeguarded_sse_stream_reports_upstream_buffering_and_closes_the_ assert SafeguardedAnthropicSSEStream( _chunked(b"", 1), _safeguards_evaluator(_OneVerdictClassifier("{}")) ).has_buffered_provider_output is False + + +def _finite_stream(chunks: int) -> AsyncIterator[dict[str, object]]: + async def _stream() -> AsyncIterator[dict[str, object]]: + yield {"type": "message_start", "message": {"id": "m", "usage": {"input_tokens": 5, "output_tokens": 1}}} + for i in range(chunks): + yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": f"t{i}"}} + yield {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 42}} + yield {"type": "message_stop"} + + return _stream() + + +async def _wait_until_no_live_pumps() -> None: + for _ in range(300): + if not any(not task.done() for task in streaming_iterator_module._UPSTREAM_PUMP_TASKS): + return + await asyncio.sleep(0.01) + raise AssertionError("pump task is still alive") + + +@pytest.mark.asyncio +async def test_pump_exits_when_relay_generator_is_never_closed_and_queue_is_full(monkeypatch: pytest.MonkeyPatch): + """ + Regression: a client that resets its connection while the relay queue is full + can leave the relay generator unfinalized, so client_detached is never set. + The pump must give up on the stalled client after the stall timeout instead of + waiting forever, then bill what it drained. + """ + monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", 2) + monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_STREAM_CLIENT_STALL_TIMEOUT_SECONDS", 0.05) + + iterator = _RecordingLoggingIterator(litellm_logging_obj=_make_logging_obj("pump_stall_full_queue"), request_body={}) + gen = iterator.async_sse_wrapper(_finite_stream(50)) + await gen.__anext__() + + await _wait_until_no_live_pumps() + + assert iterator.logging_call_count == 1 + assert any(c.startswith(b"event: message_stop\n") for c in iterator.logged_chunks) + + +@pytest.mark.asyncio +async def test_pump_keeps_streaming_to_slow_client_that_stays_within_stall_timeout(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", 2) + monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_STREAM_CLIENT_STALL_TIMEOUT_SECONDS", 5) + + iterator = _RecordingLoggingIterator(litellm_logging_obj=_make_logging_obj("pump_slow_client"), request_body={}) + received = [] + async for chunk in iterator.async_sse_wrapper(_finite_stream(20)): + received.append(chunk) + await asyncio.sleep(0.01) + + assert any(c.startswith(b"event: message_stop\n") for c in received) + assert len(received) == 23 + + +@pytest.mark.asyncio +async def test_pump_exits_when_queued_upstream_error_is_never_consumed_and_relay_is_never_closed( + monkeypatch: pytest.MonkeyPatch, +): + """ + Regression: the pump waits for the relay to consume a queued upstream error. + If the relay generator is never finalized, that wait must time out and fire + the detached failure hook instead of hanging forever. + """ + monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_STREAM_CLIENT_STALL_TIMEOUT_SECONDS", 0.05) + upstream_errored = asyncio.Event() + + async def _failing_stream() -> AsyncIterator[dict[str, object]]: + yield {"type": "message_start", "message": {"id": "msg_1", "usage": {"input_tokens": 52, "output_tokens": 1}}} + upstream_errored.set() + raise _ProviderStreamError("mid-stream failure", status_code=500) + + iterator = _RecordingLoggingIterator( + litellm_logging_obj=_make_logging_obj("pump_stall_queued_error", _FailureRecorder()), request_body={} + ) + detached_hook = _DetachedFailureRecorder() + iterator.litellm_logging_obj.on_detached_stream_failure = detached_hook + + gen = iterator.async_sse_wrapper(_failing_stream()) + await gen.__anext__() + await upstream_errored.wait() + + await _wait_until_no_live_pumps() + + assert len(detached_hook.exceptions) == 1 + + +@pytest.mark.asyncio +async def test_slow_client_that_resumes_after_an_unconsumed_upstream_error_still_gets_its_buffered_chunks( + monkeypatch: pytest.MonkeyPatch, +): + """ + Regression: when the pump gave up waiting for a stalled client to consume a queued provider + error, it emptied the whole relay queue to check whether the error was still in it, so a client + that resumed lost every chunk buffered ahead of the error. Only the error may be dropped. + """ + monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_STREAM_CLIENT_STALL_TIMEOUT_SECONDS", 0) + deltas: Final = tuple( + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": f"t{i}"}} for i in range(3) + ) + + async def _failing_stream() -> AsyncIterator[dict[str, object]]: + yield {"type": "message_start", "message": {"id": "msg_1", "usage": {"input_tokens": 52, "output_tokens": 1}}} + for delta in deltas: + yield delta + raise _ProviderStreamError("mid-stream failure", status_code=500) + + iterator = _RecordingLoggingIterator( + litellm_logging_obj=_make_logging_obj("resumed_client_queued_error", _FailureRecorder()), request_body={} + ) + detached_hook = _DetachedFailureRecorder() + iterator.litellm_logging_obj.on_detached_stream_failure = detached_hook + + gen = iterator.async_sse_wrapper(_failing_stream()) + first_frame: Final = await gen.__anext__() + await _wait_until_no_live_pumps() + + async def _read_rest() -> list[bytes]: + return [frame async for frame in gen] + + rest: Final = await asyncio.wait_for(_read_rest(), timeout=5) + + assert first_frame.startswith(b"event: message_start\n") + assert rest == [ + *(iterator._convert_chunk_to_sse_format(delta) for delta in deltas), + _incomplete_stream_error_sse_event(), + ] + assert len(detached_hook.exceptions) == 1 + + +class _BilledChunkCountIterator(BaseAnthropicMessagesStreamingIterator): + def __init__(self, litellm_logging_obj: LiteLLMLoggingObj, billed_chunk_counts: list[int]): + super().__init__(litellm_logging_obj=litellm_logging_obj, request_body={}) + self._billed_chunk_counts: Final = billed_chunk_counts + + async def _handle_streaming_logging(self, collected_chunks: list[bytes], *, stream_teardown: bool = False) -> None: + self._billed_chunk_counts.append(len(collected_chunks)) + + +@pytest.mark.asyncio +async def test_async_sse_wrapper_releases_stalled_upstream_after_client_disconnect(monkeypatch: pytest.MonkeyPatch): + """ + Regression: when the provider stalls after the client disconnects, the detached pump kept the + provider stream open and every chunk it had buffered for billing alive until the provider read + timed out, so each such request leaked its whole response. The drain must stop at its time + limit, bill what it read, close the provider stream and free the buffer. + """ + monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_DETACHED_STREAM_DRAIN_TIMEOUT_SECONDS", 0) + requests: Final = 5 + chunks_before_stall: Final = 50 + payload: Final = "x" * 65536 + response_bytes: Final = chunks_before_stall * len(payload) + provider_stalled: Final = asyncio.Event() + upstream_closed: Final = tuple(asyncio.Event() for _ in range(requests)) + billed_chunk_counts: Final[list[int]] = [] + + async def _stalling_stream(closed: asyncio.Event) -> AsyncIterator[dict[str, object]]: + try: + yield {"type": "message_start", "message": {"id": "m", "usage": {"input_tokens": 5, "output_tokens": 1}}} + for _ in range(chunks_before_stall): + yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": payload}} + await provider_stalled.wait() + finally: + closed.set() + + gc.collect() + tracemalloc.start() + try: + baseline: Final = tracemalloc.get_traced_memory()[0] + for index, closed in enumerate(upstream_closed): + iterator = _BilledChunkCountIterator(_make_logging_obj(f"stalled_drain_{index}"), billed_chunk_counts) + gen = iterator.async_sse_wrapper(_stalling_stream(closed)) + await gen.__anext__() + await gen.aclose() + try: + await asyncio.wait_for(asyncio.gather(*(closed.wait() for closed in upstream_closed)), timeout=5) + except asyncio.TimeoutError: + pass + gc.collect() + retained: Final = tracemalloc.get_traced_memory()[0] - baseline + closed_while_stalled: Final = sum(closed.is_set() for closed in upstream_closed) + billed_while_stalled: Final = tuple(billed_chunk_counts) + finally: + tracemalloc.stop() + provider_stalled.set() + await asyncio.gather(*tuple(streaming_iterator_module._UPSTREAM_PUMP_TASKS), return_exceptions=True) + + assert closed_while_stalled == requests, f"only {closed_while_stalled}/{requests} stalled provider streams closed" + assert retained < response_bytes, f"{retained} bytes still held after {requests} disconnected stalled streams" + assert billed_while_stalled == (chunks_before_stall + 1,) * requests + + +@pytest.mark.asyncio +async def test_slow_client_that_resumes_after_the_stall_timeout_gets_an_error_instead_of_hanging( + monkeypatch: pytest.MonkeyPatch, +): + """ + Regression: once the pump gave up on a stalled client it stopped feeding the relay without an + end-of-stream sentinel, so a client that started reading again drained the queue and then + waited forever. It must get the incomplete-stream error and an ended stream. + """ + monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", 1) + monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_STREAM_CLIENT_STALL_TIMEOUT_SECONDS", 0) + iterator = _RecordingLoggingIterator(litellm_logging_obj=_make_logging_obj("resumed_slow_client"), request_body={}) + gen = iterator.async_sse_wrapper(_finite_stream(20)) + first_frame: Final = await gen.__anext__() + await _wait_until_no_live_pumps() + + async def _read_rest() -> list[bytes]: + return [frame async for frame in gen] + + rest: Final = await asyncio.wait_for(_read_rest(), timeout=5) + + assert first_frame.startswith(b"event: message_start\n") + assert rest[-1] == _incomplete_stream_error_sse_event() + assert any(c.startswith(b"event: message_stop\n") for c in iterator.logged_chunks) + + + +@pytest.mark.asyncio +async def test_pump_closes_a_stalled_upstream_after_giving_up_on_a_client_that_never_closes( + monkeypatch: pytest.MonkeyPatch, +): + """ + Regression: a client that stops reading without closing its relay is detached by the stall + timeout, and the pump then drains on its own. If the provider stalls too, nothing else ever + ends that drain, so the drain time limit has to start when the stall timeout gives up. + """ + monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", 1) + monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_STREAM_CLIENT_STALL_TIMEOUT_SECONDS", 0) + monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_DETACHED_STREAM_DRAIN_TIMEOUT_SECONDS", 0) + provider_stalled: Final = asyncio.Event() + upstream_closed: Final = asyncio.Event() + + async def _stalling_stream() -> AsyncIterator[dict[str, object]]: + try: + yield {"type": "message_start", "message": {"id": "m", "usage": {"input_tokens": 5, "output_tokens": 1}}} + for i in range(5): + yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": f"t{i}"}} + await provider_stalled.wait() + finally: + upstream_closed.set() + + iterator = _RecordingLoggingIterator(litellm_logging_obj=_make_logging_obj("client_never_closes"), request_body={}) + gen = iterator.async_sse_wrapper(_stalling_stream()) + await gen.__anext__() + try: + await asyncio.wait_for(upstream_closed.wait(), timeout=5) + except asyncio.TimeoutError: + pass + closed_while_stalled: Final = upstream_closed.is_set() + provider_stalled.set() + await _wait_until_no_live_pumps() + + assert closed_while_stalled, "the stalled provider stream stayed open after the client was given up on" + assert iterator.logging_call_count == 1 + + +@pytest.mark.asyncio +async def test_pump_gives_up_on_the_end_sentinel_and_a_resumed_client_still_gets_a_clean_complete_stream( + monkeypatch: pytest.MonkeyPatch, +): + """ + Regression: when the upstream had already finished and only the end-of-stream sentinel was + left, a stalled client with a full relay queue kept the pump waiting on that last put. The + pump must give up at the stall timeout with the request already billed, and a client that + resumes must still get every chunk, ending at message_stop with no error event after it. + """ + content_chunks: Final = 3 + remaining_events: Final = content_chunks + 2 + monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", remaining_events) + monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_STREAM_CLIENT_STALL_TIMEOUT_SECONDS", 0) + client_stopped_reading: Final = asyncio.Event() + + async def _stream_that_finishes_after_the_client_stalls() -> AsyncIterator[dict[str, object]]: + yield {"type": "message_start", "message": {"id": "m", "usage": {"input_tokens": 5, "output_tokens": 1}}} + await client_stopped_reading.wait() + for i in range(content_chunks): + yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": f"t{i}"}} + yield {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 42}} + yield {"type": "message_stop"} + + iterator = _RecordingLoggingIterator(litellm_logging_obj=_make_logging_obj("blocked_end_sentinel"), request_body={}) + gen = iterator.async_sse_wrapper(_stream_that_finishes_after_the_client_stalls()) + first_frame: Final = await gen.__anext__() + client_stopped_reading.set() + await _wait_until_no_live_pumps() + billed_before_resume: Final = iterator.logging_call_count + + async def _read_rest() -> list[bytes]: + return [frame async for frame in gen] + + frames: Final = [first_frame, *await asyncio.wait_for(_read_rest(), timeout=5)] + + assert billed_before_resume == 1 + assert len(frames) == remaining_events + 1 + assert frames[-1].startswith(b"event: message_stop\n") + assert not any(frame.startswith(b"event: error\n") for frame in frames)