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