mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
fix(anthropic_messages): bound the detached pump so it cannot outlive its stream (#51268)
* fix(anthropic): time-limit /v1/messages pump waits on a stalled client The upstream pump waited on client_detached with no timeout, both when the relay queue was full and when waiting for the relay to consume a queued upstream error. client_detached is only set by the relay generator's finally block, so a client that resets its connection while stalled could leave the pump parked forever, holding its request, logging object and collected chunks. Both waits now give up after ANTHROPIC_MESSAGES_STREAM_CLIENT_STALL_TIMEOUT_SECONDS (default 120), treat the client as gone, and fall through to the existing detached billing path. * fix(anthropic): bound the detached /v1/messages drain and end a relay whose pump gave up A provider that stalled after the client left kept the detached pump, the open provider stream and every chunk it had buffered for billing alive until the provider read timed out, if it ever did. Pumps that never read another chunk after the disconnect are not counted against the detached-drain cap, so nothing limited them A detached drain still reading ANTHROPIC_MESSAGES_DETACHED_STREAM_DRAIN_TIMEOUT_SECONDS (default 600) after the client left now bills what it read and closes the provider stream. The deadline starts when the relay closes or when the client stall timeout gives up on the client, so it also covers a read that was already stalled Once the stall timeout gives up, the pump stops feeding the relay. A client that resumed reading then waited on an empty queue forever. The relay now ends with the incomplete-stream error once the pump is gone. The stall timeout is read with get_env_float so fractional values parse Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(anthropic): end a stream cleanly when the pump gives up on its last put When the upstream had already finished and only the end-of-stream sentinel was left, a client that resumed after the stall timeout got every chunk through message_stop and then an extra incomplete-stream error event. The relay now only adds that error when the client has not already received a message_stop or error event The new test pins the case where the stall timeout fires on the final sentinel put, so the pump exits with the request billed and a resumed client still sees a clean, complete stream Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(anthropic): give a stalled /v1/messages client 600 s before the pump gives up on it The 120 s default could cut off a client that was only slow. 600 s matches the usual request_timeout, so a stalled client and a stalled provider get the same budget. SSE keepalive pings do not change this, because they block on the same stalled write. A stuck pump now lives at most about 20 minutes, which keeps only a handful alive at once Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(anthropic_messages): keep queued chunks when dropping an unconsumed upstream error _exception_left_unconsumed emptied the whole relay queue to check whether the provider error was still in it, so a client that resumed after the error wait expired lost every chunk buffered ahead of the error. Re-queue everything but the error, add a regression test for the resumed client, and type the new test helpers. --------- Co-authored-by: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
nuernber
GitHub
parent
208a5863a5
commit
6116adeb74
4 changed files with 413 additions and 27 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue