fix(anthropic_messages): log spend for interrupted /v1/messages streams

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

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

@@ -1,4 +1,3 @@
-import asyncio
 import json
 from collections.abc import AsyncIterator
 from datetime import datetime
@@ -10,6 +9,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 +134,8 @@ 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(
+        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,16 +197,17 @@ 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
 
-        if not saw_terminal_event:
-            yield _incomplete_stream_error_sse_event()
-
-        # Handle logging after all chunks are processed
-        await self._handle_streaming_logging(collected_chunks)
+            if not saw_terminal_event:
+                yield _incomplete_stream_error_sse_event()
+        finally:
+            if collected_chunks:
+                await self._handle_streaming_logging(collected_chunks)

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

@@ -757,12 +757,12 @@ class AmazonAnthropicClaudeMessagesConfig(
             request_body=request_body,
         )
 
-    async def bedrock_sse_wrapper(
+    def bedrock_sse_wrapper(
         self,
         completion_stream: AsyncIterator[bytes | GenericStreamingChunk | ModelResponseStream | dict],
         litellm_logging_obj: LiteLLMLoggingObj,
         request_body: dict,
-    ):
+    ) -> AsyncIterator[bytes]:
         """
         Bedrock invoke does not return SSE formatted data. This function is a wrapper to ensure litellm chunks are SSE formatted.
 
@@ -786,8 +786,7 @@ class AmazonAnthropicClaudeMessagesConfig(
 
         patched_stream: Final = self._promote_message_stop_usage(completion_stream)
 
-        async for chunk in handler.async_sse_wrapper(patched_stream):
-            yield chunk
+        return handler.async_sse_wrapper(patched_stream)
 
     @staticmethod
     def _merge_message_start_cache_into_delta_usage(

tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py

@@ -233,6 +233,48 @@ 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)
 
 
+@pytest.mark.asyncio
+async def test_async_sse_wrapper_logs_collected_chunks_when_client_disconnects_mid_stream():
+    """
+    Regression test for issue #35958: a client that interrupts a streaming
+    /v1/messages response used to get no spend log at all, because logging
+    only ran after the upstream stream was fully consumed and a disconnect
+    raises GeneratorExit into the wrapper instead.
+    """
+
+    async def _slow_stream():
+        for i in range(50):
+            yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": f"tok{i}"}}
+
+    iterator = _RecordingLoggingIterator(
+        litellm_logging_obj=_make_logging_obj("test_disconnect_logs_partial_chunks"),
+        request_body={},
+    )
+
+    stream = iterator.async_sse_wrapper(_slow_stream())
+    streamed = [await stream.__anext__() for _ in range(3)]
+    await stream.aclose()
+
+    assert iterator.logged_chunks == streamed
+
+
+@pytest.mark.asyncio
+async def test_async_sse_wrapper_does_not_log_when_client_disconnects_before_first_chunk():
+    async def _never_yields():
+        return
+        yield
+
+    iterator = _RecordingLoggingIterator(
+        litellm_logging_obj=_make_logging_obj("test_disconnect_before_first_chunk"),
+        request_body={},
+    )
+
+    stream = iterator.async_sse_wrapper(_never_yields())
+    await stream.aclose()
+
+    assert iterator.logged_chunks == []
+
+
 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

@@ -171,6 +171,53 @@ async def test_bedrock_sse_wrapper_no_error_event_when_stream_ends_with_message_
     assert not any(chunk.startswith(b"event: error\n") for chunk in collected)
 
 
+@pytest.mark.asyncio
+async def test_bedrock_sse_wrapper_dispatches_logging_when_stream_is_closed_mid_stream():
+    """
+    Regression test for issue #35958: closing the wrapper mid-stream (what
+    Starlette does when a /v1/messages client disconnects) must dispatch
+    spend logging for the chunks already streamed. Logging used to run only
+    after the upstream stream was fully consumed, and an extra async-generator
+    layer around the wrapper deferred its teardown to garbage collection, so
+    interrupted Bedrock streams were billed nothing at all.
+
+    ``completion_start_time`` on the injected logging object is written by the
+    wrapper's logging dispatch and by nothing else on this path, so it doubles
+    as the observable signal that the dispatch happened.
+    """
+
+    cfg = AmazonAnthropicClaudeMessagesConfig()
+
+    async def _slow_stream():
+        for i in range(50):
+            yield {
+                "type": "content_block_delta",
+                "index": 0,
+                "delta": {"type": "text_delta", "text": f"tok{i}"},
+            }
+
+    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",
+    )
+    stream = cfg.bedrock_sse_wrapper(
+        _slow_stream(),
+        litellm_logging_obj=logging_obj,
+        request_body={},
+    )
+    await stream.__anext__()
+    assert logging_obj.model_call_details.get("completion_start_time") is None
+
+    await stream.aclose()
+
+    assert logging_obj.model_call_details.get("completion_start_time") is not None
+
+
 @pytest.mark.asyncio
 async def test_bedrock_sse_wrapper_keeps_usage_in_message_start_and_message_delta():
     """Regression test: usage should be available on both message_start and message_delta SSE events."""