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.

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()

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}"
+        )

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)

type-discipline-budget.json

@@ -3,7 +3,7 @@
     "limit": 23256
   },
   "LIT002": {
-    "limit": 27213
+    "limit": 27212
   },
   "LIT003": {
     "limit": 269