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."""