From 1b401af716ee8efdd1f5ca0b3a7437dbc59cc85f Mon Sep 17 00:00:00 2001 From: nuernber Date: Thu, 6 Aug 2026 08:20:30 -0700 Subject: [PATCH 01/15] fix(anthropic_messages): drain upstream in a detached pump so client disconnect doesn't undercount Bedrock spend On the /v1/messages -> bedrock/ invoke streaming path a client disconnect raises CancelledError inside the httpx socket read, which unwinds the whole upstream generator chain before any finally can drain it. Bedrock keeps generating and billing the full response, so spend tracking logged only the truncated partial the client drained (output tokens ~1-15 vs the real count) and undercounted against AWS invocation logs. Move the upstream read into a detached background task that fully drains the provider stream to its terminal message_delta/message_stop and bills there. The client-facing generator only relays chunks off a queue, so a disconnect tears down the relay but not the pump. A client_detached event stops enqueueing after disconnect so the queue can't grow unbounded. --- .../messages/streaming_iterator.py | 82 +++++++-- ..._v1_messages_streaming_disconnect_spend.py | 167 ++++++++++++++++++ .../messages/test_streaming_iterator.py | 121 +++++++++++++ type-discipline-budget.json | 2 +- 4 files changed, 358 insertions(+), 14 deletions(-) create mode 100644 tests/pass_through_tests/test_v1_messages_streaming_disconnect_spend.py diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py index f999eae1be6..abcd817ea18 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py @@ -18,6 +18,13 @@ from litellm.types.utils import GenericStreamingChunk, ModelResponseStream GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ: Final = PassThroughEndpointLogging() +# asyncio holds only a weak reference to a bare create_task() result, so a +# fire-and-forget task can be garbage-collected mid-run. The upstream pump +# below must outlive the client-facing generator (which is closed on client +# disconnect), so root every pump task in a module-level set per the stdlib +# guidance and drop it again from the done callback. +_UPSTREAM_PUMP_TASKS: Final[set[asyncio.Task[None]]] = set() # mutable-ok: stdlib strong-ref set for pump tasks + INCOMPLETE_STREAM_ERROR_MESSAGE: Final = ( "Provider stream ended before emitting a message_stop event; " "the response is incomplete and any partial content (e.g. tool_use input JSON) may be truncated." @@ -192,21 +199,70 @@ class BaseAnthropicMessagesStreamingIterator: Generic async SSE wrapper that converts streaming chunks to SSE format and handles logging. + The upstream read runs in a detached background task (``_pump_upstream``) + so that a client disconnect tears down only this client-facing generator, + never the upstream drain + billing. The provider (e.g. Bedrock) keeps + generating and billing the full response regardless of the client, so + draining it to completion is what lets spend tracking see the real + terminal ``message_delta`` / ``message_stop`` usage instead of a + truncated placeholder count. Chunks reach the client through a queue; + once the client goes away the pump only buffers for billing so the queue + can't grow unbounded. + This method provides the common logic for both Anthropic and Bedrock implementations. """ - collected_chunks: Final = [] - saw_terminal_event = False + from litellm._logging import verbose_proxy_logger - 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) - yield encoded_chunk + queue: Final[asyncio.Queue[bytes | None]] = asyncio.Queue() + client_detached: Final = asyncio.Event() - if not saw_terminal_event: - yield _incomplete_stream_error_sse_event() + async def _pump_upstream() -> None: + collected_chunks: Final[list[bytes]] = [] + saw_terminal_event = False # rebind-ok: accumulates across the upstream loop + 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(): + queue.put_nowait(encoded_chunk) + except Exception as exc: # noqa: BLE001 # must still flush partial usage in finally, not crash the pump + verbose_proxy_logger.warning( + "async_sse_wrapper upstream pump stopped after %d chunks: %s(%s)", + len(collected_chunks), + type(exc).__name__, + exc, + ) + finally: + if not client_detached.is_set(): + if not saw_terminal_event: + queue.put_nowait(_incomplete_stream_error_sse_event()) + queue.put_nowait(None) + try: + await self._handle_streaming_logging(collected_chunks) + except Exception as exc: # noqa: BLE001 # billing is best-effort; never crash the pump + verbose_proxy_logger.warning( + "async_sse_wrapper billing failed after %d chunks: %s(%s)", + len(collected_chunks), + type(exc).__name__, + exc, + ) - # Handle logging after all chunks are processed - await self._handle_streaming_logging(collected_chunks) + pump_task: Final = asyncio.create_task(_pump_upstream()) + _UPSTREAM_PUMP_TASKS.add(pump_task) + pump_task.add_done_callback(_UPSTREAM_PUMP_TASKS.discard) + + try: + while True: + item = await queue.get() + if item is None: + break + yield item + finally: + # Client-facing generator is being torn down (normal end or a + # disconnect GeneratorExit). Signal the pump to stop enqueueing so + # the queue can't grow unbounded, but let it keep draining upstream + # to its terminal usage event for accurate billing. + client_detached.set() diff --git a/tests/pass_through_tests/test_v1_messages_streaming_disconnect_spend.py b/tests/pass_through_tests/test_v1_messages_streaming_disconnect_spend.py new file mode 100644 index 00000000000..5cd80b5ec6f --- /dev/null +++ b/tests/pass_through_tests/test_v1_messages_streaming_disconnect_spend.py @@ -0,0 +1,167 @@ +""" +Regression test: /v1/messages streaming interrupted mid-stream must still +produce a spend-log entry. + +On v1.79.1 the proxy records spend for the partially-streamed request. +A refactor on `main` broke that path, so the same scenario now produces +zero spend-log rows. + +Run against a live proxy (e.g. ``litellm --config proxy_server_config.yaml``): + + pytest tests/pass_through_tests/test_v1_messages_streaming_disconnect_spend.py -s +""" + +import asyncio +import json +import uuid + +import aiohttp +import pytest + + +BASE_URL = "http://127.0.0.1:4000" # change appropriately +ADMIN_KEY = "sk-1234" + + +async def _generate_key(session: aiohttp.ClientSession) -> str: + """Create a fresh virtual key so spend is isolated.""" + url = f"{BASE_URL}/key/generate" + headers = {"Authorization": f"Bearer {ADMIN_KEY}", "Content-Type": "application/json"} + async with session.post(url, headers=headers, json={"models": []}) as resp: + assert resp.status == 200, f"key/generate failed: {await resp.text()}" + data = await resp.json() + return data["key"] + + +async def _get_spend_logs_by_spend_id(session: aiohttp.ClientSession, api_key: str, spend_id: str): + """Query /spend/logs by api_key then filter by spend_id in metadata.""" + url = f"{BASE_URL}/spend/logs?api_key={api_key}" + headers = {"Authorization": f"Bearer {ADMIN_KEY}", "Content-Type": "application/json"} + async with session.get(url, headers=headers) as resp: + assert resp.status == 200, f"spend/logs failed: {await resp.text()}" + all_logs = await resp.json() + if not isinstance(all_logs, list): + return [] + matched = [] + for log in all_logs: + meta = log.get("metadata") + if isinstance(meta, str): + meta = json.loads(meta) + if isinstance(meta, dict): + slm = meta.get("spend_logs_metadata") or {} + if slm.get("spend_id") == spend_id: + matched.append(log) + return matched + + +@pytest.mark.asyncio +@pytest.mark.flaky(retries=3, delay=2) +async def test_v1_messages_streaming_disconnect_has_spend_log(): + """ + 1. Send a streaming POST to /v1/messages. + 2. Read a few SSE chunks, then close the connection (simulating a client + disconnect / interruption). + 3. Wait for the proxy's async spend-tracking pipeline to flush. + 4. Assert that at least one spend-log row exists for the request. + + This PASSES on v1.79.1 and FAILS on the latest main branch. + """ + async with aiohttp.ClientSession( + timeout=aiohttp.ClientTimeout(total=60) + ) as session: + key = await _generate_key(session) + + spend_id = str(uuid.uuid4()) + + headers = { + "Authorization": f"Bearer {key}", + "Content-Type": "application/json", + "x-litellm-spend-logs-metadata": '{"spend_id": "' + spend_id + '"}', + } + + payload = { + "model": "us.anthropic.claude-sonnet-4-5-20250929-v1:0", + "max_tokens": 3000, + "stream": True, + "messages": [ + { + "role": "user", + "content": ( + f"Write several detailed paragraphs (at least 500 words) about the " + f"history of the Roman Empire. Unique id: {uuid.uuid4()}" + ), + } + ], + } + + chunks_read = 0 + + # ---- send the streaming request and disconnect early ---- + async with session.post( + f"{BASE_URL}/v1/messages", json=payload, headers=headers + ) as resp: + assert resp.status == 200, f"/v1/messages failed: {await resp.text()}" + + # Read a handful of SSE chunks, then break out (closes the + # connection, which is the "interruption"). + async for raw_line in resp.content: + line = raw_line.decode("utf-8", errors="replace").strip() + if not line: + continue + chunks_read += 1 + print(f" chunk #{chunks_read}: {line[:120]}") + if chunks_read >= 5: + # We have received enough data — disconnect now. + break + + assert chunks_read >= 3, ( + f"Expected at least 3 chunks before disconnect, got {chunks_read}" + ) + + print( + f"\nDisconnected after {chunks_read} chunks. " + f"Waiting for spend pipeline to flush …" + ) + + # ---- wait & poll for the spend-log entry ---- + spend_data = None + max_retries = 4 + for attempt in range(1, max_retries + 1): + await asyncio.sleep(10) + print(f" spend-log poll attempt {attempt}/{max_retries}") + spend_data = await _get_spend_logs_by_spend_id(session, key, spend_id) + if spend_data and len(spend_data) > 0: + print(f" ✓ found {len(spend_data)} spend-log row(s)") + break + print(" … not found yet") + + # ---- assertions ---- + assert spend_data is not None and len(spend_data) > 0, ( + f"No spend-log entry found for spend_id={spend_id} " + f"after streaming disconnect. " + f"This is the regression: interrupted /v1/messages streams must " + f"still record spend." + ) + + log_entry = spend_data[0] + print( + f"\nSpend-log entry:\n{json.dumps(log_entry, indent=2, default=str)}" + ) + + # A row alone is not enough: the earlier drop-in-finally attempt logged a + # row whose completion tokens reflected only the handful of chunks the + # client drained before disconnecting (~1-15), not the full response + # Bedrock generated and billed. The prompt is written to produce a long + # completion, so the recorded completion tokens must reflect the full + # upstream stream, well above what 5 SSE chunks could carry. + prompt_tokens = log_entry.get("prompt_tokens", 0) + completion_tokens = log_entry.get("completion_tokens", 0) + assert prompt_tokens > 0, ( + "Spend-log row exists but has zero prompt tokens, so usage was not recorded." + ) + assert completion_tokens >= 100, ( + f"Spend-log completion_tokens={completion_tokens} is far below the full " + f"response Bedrock generated and billed. The interrupted stream was billed " + f"on the few chunks the client drained, not the full upstream output. " + f"chunks_read={chunks_read}, prompt_tokens={prompt_tokens}" + ) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py index 6ea9098c228..6bb061c9048 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py @@ -1,3 +1,4 @@ +import asyncio import json import os import sys @@ -243,3 +244,123 @@ def test_incomplete_stream_error_sse_event_is_valid_anthropic_error(): "error": {"type": "api_error", "message": INCOMPLETE_STREAM_ERROR_MESSAGE}, } assert event.endswith("\n\n") + + +# The full stream a provider (Bedrock invoke) generates: a short prefix the +# client reads before disconnecting, then the tail (including the terminal +# ``message_delta`` carrying the real output_tokens) that arrives only after +# the client is gone. output_tokens=64 is the authoritative billed count; a +# naive "log whatever the client drained" implementation would instead see the +# ``message_start`` placeholder (output_tokens=1) and undercount ~64x. +_STREAM_PREFIX = ( + {"type": "message_start", "message": {"id": "msg_1", "usage": {"input_tokens": 52, "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": "The Roman"}}, +) +_STREAM_TAIL = ( + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": " Empire ..."}}, + {"type": "content_block_stop", "index": 0}, + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 64}}, + {"type": "message_stop"}, +) + + +def _output_tokens_from_logged_chunks(chunks: list[bytes]) -> int | None: + """Read the last output_tokens the billing path would see from the SSE bytes.""" + latest: int | None = None + for raw in chunks: + for line in raw.decode().splitlines(): + if not line.startswith("data:"): + continue + data = json.loads(line[len("data:"):].strip()) + usage = data.get("usage") if isinstance(data, dict) else None + if isinstance(usage, dict) and usage.get("output_tokens") is not None: + latest = usage["output_tokens"] + return latest + + +@pytest.mark.asyncio +async def test_async_sse_wrapper_bills_full_stream_after_client_disconnect(): + """ + Regression: on a client disconnect mid-stream the upstream provider keeps + generating (and billing) the full response. The wrapper must keep draining + that upstream to its terminal ``message_delta`` and bill the real + output_tokens (64), not the partial count the client drained before leaving + (the message_start placeholder, 1). + + A ``tail_gated`` event holds back the stream tail until the client has + disconnected, so the tail can only be captured by a drain that survives the + client teardown - exactly the path the previous implementation dropped. + """ + tail_gated = asyncio.Event() + upstream_fully_drained = asyncio.Event() + + async def _gated_stream(): + for event in _STREAM_PREFIX: + yield event + # Block until the test releases the tail (after the client disconnects). + await tail_gated.wait() + for event in _STREAM_TAIL: + yield event + upstream_fully_drained.set() + + iterator = _RecordingLoggingIterator( + litellm_logging_obj=_make_logging_obj("test_bills_full_stream_after_disconnect"), + request_body={}, + ) + + gen = iterator.async_sse_wrapper(_gated_stream()) + + # Client reads the prefix, then disconnects (closes the generator). + client_chunks = [] + async for chunk in gen: + client_chunks.append(chunk) + if len(client_chunks) == len(_STREAM_PREFIX): + break + await gen.aclose() # client disconnect tears down the client-facing generator + + # Now let the provider finish. The detached pump must still be alive. + tail_gated.set() + await asyncio.wait_for(upstream_fully_drained.wait(), timeout=5) + # Give the pump's finally (billing) a turn to run. + for _ in range(100): + if iterator.logged_chunks: + break + await asyncio.sleep(0.01) + + # The client only ever saw the prefix. + assert len(client_chunks) == len(_STREAM_PREFIX) + + # Billing saw the WHOLE stream, including the terminal usage event. + assert iterator.logged_chunks, "pump never billed after client disconnect" + assert _output_tokens_from_logged_chunks(iterator.logged_chunks) == 64 + assert any(c.startswith(b"event: message_stop\n") for c in iterator.logged_chunks) + # No synthetic incomplete-stream error, because the real message_stop arrived. + assert not any(c.startswith(b"event: error\n") for c in iterator.logged_chunks) + + +@pytest.mark.asyncio +async def test_async_sse_wrapper_bills_full_stream_when_client_reads_all(): + """Happy path: when the client drains the whole stream, billing still sees + the terminal output_tokens (64) and the client gets every chunk.""" + tail_gated = asyncio.Event() + tail_gated.set() # no gating; full stream flows immediately + + async def _full_stream(): + for event in (*_STREAM_PREFIX, *_STREAM_TAIL): + yield event + + iterator = _RecordingLoggingIterator( + litellm_logging_obj=_make_logging_obj("test_bills_full_stream_happy_path"), + request_body={}, + ) + client_chunks = [chunk async for chunk in iterator.async_sse_wrapper(_full_stream())] + + for _ in range(100): + if iterator.logged_chunks: + break + await asyncio.sleep(0.01) + + assert len(client_chunks) == len(_STREAM_PREFIX) + len(_STREAM_TAIL) + assert _output_tokens_from_logged_chunks(iterator.logged_chunks) == 64 + assert not any(c.startswith(b"event: error\n") for c in iterator.logged_chunks) diff --git a/type-discipline-budget.json b/type-discipline-budget.json index ab8198304bb..82b562a4a59 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -3,7 +3,7 @@ "limit": 23256 }, "LIT002": { - "limit": 27213 + "limit": 27212 }, "LIT003": { "limit": 269 From ce254867025e149b16fcd8e6ac65e253866060df Mon Sep 17 00:00:00 2001 From: nuernber <> Date: Wed, 5 Aug 2026 17:35:26 -0700 Subject: [PATCH 02/15] fix(anthropic_messages): preserve provider error semantics on upstream stream failure The detached pump previously caught every upstream exception (Bedrock read, decode, provider-response, or chunk-conversion error) and terminated the client stream normally, masking the original provider exception and its status so downstream failure handling never ran. Now, when the upstream fails while the client is still connected, forward the original exception through the queue so the client-facing generator re-raises it and the proxy's failure handling (status code, post_call_failure_hook) runs unchanged. Only when the client has already disconnected, where there is no one to propagate to and no failure hook will fire, fall back to salvaging partial spend from the collected chunks. --- .../messages/streaming_iterator.py | 78 ++++++++++++------- .../messages/test_streaming_iterator.py | 76 ++++++++++++++++++ 2 files changed, 128 insertions(+), 26 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py index abcd817ea18..3458705c7f8 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py @@ -209,37 +209,24 @@ class BaseAnthropicMessagesStreamingIterator: once the client goes away the pump only buffers for billing so the queue can't grow unbounded. + An upstream failure (Bedrock read / decode / chunk-conversion error) + that happens while the client is still connected is forwarded through + the queue and re-raised here, so the original provider exception (and + its status) reaches the proxy's failure handling unchanged rather than + being masked by a generic incomplete-stream event. + This method provides the common logic for both Anthropic and Bedrock implementations. """ from litellm._logging import verbose_proxy_logger - queue: Final[asyncio.Queue[bytes | None]] = asyncio.Queue() + queue: Final[asyncio.Queue[bytes | None | BaseException]] = asyncio.Queue() client_detached: Final = asyncio.Event() async def _pump_upstream() -> None: collected_chunks: Final[list[bytes]] = [] saw_terminal_event = False # rebind-ok: accumulates across the upstream loop - 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(): - queue.put_nowait(encoded_chunk) - except Exception as exc: # noqa: BLE001 # must still flush partial usage in finally, not crash the pump - verbose_proxy_logger.warning( - "async_sse_wrapper upstream pump stopped after %d chunks: %s(%s)", - len(collected_chunks), - type(exc).__name__, - exc, - ) - finally: - if not client_detached.is_set(): - if not saw_terminal_event: - queue.put_nowait(_incomplete_stream_error_sse_event()) - queue.put_nowait(None) + + async def _bill() -> None: try: await self._handle_streaming_logging(collected_chunks) except Exception as exc: # noqa: BLE001 # billing is best-effort; never crash the pump @@ -250,6 +237,42 @@ class BaseAnthropicMessagesStreamingIterator: exc, ) + 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(): + queue.put_nowait(encoded_chunk) + except Exception as exc: # noqa: BLE001 # forward the provider error to a live client, else salvage spend + if not client_detached.is_set(): + # Preserve the provider-specific failure: hand the original + # exception to the client-facing generator so it re-raises + # and the proxy's failure handling (status code, + # post_call_failure_hook) runs. The failure path owns + # logging here, so don't also success-bill. + queue.put_nowait(exc) + return + # Client already disconnected: nothing to propagate to and no + # failure hook will run, so salvage the partial spend instead + # of dropping the request entirely. + verbose_proxy_logger.warning( + "async_sse_wrapper upstream pump failed after client disconnect (%d chunks): %s(%s)", + len(collected_chunks), + type(exc).__name__, + exc, + ) + await _bill() + return + + if not client_detached.is_set(): + if not saw_terminal_event: + queue.put_nowait(_incomplete_stream_error_sse_event()) + queue.put_nowait(None) + await _bill() + pump_task: Final = asyncio.create_task(_pump_upstream()) _UPSTREAM_PUMP_TASKS.add(pump_task) pump_task.add_done_callback(_UPSTREAM_PUMP_TASKS.discard) @@ -259,10 +282,13 @@ class BaseAnthropicMessagesStreamingIterator: item = await queue.get() if item is None: break + if isinstance(item, BaseException): + raise item yield item finally: - # Client-facing generator is being torn down (normal end or a - # disconnect GeneratorExit). Signal the pump to stop enqueueing so - # the queue can't grow unbounded, but let it keep draining upstream - # to its terminal usage event for accurate billing. + # Client-facing generator is being torn down (normal end, a + # re-raised upstream error, or a disconnect GeneratorExit). Signal + # the pump to stop enqueueing so the queue can't grow unbounded, but + # let it keep draining upstream to its terminal usage event for + # accurate billing. client_detached.set() diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py index 6bb061c9048..21eeb514094 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py @@ -364,3 +364,79 @@ async def test_async_sse_wrapper_bills_full_stream_when_client_reads_all(): assert len(client_chunks) == len(_STREAM_PREFIX) + len(_STREAM_TAIL) assert _output_tokens_from_logged_chunks(iterator.logged_chunks) == 64 assert not any(c.startswith(b"event: error\n") for c in iterator.logged_chunks) + + +class _ProviderStreamError(Exception): + """Stand-in for a provider-specific streaming failure carrying a status code.""" + + def __init__(self, message: str, status_code: int): + super().__init__(message) + self.status_code = status_code + + +@pytest.mark.asyncio +async def test_async_sse_wrapper_reraises_upstream_error_to_connected_client(): + """ + Regression: an upstream failure (Bedrock read / decode / chunk-conversion) + before message_stop must propagate the ORIGINAL provider exception to a + still-connected client, so the proxy's failure handling keeps the + provider-specific status. The pump must not swallow it into a generic + api_error event + normal termination. + """ + + async def _failing_stream(): + yield {"type": "message_start", "message": {"id": "msg_1", "usage": {"input_tokens": 52, "output_tokens": 1}}} + yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "partial"}} + raise _ProviderStreamError("bedrock stream blew up", status_code=529) + + iterator = _RecordingLoggingIterator( + litellm_logging_obj=_make_logging_obj("test_reraises_upstream_error"), + request_body={}, + ) + + received = [] + with pytest.raises(_ProviderStreamError) as excinfo: + async for chunk in iterator.async_sse_wrapper(_failing_stream()): + received.append(chunk) + + # Original exception + status preserved, not masked by a synthetic api_error. + assert excinfo.value.status_code == 529 + assert received # the client still got the pre-error chunks + assert not any(c.startswith(b"event: error\n") for c in received) + # On the failure path we do NOT success-bill (failure handling owns logging). + assert iterator.logged_chunks == [] + + +@pytest.mark.asyncio +async def test_async_sse_wrapper_salvages_partial_spend_on_upstream_error_after_disconnect(): + """ + When the upstream errors AFTER the client has already disconnected there is + no live client to re-raise to and no failure hook will run, so the pump + salvages partial spend from what it collected instead of dropping the row. + """ + tail_gated = asyncio.Event() + + async def _gated_failing_stream(): + yield {"type": "message_start", "message": {"id": "msg_1", "usage": {"input_tokens": 52, "output_tokens": 1}}} + yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "partial"}} + await tail_gated.wait() + raise _ProviderStreamError("late failure", status_code=500) + + iterator = _RecordingLoggingIterator( + litellm_logging_obj=_make_logging_obj("test_salvage_partial_on_late_error"), + request_body={}, + ) + + gen = iterator.async_sse_wrapper(_gated_failing_stream()) + received = [await gen.__anext__(), await gen.__anext__()] + await gen.aclose() # client disconnects before the upstream error + + tail_gated.set() # let the upstream raise now, after disconnect + for _ in range(100): + if iterator.logged_chunks: + break + await asyncio.sleep(0.01) + + assert len(received) == 2 + # Partial spend was still recorded rather than the whole request being dropped. + assert iterator.logged_chunks == received From 739447fa4bf93c955676d3f4ebf7e67c9caa7f31 Mon Sep 17 00:00:00 2001 From: nuernber Date: Thu, 6 Aug 2026 08:58:55 -0700 Subject: [PATCH 03/15] fix(anthropic_messages): bound streaming relay queue and cap detached drains The relay queue was unbounded, so a client reading a long stream more slowly than Bedrock produced it let the pump accumulate every pending SSE chunk in memory, and detached post-disconnect drains had no concurrency bound, so an authenticated client could open many large streams and read slowly to pin unbounded worker state. Bound the relay queue and make the pump apply backpressure while the client is connected (it blocks on a full queue, racing the disconnect signal), so a slow reader throttles the upstream read exactly as the old direct yield did. Cap how many detached drains run at once; over the cap a disconnected pump bills what it collected instead of draining further. Detached-drain lifetime is otherwise bounded by the upstream stream/read timeout. Both limits are tunable via env. --- litellm/constants.py | 12 + .../messages/streaming_iterator.py | 231 +++++++++++++----- .../messages/test_streaming_iterator.py | 127 ++++++++++ 3 files changed, 309 insertions(+), 61 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 75f12190b3e..53aedc773d7 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -445,6 +445,18 @@ FIREWORKS_AI_80_B: Final = int(os.getenv("FIREWORKS_AI_80_B", 80)) #### Logging callback constants #### REDACTED_BY_LITELM_STRING: Final = "REDACTED_BY_LITELM" MAX_LANGFUSE_INITIALIZED_CLIENTS: Final = int(os.getenv("MAX_LANGFUSE_INITIALIZED_CLIENTS", 50)) +# Backpressure + lifetime bounds for the /v1/messages streaming relay (see +# BaseAnthropicMessagesStreamingIterator.async_sse_wrapper). The relay queue is +# bounded so a slow client throttles the upstream pump instead of letting it +# buffer the whole response in memory; the detached-drain cap bounds how many +# post-disconnect drains may run concurrently so client behavior can't create +# unbounded worker state. +ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE: Final = int( + os.getenv("ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", "1024") +) +ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS: Final = int( + os.getenv("ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS", "100") +) LOGGING_WORKER_CONCURRENCY: Final = int(os.getenv("LOGGING_WORKER_CONCURRENCY", 100)) # Must be above 0 LOGGING_WORKER_MAX_QUEUE_SIZE: Final = int(os.getenv("LOGGING_WORKER_MAX_QUEUE_SIZE", 50_000)) LOGGING_WORKER_MAX_TIME_PER_COROUTINE: Final = float(os.getenv("LOGGING_WORKER_MAX_TIME_PER_COROUTINE", 20.0)) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py index 3458705c7f8..bff41b9acd2 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py @@ -8,6 +8,10 @@ import httpx from pydantic import TypeAdapter from typing_extensions import TypedDict +from litellm.constants import ( + ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS, + ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE, +) from litellm.litellm_core_utils.core_helpers import process_response_headers from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy.pass_through_endpoints.success_handler import ( @@ -25,6 +29,13 @@ GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ: Final = PassThroughEndpointLogging() # guidance and drop it again from the done callback. _UPSTREAM_PUMP_TASKS: Final[set[asyncio.Task[None]]] = set() # mutable-ok: stdlib strong-ref set for pump tasks +# Rooted set of pumps still draining upstream AFTER their client disconnected. +# Bounds how many detached drains run at once so a burst of slow/abandoned +# streams can't pin unbounded worker memory; a pump over the cap bills what it +# already collected instead of continuing to drain. Only ever touched from the +# event loop, so a plain set + len() check needs no lock. +_DETACHED_STREAM_DRAINS: Final[set[asyncio.Task[None]]] = set() # mutable-ok: bounded strong-ref set, detached drains + INCOMPLETE_STREAM_ERROR_MESSAGE: Final = ( "Provider stream ended before emitting a message_stop event; " "the response is incomplete and any partial content (e.g. tool_use input JSON) may be truncated." @@ -51,6 +62,23 @@ def _is_terminal_stream_chunk(chunk: object) -> bool: return _is_message_stop_chunk(chunk) or _is_provider_error_chunk(chunk) +def _try_claim_detached_drain_slot() -> bool: + """Claim a detached-drain slot for the current task, bounding concurrency. + + Returns True if a slot was claimed (the caller may keep draining upstream + for billing) or False if the cap is already reached (the caller should stop + and bill what it has). Only touched from the event loop, so the check + + insert need no lock. + """ + if len(_DETACHED_STREAM_DRAINS) >= ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS: + return False + current_task: Final = asyncio.current_task() + if current_task is not None: + _DETACHED_STREAM_DRAINS.add(current_task) + current_task.add_done_callback(_DETACHED_STREAM_DRAINS.discard) + return True + + def _incomplete_stream_error_sse_event() -> bytes: payload: Final = json.dumps( { @@ -205,9 +233,18 @@ class BaseAnthropicMessagesStreamingIterator: generating and billing the full response regardless of the client, so draining it to completion is what lets spend tracking see the real terminal ``message_delta`` / ``message_stop`` usage instead of a - truncated placeholder count. Chunks reach the client through a queue; - once the client goes away the pump only buffers for billing so the queue - can't grow unbounded. + truncated placeholder count. + + Chunks reach the client through a bounded queue. While the client is + connected the pump blocks on a full queue (racing the disconnect + signal), so a slow reader throttles the upstream read exactly as the old + direct ``yield`` did instead of letting the whole response buffer in + memory. Once the client goes away the pump stops enqueueing and only + 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. An upstream failure (Bedrock read / decode / chunk-conversion error) that happens while the client is still connected is forwarded through @@ -217,63 +254,12 @@ class BaseAnthropicMessagesStreamingIterator: This method provides the common logic for both Anthropic and Bedrock implementations. """ - from litellm._logging import verbose_proxy_logger - - queue: Final[asyncio.Queue[bytes | None | BaseException]] = asyncio.Queue() + queue: Final[asyncio.Queue[bytes | None | BaseException]] = asyncio.Queue( + maxsize=ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE + ) client_detached: Final = asyncio.Event() - async def _pump_upstream() -> None: - collected_chunks: Final[list[bytes]] = [] - saw_terminal_event = False # rebind-ok: accumulates across the upstream loop - - async def _bill() -> None: - try: - await self._handle_streaming_logging(collected_chunks) - except Exception as exc: # noqa: BLE001 # billing is best-effort; never crash the pump - verbose_proxy_logger.warning( - "async_sse_wrapper billing failed after %d chunks: %s(%s)", - len(collected_chunks), - type(exc).__name__, - exc, - ) - - 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(): - queue.put_nowait(encoded_chunk) - except Exception as exc: # noqa: BLE001 # forward the provider error to a live client, else salvage spend - if not client_detached.is_set(): - # Preserve the provider-specific failure: hand the original - # exception to the client-facing generator so it re-raises - # and the proxy's failure handling (status code, - # post_call_failure_hook) runs. The failure path owns - # logging here, so don't also success-bill. - queue.put_nowait(exc) - return - # Client already disconnected: nothing to propagate to and no - # failure hook will run, so salvage the partial spend instead - # of dropping the request entirely. - verbose_proxy_logger.warning( - "async_sse_wrapper upstream pump failed after client disconnect (%d chunks): %s(%s)", - len(collected_chunks), - type(exc).__name__, - exc, - ) - await _bill() - return - - if not client_detached.is_set(): - if not saw_terminal_event: - queue.put_nowait(_incomplete_stream_error_sse_event()) - queue.put_nowait(None) - await _bill() - - pump_task: Final = asyncio.create_task(_pump_upstream()) + pump_task: Final = asyncio.create_task(self._pump_upstream_to_queue(completion_stream, queue, client_detached)) _UPSTREAM_PUMP_TASKS.add(pump_task) pump_task.add_done_callback(_UPSTREAM_PUMP_TASKS.discard) @@ -288,7 +274,130 @@ class BaseAnthropicMessagesStreamingIterator: finally: # Client-facing generator is being torn down (normal end, a # re-raised upstream error, or a disconnect GeneratorExit). Signal - # the pump to stop enqueueing so the queue can't grow unbounded, but - # let it keep draining upstream to its terminal usage event for - # accurate billing. + # the pump to stop enqueueing and unblock any backpressure-blocked + # put; the pump then either finishes billing or drains detached + # (subject to the cap) for accurate usage. client_detached.set() + + async def _bill_collected_chunks( + self, + collected_chunks: list[bytes], # mutable-ok: SSE buffer forwarded to list-typed _handle_streaming_logging + ) -> None: + from litellm._logging import verbose_proxy_logger + + try: + await self._handle_streaming_logging(collected_chunks) + except Exception as exc: # noqa: BLE001 # billing is best-effort; never crash the pump + verbose_proxy_logger.warning( + "async_sse_wrapper billing failed after %d chunks: %s(%s)", + len(collected_chunks), + type(exc).__name__, + exc, + ) + + @staticmethod + async def _enqueue_for_client( + queue: "asyncio.Queue[bytes | None | BaseException]", + client_detached: "asyncio.Event", + item: bytes | None | BaseException, + ) -> bool: + """Deliver one item to the client, applying backpressure. + + Returns True if the item was queued, False if the client disconnected + before there was room (the item is then dropped, since a gone client + can't receive it). Never blocks once the client has detached. + """ + if client_detached.is_set(): + return False + try: + queue.put_nowait(item) + return True + except asyncio.QueueFull: + pass + 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) + finally: + if not detached_task.done(): + detached_task.cancel() + if put_task.done() and not put_task.cancelled(): + return True + put_task.cancel() + return False + + async def _pump_upstream_to_queue( + self, + completion_stream: AsyncIterator[bytes | GenericStreamingChunk | ModelResponseStream | dict], + queue: "asyncio.Queue[bytes | None | BaseException]", + client_detached: "asyncio.Event", + ) -> None: + """Drain the whole upstream into ``queue`` (backpressured) and bill once. + + Runs detached so a client disconnect can't interrupt the upstream read; + see ``async_sse_wrapper`` for the full rationale. Returns after billing. + """ + from litellm._logging import verbose_proxy_logger + + 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 + 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 + # Client has gone: keep draining only to reach the terminal usage + # event for billing, but claim a detached-drain slot first; over + # the cap, bill what we have rather than pinning more memory. + 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 without draining the rest of the upstream stream", + ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS, + len(collected_chunks), + ) + await self._bill_collected_chunks(collected_chunks) + return + draining_detached = True + except Exception as exc: # noqa: BLE001 # forward the provider error to a live client, else salvage spend + await self._handle_pump_upstream_error(queue, client_detached, collected_chunks, exc) + return + + if not client_detached.is_set(): + if not saw_terminal_event: + await self._enqueue_for_client(queue, client_detached, _incomplete_stream_error_sse_event()) + await self._enqueue_for_client(queue, client_detached, None) + await self._bill_collected_chunks(collected_chunks) + + async def _handle_pump_upstream_error( + self, + queue: "asyncio.Queue[bytes | None | BaseException]", + client_detached: "asyncio.Event", + collected_chunks: list[bytes], # mutable-ok: SSE buffer forwarded to list-typed _bill_collected_chunks + exc: BaseException, + ) -> None: + from litellm._logging import verbose_proxy_logger + + if not client_detached.is_set() and await self._enqueue_for_client(queue, client_detached, exc): + # Preserve the provider-specific failure: the client-facing + # generator re-raises it and the proxy's failure handling (status + # code, post_call_failure_hook) runs. That path owns logging, so + # don't also success-bill. + return + # Client already gone (or disconnected before the error reached it): no + # failure hook will run, so salvage the partial spend instead of + # dropping the request entirely. + verbose_proxy_logger.warning( + "async_sse_wrapper upstream pump failed after client disconnect (%d chunks): %s(%s)", + len(collected_chunks), + type(exc).__name__, + exc, + ) + await self._bill_collected_chunks(collected_chunks) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py index 21eeb514094..4afbe8fe833 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py @@ -9,6 +9,7 @@ import pytest sys.path.insert(0, os.path.abspath("../../../../..")) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.anthropic.experimental_pass_through.messages import streaming_iterator as streaming_iterator_module from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( INCOMPLETE_STREAM_ERROR_MESSAGE, BaseAnthropicMessagesStreamingIterator, @@ -440,3 +441,129 @@ async def test_async_sse_wrapper_salvages_partial_spend_on_upstream_error_after_ assert len(received) == 2 # Partial spend was still recorded rather than the whole request being dropped. assert iterator.logged_chunks == received + + +@pytest.mark.asyncio +async def test_async_sse_wrapper_applies_backpressure_to_slow_client(monkeypatch): + """ + Regression: the relay queue is bounded, so a slow client throttles the + upstream read instead of letting the pump buffer the whole response in + memory. With a tiny queue and a client that reads a single chunk, the pump + must stall after producing only a queue's worth of chunks ahead, not race + to the end of a large stream. + """ + monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", 2) + + total = 200 + produced = 0 + + async def _fast_stream(): + nonlocal produced + for i in range(total): + produced += 1 + yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": f"t{i}"}} + + iterator = _make_iterator("test_backpressure_slow_client") + gen = iterator.async_sse_wrapper(_fast_stream()) + try: + await gen.__anext__() # read exactly one chunk, then stall + # Let the pump run as far as the bounded queue permits. + for _ in range(500): + await asyncio.sleep(0) + # Bounded by queue maxsize + the one in-flight put + the one delivered + # chunk; nowhere near the full 200-chunk stream. + assert produced <= 2 + 3, f"pump ran ahead unthrottled: produced {produced} of {total}" + assert produced < total + finally: + await gen.aclose() + + +@pytest.mark.asyncio +async def test_async_sse_wrapper_bills_partial_when_detached_drain_cap_reached(monkeypatch): + """ + Regression: when the concurrent detached-drain cap is already reached, a + pump whose client has disconnected must bill what it collected instead of + continuing to drain (and accumulating) the rest of a large upstream stream, + so slow/abandoned clients can't pin unbounded worker state. + + The cap slot set is pre-occupied so the single slot is unavailable when this + pump reaches its first post-disconnect chunk; that isolates the cap decision + from multi-pump scheduling races. + """ + monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS", 1) + monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", 4) + + # Occupy the only detached-drain slot with a placeholder task. + async def _hold_slot(): + await asyncio.sleep(3600) + + holder = asyncio.ensure_future(_hold_slot()) + streaming_iterator_module._DETACHED_STREAM_DRAINS.add(holder) + tail_reached = False + + async def _long_stream(): + nonlocal tail_reached + yield {"type": "message_start", "message": {"id": "m", "usage": {"input_tokens": 5, "output_tokens": 1}}} + yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "x"}} + # These arrive only after the client has disconnected. + for i in range(100): + yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": f"more{i}"}} + tail_reached = True + 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("drain_cap_full"), request_body={}) + try: + gen = iterator.async_sse_wrapper(_long_stream()) + await gen.__anext__() # message_start + await gen.__anext__() # first delta + await gen.aclose() # client disconnects; 100+ chunks remain upstream + + for _ in range(200): + if iterator.logged_chunks: + break + await asyncio.sleep(0) + + # Billed a small bounded partial (the prefix plus at most a queue's + # worth the pump ran ahead before disconnect) without draining the + # 100-chunk tail. The exact count depends on how far the bounded queue + # let the pump run ahead, so assert the bound, not an exact number. + assert iterator.logged_chunks, "capped pump never billed" + assert len(iterator.logged_chunks) <= 2 + streaming_iterator_module.ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE + assert len(iterator.logged_chunks) < 100 + assert not any(c.startswith(b"event: message_stop\n") for c in iterator.logged_chunks) + assert tail_reached is False, "pump kept draining past the cap instead of stopping" + finally: + holder.cancel() + streaming_iterator_module._DETACHED_STREAM_DRAINS.discard(holder) + + +@pytest.mark.asyncio +async def test_async_sse_wrapper_drains_detached_when_cap_available(monkeypatch): + """Complement to the cap test: with a slot free, a disconnected pump drains + the full upstream and bills the terminal usage, and releases its slot after.""" + monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS", 1) + monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", 4) + + async def _stream(): + yield {"type": "message_start", "message": {"id": "m", "usage": {"input_tokens": 5, "output_tokens": 1}}} + yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "x"}} + for i in range(20): + yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": f"m{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("drain_cap_free"), request_body={}) + gen = iterator.async_sse_wrapper(_stream()) + await gen.__anext__() + await gen.__anext__() + await gen.aclose() + + for _ in range(300): + if iterator.logged_chunks: + break + await asyncio.sleep(0.01) + + assert any(c.startswith(b"event: message_stop\n") for c in iterator.logged_chunks) + # Slot released once the drain finished. + assert len(streaming_iterator_module._DETACHED_STREAM_DRAINS) == 0 From 321779138ef74cd5fd2b3f4f232fb5a0fd7a1158 Mon Sep 17 00:00:00 2001 From: nuernber Date: Thu, 6 Aug 2026 09:36:53 -0700 Subject: [PATCH 04/15] test(env_keys): exclude internal streaming tuning vars from documentation checks Add ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS and ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE to the excluded set. These are advanced internal infrastructure parameters for streaming/queue management with sensible defaults that most users should not modify. --- tests/documentation_tests/test_env_keys.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/tests/documentation_tests/test_env_keys.py b/tests/documentation_tests/test_env_keys.py index b91c404b2eb..3652378503e 100644 --- a/tests/documentation_tests/test_env_keys.py +++ b/tests/documentation_tests/test_env_keys.py @@ -33,6 +33,13 @@ EXCLUDED_ROLLOUT_FLAGS = { "LITELLM_RUST", } +# Internal infrastructure tuning parameters for streaming/queue management +# 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_STREAM_RELAY_QUEUE_MAXSIZE", +} + EXCLUDED_TERMINAL_VARS = { "TERM", "TERM_PROGRAM", @@ -50,7 +57,9 @@ EXCLUDED_TERMINAL_VARS = { "ALACRITTY_SOCKET", } -EXCLUDED_KEYS = frozenset(EXCLUDED_TERMINAL_VARS | EXCLUDED_GUARD_ONLY_VARS | EXCLUDED_ROLLOUT_FLAGS) +EXCLUDED_KEYS = frozenset( + EXCLUDED_TERMINAL_VARS | EXCLUDED_GUARD_ONLY_VARS | EXCLUDED_ROLLOUT_FLAGS | EXCLUDED_INTERNAL_TUNING_VARS +) # Directories to skip (dependencies, venvs, caches) - only scan litellm source SKIP_DIRS = { From a85a9e1186df3ca630101b1a97767ec5c586a865 Mon Sep 17 00:00:00 2001 From: nuernber Date: Thu, 6 Aug 2026 10:46:54 -0700 Subject: [PATCH 05/15] fix(anthropic_messages): strip inline comments, add abort-upstream regression test Strip net-new inline # blocks from streaming_iterator.py, the unit test file, and the live-proxy regression test to comply with the no-new-comments rule. Add test_async_sse_wrapper_aborts_upstream_when_detached_drain_cap_reached: verifies that when the detached-drain cap is already full, the pump calls aclose() on the upstream so the provider stops generating and billing instead of continuing to stream while we record only the partial prefix. Also fixes LIT001 (bare dict in AsyncIterator union) by replacing dict with Mapping[str, object] across all three stream-type annotations, and adds the required LIT003 reason strings to the three noqa: BLE001 directives. --- .../messages/streaming_iterator.py | 60 ++++++------ ..._v1_messages_streaming_disconnect_spend.py | 12 --- .../messages/test_streaming_iterator.py | 94 +++++++++++++------ 3 files changed, 94 insertions(+), 72 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py index bff41b9acd2..e302896ff5e 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py @@ -1,6 +1,6 @@ import asyncio import json -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Mapping from datetime import datetime from typing import Any, Final, Protocol, runtime_checkable @@ -22,18 +22,7 @@ from litellm.types.utils import GenericStreamingChunk, ModelResponseStream GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ: Final = PassThroughEndpointLogging() -# asyncio holds only a weak reference to a bare create_task() result, so a -# fire-and-forget task can be garbage-collected mid-run. The upstream pump -# below must outlive the client-facing generator (which is closed on client -# disconnect), so root every pump task in a module-level set per the stdlib -# guidance and drop it again from the done callback. _UPSTREAM_PUMP_TASKS: Final[set[asyncio.Task[None]]] = set() # mutable-ok: stdlib strong-ref set for pump tasks - -# Rooted set of pumps still draining upstream AFTER their client disconnected. -# Bounds how many detached drains run at once so a burst of slow/abandoned -# streams can't pin unbounded worker memory; a pump over the cap bills what it -# already collected instead of continuing to drain. Only ever touched from the -# event loop, so a plain set + len() check needs no lock. _DETACHED_STREAM_DRAINS: Final[set[asyncio.Task[None]]] = set() # mutable-ok: bounded strong-ref set, detached drains INCOMPLETE_STREAM_ERROR_MESSAGE: Final = ( @@ -221,7 +210,7 @@ class BaseAnthropicMessagesStreamingIterator: async def async_sse_wrapper( self, - completion_stream: AsyncIterator[bytes | GenericStreamingChunk | ModelResponseStream | dict], + completion_stream: AsyncIterator[bytes | GenericStreamingChunk | ModelResponseStream | Mapping[str, object]], ) -> AsyncIterator[bytes]: """ Generic async SSE wrapper that converts streaming chunks to SSE format @@ -272,11 +261,6 @@ class BaseAnthropicMessagesStreamingIterator: raise item yield item finally: - # Client-facing generator is being torn down (normal end, a - # re-raised upstream error, or a disconnect GeneratorExit). Signal - # the pump to stop enqueueing and unblock any backpressure-blocked - # put; the pump then either finishes billing or drains detached - # (subject to the cap) for accurate usage. client_detached.set() async def _bill_collected_chunks( @@ -295,6 +279,22 @@ class BaseAnthropicMessagesStreamingIterator: exc, ) + @staticmethod + async def _abort_upstream( + completion_stream: AsyncIterator[bytes | GenericStreamingChunk | ModelResponseStream | Mapping[str, object]], + ) -> None: + """Close the upstream provider stream so it stops generating and billing.""" + from litellm._logging import verbose_proxy_logger + + try: + await aclose_if_supported(completion_stream) + except Exception as exc: # noqa: BLE001 # abort is best-effort; log and continue + verbose_proxy_logger.warning( + "async_sse_wrapper failed to abort upstream stream: %s(%s)", + type(exc).__name__, + exc, + ) + @staticmethod async def _enqueue_for_client( queue: "asyncio.Queue[bytes | None | BaseException]", @@ -328,7 +328,7 @@ class BaseAnthropicMessagesStreamingIterator: async def _pump_upstream_to_queue( self, - completion_stream: AsyncIterator[bytes | GenericStreamingChunk | ModelResponseStream | dict], + completion_stream: AsyncIterator[bytes | GenericStreamingChunk | ModelResponseStream | Mapping[str, object]], queue: "asyncio.Queue[bytes | None | BaseException]", client_detached: "asyncio.Event", ) -> None: @@ -352,21 +352,19 @@ class BaseAnthropicMessagesStreamingIterator: if not client_detached.is_set(): await self._enqueue_for_client(queue, client_detached, encoded_chunk) continue - # Client has gone: keep draining only to reach the terminal usage - # event for billing, but claim a detached-drain slot first; over - # the cap, bill what we have rather than pinning more memory. 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 without draining the rest of the upstream stream", + "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) + await self._abort_upstream(completion_stream) return draining_detached = True - except Exception as exc: # noqa: BLE001 # forward the provider error to a live client, else salvage spend + 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 @@ -383,17 +381,17 @@ class BaseAnthropicMessagesStreamingIterator: collected_chunks: list[bytes], # mutable-ok: SSE buffer forwarded to list-typed _bill_collected_chunks exc: BaseException, ) -> None: + """Forward a provider error to a still-connected client, else salvage partial spend. + + Handing the original exception to the client-facing generator lets it + re-raise so the proxy's failure handling keeps the provider status and + owns logging (no success-bill). If the client already went away, no + failure hook runs, so bill the partial instead of dropping the request. + """ from litellm._logging import verbose_proxy_logger if not client_detached.is_set() and await self._enqueue_for_client(queue, client_detached, exc): - # Preserve the provider-specific failure: the client-facing - # generator re-raises it and the proxy's failure handling (status - # code, post_call_failure_hook) runs. That path owns logging, so - # don't also success-bill. return - # Client already gone (or disconnected before the error reached it): no - # failure hook will run, so salvage the partial spend instead of - # dropping the request entirely. verbose_proxy_logger.warning( "async_sse_wrapper upstream pump failed after client disconnect (%d chunks): %s(%s)", len(collected_chunks), diff --git a/tests/pass_through_tests/test_v1_messages_streaming_disconnect_spend.py b/tests/pass_through_tests/test_v1_messages_streaming_disconnect_spend.py index 5cd80b5ec6f..e69de720ea4 100644 --- a/tests/pass_through_tests/test_v1_messages_streaming_disconnect_spend.py +++ b/tests/pass_through_tests/test_v1_messages_streaming_disconnect_spend.py @@ -96,14 +96,11 @@ async def test_v1_messages_streaming_disconnect_has_spend_log(): chunks_read = 0 - # ---- send the streaming request and disconnect early ---- async with session.post( f"{BASE_URL}/v1/messages", json=payload, headers=headers ) as resp: assert resp.status == 200, f"/v1/messages failed: {await resp.text()}" - # Read a handful of SSE chunks, then break out (closes the - # connection, which is the "interruption"). async for raw_line in resp.content: line = raw_line.decode("utf-8", errors="replace").strip() if not line: @@ -111,7 +108,6 @@ async def test_v1_messages_streaming_disconnect_has_spend_log(): chunks_read += 1 print(f" chunk #{chunks_read}: {line[:120]}") if chunks_read >= 5: - # We have received enough data — disconnect now. break assert chunks_read >= 3, ( @@ -123,7 +119,6 @@ async def test_v1_messages_streaming_disconnect_has_spend_log(): f"Waiting for spend pipeline to flush …" ) - # ---- wait & poll for the spend-log entry ---- spend_data = None max_retries = 4 for attempt in range(1, max_retries + 1): @@ -135,7 +130,6 @@ async def test_v1_messages_streaming_disconnect_has_spend_log(): break print(" … not found yet") - # ---- assertions ---- assert spend_data is not None and len(spend_data) > 0, ( f"No spend-log entry found for spend_id={spend_id} " f"after streaming disconnect. " @@ -148,12 +142,6 @@ async def test_v1_messages_streaming_disconnect_has_spend_log(): f"\nSpend-log entry:\n{json.dumps(log_entry, indent=2, default=str)}" ) - # A row alone is not enough: the earlier drop-in-finally attempt logged a - # row whose completion tokens reflected only the handful of chunks the - # client drained before disconnecting (~1-15), not the full response - # Bedrock generated and billed. The prompt is written to produce a long - # completion, so the recorded completion tokens must reflect the full - # upstream stream, well above what 5 SSE chunks could carry. prompt_tokens = log_entry.get("prompt_tokens", 0) completion_tokens = log_entry.get("completion_tokens", 0) assert prompt_tokens > 0, ( diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py index 4afbe8fe833..6ad2f1774da 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py @@ -247,12 +247,6 @@ def test_incomplete_stream_error_sse_event_is_valid_anthropic_error(): assert event.endswith("\n\n") -# The full stream a provider (Bedrock invoke) generates: a short prefix the -# client reads before disconnecting, then the tail (including the terminal -# ``message_delta`` carrying the real output_tokens) that arrives only after -# the client is gone. output_tokens=64 is the authoritative billed count; a -# naive "log whatever the client drained" implementation would instead see the -# ``message_start`` placeholder (output_tokens=1) and undercount ~64x. _STREAM_PREFIX = ( {"type": "message_start", "message": {"id": "msg_1", "usage": {"input_tokens": 52, "output_tokens": 1}}}, {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, @@ -299,7 +293,6 @@ async def test_async_sse_wrapper_bills_full_stream_after_client_disconnect(): async def _gated_stream(): for event in _STREAM_PREFIX: yield event - # Block until the test releases the tail (after the client disconnects). await tail_gated.wait() for event in _STREAM_TAIL: yield event @@ -312,31 +305,25 @@ async def test_async_sse_wrapper_bills_full_stream_after_client_disconnect(): gen = iterator.async_sse_wrapper(_gated_stream()) - # Client reads the prefix, then disconnects (closes the generator). client_chunks = [] async for chunk in gen: client_chunks.append(chunk) if len(client_chunks) == len(_STREAM_PREFIX): break - await gen.aclose() # client disconnect tears down the client-facing generator + await gen.aclose() - # Now let the provider finish. The detached pump must still be alive. tail_gated.set() await asyncio.wait_for(upstream_fully_drained.wait(), timeout=5) - # Give the pump's finally (billing) a turn to run. for _ in range(100): if iterator.logged_chunks: break await asyncio.sleep(0.01) - # The client only ever saw the prefix. assert len(client_chunks) == len(_STREAM_PREFIX) - # Billing saw the WHOLE stream, including the terminal usage event. assert iterator.logged_chunks, "pump never billed after client disconnect" assert _output_tokens_from_logged_chunks(iterator.logged_chunks) == 64 assert any(c.startswith(b"event: message_stop\n") for c in iterator.logged_chunks) - # No synthetic incomplete-stream error, because the real message_stop arrived. assert not any(c.startswith(b"event: error\n") for c in iterator.logged_chunks) @@ -400,11 +387,9 @@ async def test_async_sse_wrapper_reraises_upstream_error_to_connected_client(): async for chunk in iterator.async_sse_wrapper(_failing_stream()): received.append(chunk) - # Original exception + status preserved, not masked by a synthetic api_error. assert excinfo.value.status_code == 529 - assert received # the client still got the pre-error chunks + assert received assert not any(c.startswith(b"event: error\n") for c in received) - # On the failure path we do NOT success-bill (failure handling owns logging). assert iterator.logged_chunks == [] @@ -439,7 +424,6 @@ async def test_async_sse_wrapper_salvages_partial_spend_on_upstream_error_after_ await asyncio.sleep(0.01) assert len(received) == 2 - # Partial spend was still recorded rather than the whole request being dropped. assert iterator.logged_chunks == received @@ -466,12 +450,9 @@ async def test_async_sse_wrapper_applies_backpressure_to_slow_client(monkeypatch iterator = _make_iterator("test_backpressure_slow_client") gen = iterator.async_sse_wrapper(_fast_stream()) try: - await gen.__anext__() # read exactly one chunk, then stall - # Let the pump run as far as the bounded queue permits. + await gen.__anext__() for _ in range(500): await asyncio.sleep(0) - # Bounded by queue maxsize + the one in-flight put + the one delivered - # chunk; nowhere near the full 200-chunk stream. assert produced <= 2 + 3, f"pump ran ahead unthrottled: produced {produced} of {total}" assert produced < total finally: @@ -493,7 +474,6 @@ async def test_async_sse_wrapper_bills_partial_when_detached_drain_cap_reached(m monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS", 1) monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", 4) - # Occupy the only detached-drain slot with a placeholder task. async def _hold_slot(): await asyncio.sleep(3600) @@ -505,7 +485,6 @@ async def test_async_sse_wrapper_bills_partial_when_detached_drain_cap_reached(m nonlocal tail_reached yield {"type": "message_start", "message": {"id": "m", "usage": {"input_tokens": 5, "output_tokens": 1}}} yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "x"}} - # These arrive only after the client has disconnected. for i in range(100): yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": f"more{i}"}} tail_reached = True @@ -524,10 +503,6 @@ async def test_async_sse_wrapper_bills_partial_when_detached_drain_cap_reached(m break await asyncio.sleep(0) - # Billed a small bounded partial (the prefix plus at most a queue's - # worth the pump ran ahead before disconnect) without draining the - # 100-chunk tail. The exact count depends on how far the bounded queue - # let the pump run ahead, so assert the bound, not an exact number. assert iterator.logged_chunks, "capped pump never billed" assert len(iterator.logged_chunks) <= 2 + streaming_iterator_module.ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE assert len(iterator.logged_chunks) < 100 @@ -538,6 +513,68 @@ async def test_async_sse_wrapper_bills_partial_when_detached_drain_cap_reached(m streaming_iterator_module._DETACHED_STREAM_DRAINS.discard(holder) +@pytest.mark.asyncio +async def test_async_sse_wrapper_aborts_upstream_when_detached_drain_cap_reached(monkeypatch): + """ + Regression: when the cap is full and a disconnected pump bails, it must call + aclose on the upstream stream so the provider stops generating and billing, + not continue running the stream while we record only the partial prefix. + """ + monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS", 1) + monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", 4) + + async def _hold_slot(): + await asyncio.sleep(3600) + + holder = asyncio.ensure_future(_hold_slot()) + streaming_iterator_module._DETACHED_STREAM_DRAINS.add(holder) + + class _AbortableStream: + def __init__(self): + self.aclose_called = False + self._remaining = iter( + ( + {"type": "message_start", "message": {"id": "m", "usage": {"input_tokens": 5, "output_tokens": 1}}}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "x"}}, + ) + + tuple( + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": f"t{i}"}} + for i in range(50) + ) + ) + + def __aiter__(self): + return self + + async def __anext__(self): + try: + return next(self._remaining) + except StopIteration: + raise StopAsyncIteration + + async def aclose(self): + self.aclose_called = True + + stream = _AbortableStream() + iterator = _RecordingLoggingIterator(litellm_logging_obj=_make_logging_obj("abort_upstream_at_cap"), request_body={}) + try: + gen = iterator.async_sse_wrapper(stream) + await gen.__anext__() + await gen.__anext__() + await gen.aclose() + + for _ in range(200): + if iterator.logged_chunks: + break + await asyncio.sleep(0) + + assert iterator.logged_chunks, "capped pump never billed" + assert stream.aclose_called, "upstream aclose was not called when the detached-drain cap was reached" + finally: + holder.cancel() + streaming_iterator_module._DETACHED_STREAM_DRAINS.discard(holder) + + @pytest.mark.asyncio async def test_async_sse_wrapper_drains_detached_when_cap_available(monkeypatch): """Complement to the cap test: with a slot free, a disconnected pump drains @@ -565,5 +602,4 @@ async def test_async_sse_wrapper_drains_detached_when_cap_available(monkeypatch) await asyncio.sleep(0.01) assert any(c.startswith(b"event: message_stop\n") for c in iterator.logged_chunks) - # Slot released once the drain finished. assert len(streaming_iterator_module._DETACHED_STREAM_DRAINS) == 0 From 3db47daefe4ae8c938a37bcb68076b8a05f586b5 Mon Sep 17 00:00:00 2001 From: nuernber Date: Thu, 6 Aug 2026 11:17:03 -0700 Subject: [PATCH 06/15] test(anthropic_messages): add unit tests for _abort_upstream and _enqueue_for_client edge cases Add test_abort_upstream_logs_warning_when_aclose_raises: verifies that _abort_upstream swallows and logs any exception raised by the upstream's aclose() method instead of propagating it. Add test_enqueue_for_client_returns_false_when_already_detached: verifies that _enqueue_for_client returns False immediately without touching the queue when client_detached is already set before the call. Add test_enqueue_for_ --- .../messages/test_streaming_iterator.py | 62 +++++++++++++++++++ 1 file changed, 62 insertions(+) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py index 6ad2f1774da..f5399893a3b 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py @@ -575,6 +575,68 @@ async def test_async_sse_wrapper_aborts_upstream_when_detached_drain_cap_reached streaming_iterator_module._DETACHED_STREAM_DRAINS.discard(holder) +@pytest.mark.asyncio +async def test_abort_upstream_logs_warning_when_aclose_raises(caplog): + """_abort_upstream must swallow and log any exception from aclose().""" + import logging + + class _ExplodingStream: + def __aiter__(self): + return self + + async def __anext__(self): + raise StopAsyncIteration + + async def aclose(self): + raise RuntimeError("aclose exploded") + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + await BaseAnthropicMessagesStreamingIterator._abort_upstream(_ExplodingStream()) + + assert any("abort" in r.message and "RuntimeError" in r.message for r in caplog.records) + + +@pytest.mark.asyncio +async def test_enqueue_for_client_returns_false_when_already_detached(): + """_enqueue_for_client must return False immediately (without touching the queue) + when client_detached is already set before the call.""" + from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( + BaseAnthropicMessagesStreamingIterator, + ) + + queue: asyncio.Queue[bytes | None | BaseException] = asyncio.Queue(maxsize=1) + client_detached = asyncio.Event() + client_detached.set() + + result = await BaseAnthropicMessagesStreamingIterator._enqueue_for_client(queue, client_detached, b"chunk") + assert result is False + assert queue.empty() + + +@pytest.mark.asyncio +async def test_enqueue_for_client_returns_false_when_client_detaches_while_queue_full(): + """_enqueue_for_client must return False (and cancel the put) when the queue + is full and client_detached fires before space becomes available.""" + from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( + BaseAnthropicMessagesStreamingIterator, + ) + + queue: asyncio.Queue[bytes | None | BaseException] = asyncio.Queue(maxsize=1) + queue.put_nowait(b"already-full") + + client_detached = asyncio.Event() + + async def _set_detached_soon(): + await asyncio.sleep(0.01) + client_detached.set() + + asyncio.create_task(_set_detached_soon()) + result = await BaseAnthropicMessagesStreamingIterator._enqueue_for_client(queue, client_detached, b"new-chunk") + assert result is False + assert queue.qsize() == 1 + assert queue.get_nowait() == b"already-full" + + @pytest.mark.asyncio async def test_async_sse_wrapper_drains_detached_when_cap_available(monkeypatch): """Complement to the cap test: with a slot free, a disconnected pump drains From 13370adbf4bc3acc0af3b8d7916d834e71931f7f Mon Sep 17 00:00:00 2001 From: nuernber Date: Tue, 11 Aug 2026 14:43:40 -0700 Subject: [PATCH 07/15] chore: ratchet down basedpyright-code-budget after merge --- basedpyright-code-budget.json | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 6e3cbdff9d0..a0042234935 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -57,7 +57,7 @@ "limit": 5719 }, "reportMissingTypeArgument": { - "limit": 15657 + "limit": 15656 }, "reportMissingTypeStubs": { "limit": 40 @@ -99,19 +99,19 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 44832 + "limit": 44831 }, "reportUnknownLambdaType": { "limit": 113 }, "reportUnknownMemberType": { - "limit": 39269 + "limit": 39267 }, "reportUnknownParameterType": { - "limit": 19988 + "limit": 19987 }, "reportUnknownVariableType": { - "limit": 30923 + "limit": 30922 }, "reportUnnecessaryCast": { "limit": 118 From b2df72f980c1f65dbf2e7200ecef4aa1a2a45aa5 Mon Sep 17 00:00:00 2001 From: nuernber Date: Mon, 31 Aug 2026 09:05:37 -0700 Subject: [PATCH 08/15] chore: ratchet down lint/type budgets after disconnect-billing merge fix --- basedpyright-code-budget.json | 10 +++++----- type-discipline-budget.json | 6 +++--- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index a542d025e74..f62d5f95256 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -57,7 +57,7 @@ "limit": 5658 }, "reportMissingTypeArgument": { - "limit": 15656 + "limit": 15655 }, "reportMissingTypeStubs": { "limit": 40 @@ -99,19 +99,19 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 44831 + "limit": 44829 }, "reportUnknownLambdaType": { "limit": 109 }, "reportUnknownMemberType": { - "limit": 39267 + "limit": 39265 }, "reportUnknownParameterType": { - "limit": 19987 + "limit": 19986 }, "reportUnknownVariableType": { - "limit": 30922 + "limit": 30921 }, "reportUnnecessaryCast": { "limit": 117 diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 1f2651f5f25..c5349689ec0 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 23002 + "limit": 23001 }, "LIT002": { - "limit": 27145 + "limit": 27144 }, "LIT003": { "limit": 269 @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16730 + "limit": 16729 }, "LIT011": { "limit": 5577 From e1fece511a20298f16926a8b595850a093cfe34f Mon Sep 17 00:00:00 2001 From: nuernber Date: Mon, 31 Aug 2026 09:11:51 -0700 Subject: [PATCH 09/15] test(anthropic): fix PT012 lint violation in upstream-error regression test --- .../messages/test_streaming_iterator.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py index f0a3c7fdff0..6a1b6f417a2 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py @@ -501,10 +501,14 @@ async def test_async_sse_wrapper_reraises_upstream_error_to_connected_client(): ) received = [] - with pytest.raises(_ProviderStreamError) as excinfo: + + async def _drain(): async for chunk in iterator.async_sse_wrapper(_failing_stream()): received.append(chunk) + with pytest.raises(_ProviderStreamError) as excinfo: + await _drain() + assert excinfo.value.status_code == 529 assert received assert not any(c.startswith(b"event: error\n") for c in received) From 79fd2f4872ca142cc9a4df8ad9053e93905eb315 Mon Sep 17 00:00:00 2001 From: nuernber Date: Mon, 31 Aug 2026 09:15:36 -0700 Subject: [PATCH 10/15] test(anthropic): cover ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS=0 fallback to partial billing --- litellm/constants.py | 5 ++- .../messages/test_streaming_iterator.py | 42 +++++++++++++++++++ 2 files changed, 46 insertions(+), 1 deletion(-) diff --git a/litellm/constants.py b/litellm/constants.py index 2c46153cff6..b03e122d09c 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -489,7 +489,10 @@ MAX_LANGFUSE_INITIALIZED_CLIENTS: Final = int(os.getenv("MAX_LANGFUSE_INITIALIZE # bounded so a slow client throttles the upstream pump instead of letting it # buffer the whole response in memory; the detached-drain cap bounds how many # post-disconnect drains may run concurrently so client behavior can't create -# unbounded worker state. +# unbounded worker state. Setting the cap to 0 disables detached draining +# entirely: every post-disconnect pump bills whatever partial output it has +# already collected and aborts the upstream stream immediately, instead of +# continuing to drain for the real terminal usage. ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE: Final = int( os.getenv("ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", "1024") ) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py index 6a1b6f417a2..8c3b2852345 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py @@ -635,6 +635,48 @@ async def test_async_sse_wrapper_bills_partial_when_detached_drain_cap_reached(m streaming_iterator_module._DETACHED_STREAM_DRAINS.discard(holder) +@pytest.mark.asyncio +async def test_async_sse_wrapper_bills_partial_when_detached_drains_disabled(monkeypatch): + """ + Regression: ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS=0 must disable + detached draining entirely, not just shrink the cap. With no slots ever + available, the very first post-disconnect chunk must fall back to partial + spend logging instead of hanging on a cap that's unreachable. + """ + monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS", 0) + monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", 4) + + tail_reached = False + + async def _long_stream(): + nonlocal tail_reached + yield {"type": "message_start", "message": {"id": "m", "usage": {"input_tokens": 5, "output_tokens": 1}}} + yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "x"}} + for i in range(100): + yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": f"more{i}"}} + tail_reached = True + 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("drains_disabled"), request_body={}) + gen = iterator.async_sse_wrapper(_long_stream()) + await gen.__anext__() # message_start + await gen.__anext__() # first delta + await gen.aclose() # client disconnects; 100+ chunks remain upstream + + for _ in range(200): + if iterator.logged_chunks: + break + await asyncio.sleep(0) + + assert iterator.logged_chunks, "pump never billed with detached drains disabled" + assert len(iterator.logged_chunks) <= 2 + streaming_iterator_module.ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE + assert len(iterator.logged_chunks) < 100 + assert not any(c.startswith(b"event: message_stop\n") for c in iterator.logged_chunks) + assert tail_reached is False, "pump kept draining despite detached drains being disabled" + assert len(streaming_iterator_module._DETACHED_STREAM_DRAINS) == 0 + + @pytest.mark.asyncio async def test_async_sse_wrapper_aborts_upstream_when_detached_drain_cap_reached(monkeypatch): """ From 8e92f989051db0a518ef244e9961efc7b487a2fe Mon Sep 17 00:00:00 2001 From: nuernber Date: Mon, 31 Aug 2026 09:18:12 -0700 Subject: [PATCH 11/15] docs(constants): clarify which env var the cap=0 fallback note applies to --- litellm/constants.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index b03e122d09c..7a295c37010 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -489,13 +489,14 @@ MAX_LANGFUSE_INITIALIZED_CLIENTS: Final = int(os.getenv("MAX_LANGFUSE_INITIALIZE # bounded so a slow client throttles the upstream pump instead of letting it # buffer the whole response in memory; the detached-drain cap bounds how many # post-disconnect drains may run concurrently so client behavior can't create -# unbounded worker state. Setting the cap to 0 disables detached draining -# entirely: every post-disconnect pump bills whatever partial output it has -# already collected and aborts the upstream stream immediately, instead of -# continuing to drain for the real terminal usage. +# unbounded worker state. ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE: Final = int( os.getenv("ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", "1024") ) +# Setting this to 0 disables detached draining entirely: every post-disconnect +# pump bills whatever partial output it has already collected and aborts the +# upstream stream immediately, instead of continuing to drain for the real +# terminal usage. ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS: Final = int( os.getenv("ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS", "100") ) From 2f50988fed00657dc1c0484828dde37590701930 Mon Sep 17 00:00:00 2001 From: nuernber Date: Mon, 31 Aug 2026 09:30:29 -0700 Subject: [PATCH 12/15] fix(anthropic): resolve TRY300 lint violation and ratchet budgets after litellm_internal_staging merge --- basedpyright-code-budget.json | 10 +++++----- .../messages/streaming_iterator.py | 3 ++- type-discipline-budget.json | 6 +++--- 3 files changed, 10 insertions(+), 9 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index d60c3e9c0af..7d105cdcfe0 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -57,7 +57,7 @@ "limit": 5611 }, "reportMissingTypeArgument": { - "limit": 15350 + "limit": 15349 }, "reportMissingTypeStubs": { "limit": 40 @@ -99,19 +99,19 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 44368 + "limit": 44366 }, "reportUnknownLambdaType": { "limit": 109 }, "reportUnknownMemberType": { - "limit": 38468 + "limit": 38466 }, "reportUnknownParameterType": { - "limit": 19665 + "limit": 19664 }, "reportUnknownVariableType": { - "limit": 30066 + "limit": 30065 }, "reportUnnecessaryCast": { "limit": 111 diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py index 9e115c6624d..d822d09771b 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py @@ -541,9 +541,10 @@ class BaseAnthropicMessagesStreamingIterator: return False try: queue.put_nowait(item) - return True except asyncio.QueueFull: pass + else: + return True put_task: Final = asyncio.ensure_future(queue.put(item)) detached_task: Final = asyncio.ensure_future(client_detached.wait()) try: diff --git a/type-discipline-budget.json b/type-discipline-budget.json index ab34775c460..143cbc91787 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22521 + "limit": 22520 }, "LIT002": { - "limit": 26820 + "limit": 26819 }, "LIT003": { "limit": 269 @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16546 + "limit": 16545 }, "LIT011": { "limit": 5575 From 0e78c5bff7cca880dca8d3c726df4606ffb71a95 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 10:17:27 -0700 Subject: [PATCH 13/15] fix(anthropic_messages): dispatch deferred spend logging when the client disconnects mid-relay When the pump finishes draining while the client is still connected, billing is deferred to the proxy's post-response hook, which only fires on a normally completed response. A client disconnect before the relay consumed the queued tail tore the generator down past that hook, so the request logged no spend at all. The relay teardown now dispatches the stored deferred billing whenever it never reached the end-of-stream sentinel. Also drops the live pass_through_tests script: that CI job runs against a fixed config with no Bedrock model or AWS credentials, so it could only fail there. The scenario is covered by unit tests on the relay/pump seam. --- .../messages/streaming_iterator.py | 24 ++- ..._v1_messages_streaming_disconnect_spend.py | 155 ------------------ .../messages/test_streaming_iterator.py | 52 ++++++ 3 files changed, 75 insertions(+), 156 deletions(-) delete mode 100644 tests/pass_through_tests/test_v1_messages_streaming_disconnect_spend.py diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py index 9e115c6624d..b54dac18c95 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py @@ -480,16 +480,37 @@ class BaseAnthropicMessagesStreamingIterator: _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 try: while True: item = await queue.get() if item is None: + reached_end = True break if isinstance(item, BaseException): raise item yield item finally: client_detached.set() + if not reached_end: + self._dispatch_pending_deferred_logging() + + def _dispatch_pending_deferred_logging(self) -> None: + """Fire deferred billing that a torn-down response would otherwise drop. + + When the pump finishes draining while the client is still connected it + stores the logging coroutine for ProxyLogging._fire_deferred_stream_logging, + which the proxy only fires on a normally completed response: a client + disconnect (GeneratorExit / CancelledError) re-raises past it. Without + this dispatch that window loses the spend row entirely. + """ + deferred_cb: Final = getattr(self.litellm_logging_obj, "_on_deferred_stream_complete", None) + deferred_args: Final = getattr(self.litellm_logging_obj, "_deferred_stream_complete_args", None) + if deferred_cb is None or deferred_args is None: + return + self.litellm_logging_obj._on_deferred_stream_complete = None + self.litellm_logging_obj._deferred_stream_complete_args = None + GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=deferred_cb(*deferred_args)) async def _bill_collected_chunks( self, @@ -541,9 +562,10 @@ class BaseAnthropicMessagesStreamingIterator: return False try: queue.put_nowait(item) - return True except asyncio.QueueFull: pass + else: + return True put_task: Final = asyncio.ensure_future(queue.put(item)) detached_task: Final = asyncio.ensure_future(client_detached.wait()) try: diff --git a/tests/pass_through_tests/test_v1_messages_streaming_disconnect_spend.py b/tests/pass_through_tests/test_v1_messages_streaming_disconnect_spend.py deleted file mode 100644 index e69de720ea4..00000000000 --- a/tests/pass_through_tests/test_v1_messages_streaming_disconnect_spend.py +++ /dev/null @@ -1,155 +0,0 @@ -""" -Regression test: /v1/messages streaming interrupted mid-stream must still -produce a spend-log entry. - -On v1.79.1 the proxy records spend for the partially-streamed request. -A refactor on `main` broke that path, so the same scenario now produces -zero spend-log rows. - -Run against a live proxy (e.g. ``litellm --config proxy_server_config.yaml``): - - pytest tests/pass_through_tests/test_v1_messages_streaming_disconnect_spend.py -s -""" - -import asyncio -import json -import uuid - -import aiohttp -import pytest - - -BASE_URL = "http://127.0.0.1:4000" # change appropriately -ADMIN_KEY = "sk-1234" - - -async def _generate_key(session: aiohttp.ClientSession) -> str: - """Create a fresh virtual key so spend is isolated.""" - url = f"{BASE_URL}/key/generate" - headers = {"Authorization": f"Bearer {ADMIN_KEY}", "Content-Type": "application/json"} - async with session.post(url, headers=headers, json={"models": []}) as resp: - assert resp.status == 200, f"key/generate failed: {await resp.text()}" - data = await resp.json() - return data["key"] - - -async def _get_spend_logs_by_spend_id(session: aiohttp.ClientSession, api_key: str, spend_id: str): - """Query /spend/logs by api_key then filter by spend_id in metadata.""" - url = f"{BASE_URL}/spend/logs?api_key={api_key}" - headers = {"Authorization": f"Bearer {ADMIN_KEY}", "Content-Type": "application/json"} - async with session.get(url, headers=headers) as resp: - assert resp.status == 200, f"spend/logs failed: {await resp.text()}" - all_logs = await resp.json() - if not isinstance(all_logs, list): - return [] - matched = [] - for log in all_logs: - meta = log.get("metadata") - if isinstance(meta, str): - meta = json.loads(meta) - if isinstance(meta, dict): - slm = meta.get("spend_logs_metadata") or {} - if slm.get("spend_id") == spend_id: - matched.append(log) - return matched - - -@pytest.mark.asyncio -@pytest.mark.flaky(retries=3, delay=2) -async def test_v1_messages_streaming_disconnect_has_spend_log(): - """ - 1. Send a streaming POST to /v1/messages. - 2. Read a few SSE chunks, then close the connection (simulating a client - disconnect / interruption). - 3. Wait for the proxy's async spend-tracking pipeline to flush. - 4. Assert that at least one spend-log row exists for the request. - - This PASSES on v1.79.1 and FAILS on the latest main branch. - """ - async with aiohttp.ClientSession( - timeout=aiohttp.ClientTimeout(total=60) - ) as session: - key = await _generate_key(session) - - spend_id = str(uuid.uuid4()) - - headers = { - "Authorization": f"Bearer {key}", - "Content-Type": "application/json", - "x-litellm-spend-logs-metadata": '{"spend_id": "' + spend_id + '"}', - } - - payload = { - "model": "us.anthropic.claude-sonnet-4-5-20250929-v1:0", - "max_tokens": 3000, - "stream": True, - "messages": [ - { - "role": "user", - "content": ( - f"Write several detailed paragraphs (at least 500 words) about the " - f"history of the Roman Empire. Unique id: {uuid.uuid4()}" - ), - } - ], - } - - chunks_read = 0 - - async with session.post( - f"{BASE_URL}/v1/messages", json=payload, headers=headers - ) as resp: - assert resp.status == 200, f"/v1/messages failed: {await resp.text()}" - - async for raw_line in resp.content: - line = raw_line.decode("utf-8", errors="replace").strip() - if not line: - continue - chunks_read += 1 - print(f" chunk #{chunks_read}: {line[:120]}") - if chunks_read >= 5: - break - - assert chunks_read >= 3, ( - f"Expected at least 3 chunks before disconnect, got {chunks_read}" - ) - - print( - f"\nDisconnected after {chunks_read} chunks. " - f"Waiting for spend pipeline to flush …" - ) - - spend_data = None - max_retries = 4 - for attempt in range(1, max_retries + 1): - await asyncio.sleep(10) - print(f" spend-log poll attempt {attempt}/{max_retries}") - spend_data = await _get_spend_logs_by_spend_id(session, key, spend_id) - if spend_data and len(spend_data) > 0: - print(f" ✓ found {len(spend_data)} spend-log row(s)") - break - print(" … not found yet") - - assert spend_data is not None and len(spend_data) > 0, ( - f"No spend-log entry found for spend_id={spend_id} " - f"after streaming disconnect. " - f"This is the regression: interrupted /v1/messages streams must " - f"still record spend." - ) - - log_entry = spend_data[0] - print( - f"\nSpend-log entry:\n{json.dumps(log_entry, indent=2, default=str)}" - ) - - prompt_tokens = log_entry.get("prompt_tokens", 0) - completion_tokens = log_entry.get("completion_tokens", 0) - assert prompt_tokens > 0, ( - "Spend-log row exists but has zero prompt tokens, so usage was not recorded." - ) - assert completion_tokens >= 100, ( - f"Spend-log completion_tokens={completion_tokens} is far below the full " - f"response Bedrock generated and billed. The interrupted stream was billed " - f"on the few chunks the client drained, not the full upstream output. " - f"chunks_read={chunks_read}, prompt_tokens={prompt_tokens}" - ) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py index 8c3b2852345..2135397302e 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py @@ -472,6 +472,58 @@ async def test_async_sse_wrapper_bills_full_stream_when_client_reads_all(): assert not any(c.startswith(b"event: error\n") for c in iterator.logged_chunks) +@pytest.mark.asyncio +async def test_async_sse_wrapper_dispatches_deferred_logging_when_client_disconnects_mid_tail(): + """ + Regression: when the pump finishes draining while the client is still + connected, ``_handle_streaming_logging`` defers billing for the proxy's + post-response hook (``ProxyLogging._fire_deferred_stream_logging``), which + only fires on a normally completed response. If the client then disconnects + before consuming the queued tail, the response generator tears down via + GeneratorExit and that hook never runs. The relay teardown must dispatch + the stored deferred billing itself, or the request logs no spend at all. + """ + dispatched = [] + deferred_fired = asyncio.Event() + + def _deferred_stream_complete(logging_coroutine): + dispatched.append(logging_coroutine) + + async def _consume(): + logging_coroutine.close() + deferred_fired.set() + + return _consume() + + logging_obj = _make_logging_obj("test_deferred_dispatch_on_disconnect_mid_tail") + logging_obj._on_deferred_stream_complete = _deferred_stream_complete + iterator = BaseAnthropicMessagesStreamingIterator(litellm_logging_obj=logging_obj, request_body={}) + + async def _full_stream(): + for event in (*_STREAM_PREFIX, *_STREAM_TAIL): + yield event + + gen = iterator.async_sse_wrapper(_full_stream()) + client_chunks = [] + async for chunk in gen: + client_chunks.append(chunk) + if len(client_chunks) == len(_STREAM_PREFIX): + break + + for _ in range(100): + if getattr(logging_obj, "_deferred_stream_complete_args", None) is not None: + break + await asyncio.sleep(0.01) + assert getattr(logging_obj, "_deferred_stream_complete_args", None) is not None, "pump never deferred billing" + + await gen.aclose() + + assert len(dispatched) == 1, "relay teardown did not dispatch the deferred billing" + assert logging_obj._on_deferred_stream_complete is None + assert logging_obj._deferred_stream_complete_args is None + await asyncio.wait_for(deferred_fired.wait(), timeout=5) + + class _ProviderStreamError(Exception): """Stand-in for a provider-specific streaming failure carrying a status code.""" From 46e090d2f333e9974e0240cac4a4ce8fdb8840e4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 12:47:26 -0700 Subject: [PATCH 14/15] fix(anthropic_messages): bill partial spend when a queued pump error is never consumed When the upstream errors while the client is still connected, the pump forwards the exception through the relay queue so the proxy's failure handling re-raises it. If the client disconnects before consuming that queued exception, neither the failure hook nor billing ran and the spend row was lost. The pump now waits for client detach and, if the exception was never consumed, salvages partial spend like the post-disconnect error path. Also rewrites the bedrock disconnect logging test to the detached-pump contract: billing fires after the upstream drain completes, not synchronously at aclose(). --- .../messages/streaming_iterator.py | 20 ++++++++-- .../messages/test_streaming_iterator.py | 37 +++++++++++++++++++ .../test_anthropic_claude3_transformation.py | 23 ++++++++---- 3 files changed, 70 insertions(+), 10 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py index b54dac18c95..55a64dedee4 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py @@ -157,6 +157,17 @@ def _try_claim_detached_drain_slot() -> bool: return True +def _exception_left_unconsumed(queue: "asyncio.Queue[bytes | None | BaseException]", exc: BaseException) -> bool: + """After client detach the relay never reads the queue again, so drain it here. + + The forwarded exception still sitting in the queue means the relay tore + down before re-raising it, so the proxy's failure handling never ran and + the caller must salvage spend itself. + """ + remaining: Final = tuple(queue.get_nowait() for _ in range(queue.qsize())) + return any(item is exc for item in remaining) + + def _sse_event(event_type: str, payload: Mapping[str, object]) -> bytes: return f"event: {event_type}\ndata: {json.dumps(payload)}\n\n".encode() @@ -637,13 +648,16 @@ class BaseAnthropicMessagesStreamingIterator: Handing the original exception to the client-facing generator lets it re-raise so the proxy's failure handling keeps the provider status and - owns logging (no success-bill). If the client already went away, no - failure hook runs, so bill the partial instead of dropping the request. + owns logging (no success-bill). If the client already went away, or + disconnects before ever consuming the queued exception, no failure hook + runs, so bill the partial instead of dropping the request. """ from litellm._logging import verbose_proxy_logger if not client_detached.is_set() and await self._enqueue_for_client(queue, client_detached, exc): - return + await client_detached.wait() + if not _exception_left_unconsumed(queue, exc): + return verbose_proxy_logger.warning( "async_sse_wrapper upstream pump failed after client disconnect (%d chunks): %s(%s)", len(collected_chunks), diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py index 2135397302e..e8099a4217b 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py @@ -601,6 +601,43 @@ async def test_async_sse_wrapper_salvages_partial_spend_on_upstream_error_after_ assert iterator.logged_chunks == received +@pytest.mark.asyncio +async def test_async_sse_wrapper_salvages_spend_when_queued_error_is_never_consumed(): + """ + When the upstream errors while the client is still connected, the pump + forwards the exception through the queue expecting the relay to re-raise it + into the proxy's failure handling. If the client disconnects before + consuming that queued exception, the handoff never happens and no failure + hook runs, so the pump must notice the unconsumed exception at teardown and + salvage partial spend instead of dropping the row entirely. + """ + upstream_errored = asyncio.Event() + + async def _failing_stream(): + yield {"type": "message_start", "message": {"id": "msg_1", "usage": {"input_tokens": 52, "output_tokens": 1}}} + yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "partial"}} + upstream_errored.set() + raise _ProviderStreamError("mid-stream failure", status_code=500) + + iterator = _RecordingLoggingIterator( + litellm_logging_obj=_make_logging_obj("test_salvage_on_unconsumed_queued_error"), + request_body={}, + ) + + gen = iterator.async_sse_wrapper(_failing_stream()) + received = [await gen.__anext__(), await gen.__anext__()] + await upstream_errored.wait() # exception is now queued behind the consumed chunks + await gen.aclose() # client disconnects without ever consuming the queued exception + + for _ in range(100): + if iterator.logged_chunks: + break + await asyncio.sleep(0.01) + + assert iterator.logging_call_count == 1 + assert iterator.logged_chunks == received + + @pytest.mark.asyncio async def test_async_sse_wrapper_applies_backpressure_to_slow_client(monkeypatch): """ diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index 1e09afd6919..8d07d38b1b6 100644 --- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -3066,17 +3066,21 @@ def test_bedrock_invoke_messages_allows_converted_websearch_function_tool(): async def test_bedrock_sse_wrapper_dispatches_logging_on_client_disconnect(): """ Regression test for LIT-5839: closing the outer bedrock_sse_wrapper - mid-stream (what the proxy does on a client disconnect) must close the - inner async_sse_wrapper deterministically so the partial-stream logging - fires. `completion_start_time` is only stamped on the logging object by - that dispatch, so it observing a value proves the whole chain ran. + mid-stream (what the proxy does on a client disconnect) must not lose the + stream's spend logging. Since the detached-pump relay, the upstream read + survives the disconnect and billing fires once the provider stream ends, + so the dispatch is awaited after releasing the upstream instead of being + observed synchronously at aclose(). `completion_start_time` is only + stamped on the logging object by that dispatch, so it observing a value + proves the whole chain ran. """ cfg = AmazonAnthropicClaudeMessagesConfig() + release_upstream = asyncio.Event() - async def _hanging_stream(): + async def _gated_stream(): yield {"type": "message_start", "message": {"id": "msg_1", "usage": {"input_tokens": 25, "output_tokens": 1}}} yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "partial"}} - await asyncio.Event().wait() + await release_upstream.wait() logging_obj = LiteLLMLoggingObj( model="bedrock/invoke/anthropic.claude-3-sonnet-20240229-v1:0", @@ -3087,11 +3091,16 @@ async def test_bedrock_sse_wrapper_dispatches_logging_on_client_disconnect(): litellm_call_id="test_bedrock_sse_wrapper_disconnect_logging", function_id="test_bedrock_sse_wrapper_disconnect_logging", ) - wrapped = cfg.bedrock_sse_wrapper(_hanging_stream(), litellm_logging_obj=logging_obj, request_body={}) + wrapped = cfg.bedrock_sse_wrapper(_gated_stream(), litellm_logging_obj=logging_obj, request_body={}) await wrapped.__anext__() await wrapped.__anext__() assert logging_obj.completion_start_time is None await wrapped.aclose() + release_upstream.set() + for _ in range(500): + if logging_obj.completion_start_time is not None: + break + await asyncio.sleep(0.01) assert logging_obj.completion_start_time is not None From a99f62d1bc70cdc8c6a6ad0b8b358ed320ad919e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 16:46:12 -0700 Subject: [PATCH 15/15] fix(anthropic_messages): park deferred billing before the end-of-stream sentinel At end of drain the pump enqueued the sentinel first and picked the billing mode from client_detached afterward, so a client that consumed the sentinel and tore the relay down before the pump resumed (possible whenever the sentinel enqueue hit a full queue) had its fully delivered response billed through the teardown path, skipping the proxy's post-response hook. Bill or park before the sentinel goes out, and let an unconsumed sentinel fall back to dispatching the parked billing. --- .../messages/streaming_iterator.py | 23 +++- .../messages/test_streaming_iterator.py | 120 ++++++++++++++++++ 2 files changed, 137 insertions(+), 6 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py index 55a64dedee4..45c7825344b 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py @@ -598,7 +598,12 @@ class BaseAnthropicMessagesStreamingIterator: """Drain the whole upstream into ``queue`` (backpressured) and bill once. Runs detached so a client disconnect can't interrupt the upstream read; - see ``async_sse_wrapper`` for the full rationale. Returns after billing. + see ``async_sse_wrapper`` for the full rationale. On a completed drain + the success billing (or deferred park) happens before the end-of-stream + sentinel is enqueued: the relay can only tear down after consuming the + sentinel, so its teardown can never outrun the park and get mistaken + for a client disconnect, and a sentinel the client never consumes falls + back to dispatching the parked billing here. """ from litellm._logging import verbose_proxy_logger @@ -631,11 +636,17 @@ class BaseAnthropicMessagesStreamingIterator: await self._handle_pump_upstream_error(queue, client_detached, collected_chunks, exc) return - if not client_detached.is_set(): - if not saw_terminal_event: - await self._enqueue_for_client(queue, client_detached, _incomplete_stream_error_sse_event()) - await self._enqueue_for_client(queue, client_detached, None) - await self._bill_collected_chunks(collected_chunks, stream_teardown=client_detached.is_set()) + if client_detached.is_set(): + await self._bill_collected_chunks(collected_chunks, stream_teardown=True) + return + if not saw_terminal_event and not await self._enqueue_for_client( + queue, client_detached, _incomplete_stream_error_sse_event() + ): + await self._bill_collected_chunks(collected_chunks, stream_teardown=True) + return + await self._bill_collected_chunks(collected_chunks, stream_teardown=False) + if not await self._enqueue_for_client(queue, client_detached, None): + self._dispatch_pending_deferred_logging() async def _handle_pump_upstream_error( self, diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py index e8099a4217b..11a048edc1f 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py @@ -1156,3 +1156,123 @@ async def test_normal_end_without_deferred_dispatch_enqueues_immediately(monkeyp assert len(worker.enqueued) == 1 assert getattr(iterator.litellm_logging_obj, "_deferred_stream_complete_args", None) is None worker.close_enqueued() + + +def _backpressured_wrapper(iterator, upstream_exhausted: asyncio.Event): + async def _stream(): + try: + for event in COMPLETE_STREAM_EVENTS: + yield event + finally: + upstream_exhausted.set() + + return iterator.async_sse_wrapper(_stream()) + + +async def _drain_with_pauses_until_upstream_exhausted(gen, upstream_exhausted: asyncio.Event) -> list: + received = [] + while not upstream_exhausted.is_set(): + received.append(await gen.__anext__()) + for _ in range(25): + await asyncio.sleep(0) + assert len(received) <= len(COMPLETE_STREAM_EVENTS) + return received + + +@pytest.mark.asyncio +async def test_normal_end_parks_deferred_logging_even_when_sentinel_enqueue_backpressured(monkeypatch): + """ + Regression: with a full relay queue at end of stream, the pump suspends + while enqueueing the end-of-stream sentinel, and a client that then drains + the whole tail tears the relay down (setting ``client_detached``) before + the pump resumes. That teardown is a normally completed response, not a + disconnect: billing must still park for the proxy's post-response hook + (preserving post_call decoration such as guardrail_information) instead of + enqueueing immediately through the teardown path. + """ + monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", 2) + worker = _RecordingLoggingWorker() + monkeypatch.setattr(streaming_iterator_module, "GLOBAL_LOGGING_WORKER", worker) + + dispatched = [] + + async def _deferred_stream_complete(logging_coroutine): + dispatched.append(logging_coroutine) + logging_coroutine.close() + + iterator = _make_iterator("test_sentinel_backpressure_normal_end") + iterator.litellm_logging_obj._on_deferred_stream_complete = _deferred_stream_complete + + upstream_exhausted = asyncio.Event() + gen = _backpressured_wrapper(iterator, upstream_exhausted) + received = await _drain_with_pauses_until_upstream_exhausted(gen, upstream_exhausted) + + while True: + try: + received.append(await gen.__anext__()) + except StopAsyncIteration: + break + + for _ in range(100): + if worker.enqueued or getattr(iterator.litellm_logging_obj, "_deferred_stream_complete_args", None): + break + await asyncio.sleep(0.01) + + assert len(received) == len(COMPLETE_STREAM_EVENTS) + assert worker.enqueued == [], "fully delivered stream billed through the teardown path" + assert dispatched == [] + parked = getattr(iterator.litellm_logging_obj, "_deferred_stream_complete_args", None) + assert parked is not None, "pump never parked deferred billing" + parked[0].close() + + +@pytest.mark.asyncio +async def test_relay_teardown_dispatches_deferred_billing_when_sentinel_never_consumed(monkeypatch): + """ + Regression: when the pump has parked deferred billing but its end-of-stream + sentinel never fits in the full relay queue (the client disconnects without + draining the tail), the proxy's post-response hook never fires. Exactly one + of the relay teardown or the pump's fallback must dispatch the parked + billing, or the request logs no spend at all. + """ + monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", 2) + worker = _RecordingLoggingWorker() + monkeypatch.setattr(streaming_iterator_module, "GLOBAL_LOGGING_WORKER", worker) + + dispatched = [] + deferred_fired = asyncio.Event() + + def _deferred_stream_complete(logging_coroutine): + dispatched.append(logging_coroutine) + + async def _consume(): + logging_coroutine.close() + deferred_fired.set() + + return _consume() + + iterator = _make_iterator("test_sentinel_never_consumed_dispatch") + iterator.litellm_logging_obj._on_deferred_stream_complete = _deferred_stream_complete + + upstream_exhausted = asyncio.Event() + gen = _backpressured_wrapper(iterator, upstream_exhausted) + await _drain_with_pauses_until_upstream_exhausted(gen, upstream_exhausted) + + for _ in range(100): + if getattr(iterator.litellm_logging_obj, "_deferred_stream_complete_args", None) is not None: + break + await asyncio.sleep(0.01) + + await gen.aclose() + + for _ in range(100): + if dispatched: + break + await asyncio.sleep(0.01) + + assert len(dispatched) == 1, "parked billing was never dispatched" + assert getattr(iterator.litellm_logging_obj, "_deferred_stream_complete_args", None) is None + assert getattr(iterator.litellm_logging_obj, "_on_deferred_stream_complete", None) is None + assert len(worker.enqueued) == 1, "teardown billing enqueued alongside the deferred dispatch" + await worker.enqueued[0] + assert deferred_fired.is_set()