fix(anthropic): log partial stream spend when a /v1/messages client disconnects mid-stream (#37558)

litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py

@@ -10,6 +10,7 @@ from typing_extensions import TypedDict
 
 from litellm.litellm_core_utils.core_helpers import process_response_headers
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
 from litellm.proxy.pass_through_endpoints.success_handler import (
     PassThroughEndpointLogging,
 )
@@ -134,8 +135,11 @@ class BaseAnthropicMessagesStreamingIterator:
         if self.completion_start_time is not None:
             self.litellm_logging_obj.completion_start_time = self.completion_start_time
             self.litellm_logging_obj.model_call_details["completion_start_time"] = self.completion_start_time
-        asyncio.create_task(
-            PassThroughStreamingHandler._route_streaming_logging_to_handler(
+        # Enqueue on the rooted logging worker rather than asyncio.create_task:
+        # this also runs during generator teardown after a client disconnect,
+        # where an unrooted task could be garbage-collected before it bills.
+        GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(
+            async_coroutine=PassThroughStreamingHandler._route_streaming_logging_to_handler(
                 litellm_logging_obj=self.litellm_logging_obj,
                 passthrough_success_handler_obj=GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ,
                 url_route="/v1/messages",
@@ -197,13 +201,21 @@ class BaseAnthropicMessagesStreamingIterator:
         collected_chunks: Final = []
         saw_terminal_event = False
 
-        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
+        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)
+                yield encoded_chunk
+        except (GeneratorExit, asyncio.CancelledError):
+            # A client disconnect tears the generator down at the yield, so the
+            # post-loop logging below never runs and the tokens already streamed
+            # (and billed by the provider) would never reach spend tracking. See LIT-5839.
+            if collected_chunks:
+                await self._handle_streaming_logging(collected_chunks)
+            raise
 
         if not saw_terminal_event:
             yield _incomplete_stream_error_sse_event()

litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py

@@ -842,8 +842,15 @@ class AmazonAnthropicClaudeMessagesConfig(
 
         patched_stream: Final = self._promote_message_stop_usage(completion_stream)
 
-        async for chunk in handler.async_sse_wrapper(patched_stream):
-            yield chunk
+        sse_stream: Final = handler.async_sse_wrapper(patched_stream)
+        try:
+            async for chunk in sse_stream:
+                yield chunk
+        finally:
+            # Close the inner generator deterministically so a client disconnect
+            # (GeneratorExit here) reaches async_sse_wrapper's partial-spend logging
+            # now instead of at garbage collection. See LIT-5839.
+            await sse_stream.aclose()
 
     @staticmethod
     def _merge_message_start_cache_into_delta_usage(

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
@@ -20,9 +21,11 @@ class _RecordingLoggingIterator(BaseAnthropicMessagesStreamingIterator):
     def __init__(self, litellm_logging_obj: LiteLLMLoggingObj, request_body: dict):
         super().__init__(litellm_logging_obj=litellm_logging_obj, request_body=request_body)
         self.logged_chunks: list = []
+        self.logging_call_count: int = 0
 
     async def _handle_streaming_logging(self, collected_chunks):
         self.logged_chunks = list(collected_chunks)
+        self.logging_call_count += 1
 
 
 def _make_logging_obj(test_name: str) -> LiteLLMLoggingObj:
@@ -233,6 +236,70 @@ async def test_async_sse_wrapper_excludes_synthetic_error_event_from_logged_chun
     assert not any(chunk.startswith(b"event: error\n") for chunk in iterator.logged_chunks)
 
 
+async def _events_then_hang(events):
+    for event in events:
+        yield event
+    await asyncio.Event().wait()
+
+
+@pytest.mark.asyncio
+async def test_async_sse_wrapper_logs_partial_chunks_on_client_disconnect():
+    """
+    Regression test for LIT-5839: a client disconnect tears the generator
+    down with GeneratorExit at the yield, which used to skip the post-loop
+    logging dispatch entirely, so the partial output tokens the provider
+    already generated (and billed) never reached spend tracking.
+    """
+    iterator = _RecordingLoggingIterator(
+        litellm_logging_obj=_make_logging_obj("test_disconnect_logs_partial_chunks"),
+        request_body={},
+    )
+    wrapped = iterator.async_sse_wrapper(_events_then_hang(TRUNCATED_TOOL_USE_EVENTS))
+    streamed = [await wrapped.__anext__() for _ in range(len(TRUNCATED_TOOL_USE_EVENTS))]
+    assert iterator.logging_call_count == 0
+
+    await wrapped.aclose()
+
+    assert iterator.logging_call_count == 1
+    assert iterator.logged_chunks == streamed
+
+
+@pytest.mark.asyncio
+async def test_async_sse_wrapper_logs_partial_chunks_on_cancellation():
+    iterator = _RecordingLoggingIterator(
+        litellm_logging_obj=_make_logging_obj("test_cancellation_logs_partial_chunks"),
+        request_body={},
+    )
+    wrapped = iterator.async_sse_wrapper(_events_then_hang(TRUNCATED_TOOL_USE_EVENTS))
+    streamed = [await wrapped.__anext__() for _ in range(len(TRUNCATED_TOOL_USE_EVENTS))]
+
+    consume_task = asyncio.ensure_future(wrapped.__anext__())
+    await asyncio.sleep(0.01)
+    consume_task.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await consume_task
+
+    assert iterator.logging_call_count == 1
+    assert iterator.logged_chunks == streamed
+
+
+@pytest.mark.asyncio
+async def test_async_sse_wrapper_skips_logging_on_disconnect_before_first_chunk():
+    iterator = _RecordingLoggingIterator(
+        litellm_logging_obj=_make_logging_obj("test_disconnect_before_first_chunk"),
+        request_body={},
+    )
+    wrapped = iterator.async_sse_wrapper(_events_then_hang(()))
+
+    consume_task = asyncio.ensure_future(wrapped.__anext__())
+    await asyncio.sleep(0.01)
+    consume_task.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await consume_task
+
+    assert iterator.logging_call_count == 0
+
+
 def test_incomplete_stream_error_sse_event_is_valid_anthropic_error():
     event = _incomplete_stream_error_sse_event().decode()
     lines = event.split("\n")

tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py

@@ -2948,3 +2948,38 @@ def test_bedrock_invoke_messages_allows_converted_websearch_function_tool():
         headers={},
     )
     assert result["tools"][0]["name"] == "litellm_web_search"
+
+
+@pytest.mark.asyncio
+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.
+    """
+    cfg = AmazonAnthropicClaudeMessagesConfig()
+
+    async def _hanging_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()
+
+    logging_obj = LiteLLMLoggingObj(
+        model="bedrock/invoke/anthropic.claude-3-sonnet-20240229-v1:0",
+        messages=[{"role": "user", "content": "hi"}],
+        stream=True,
+        call_type="chat",
+        start_time=datetime.now(),
+        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={})
+    await wrapped.__anext__()
+    await wrapped.__anext__()
+    assert logging_obj.completion_start_time is None
+
+    await wrapped.aclose()
+
+    assert logging_obj.completion_start_time is not None