From fce256f584ae7aed41766c8d20a6ab6f39a63e85 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 5 Aug 2026 18:39:19 +0000 Subject: [PATCH] fix(anthropic_messages): log spend for interrupted /v1/messages streams Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../messages/streaming_iterator.py | 33 ++++++------- .../anthropic_claude3_transformation.py | 7 ++- .../messages/test_streaming_iterator.py | 42 +++++++++++++++++ .../test_anthropic_claude3_transformation.py | 47 +++++++++++++++++++ 4 files changed, 109 insertions(+), 20 deletions(-) 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..78d59b17c95 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py +++ b/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 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 @@ async def _handle_streaming_logging(self, collected_chunks: list[bytes]): 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 @@ async def async_sse_wrapper( 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 - - 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) + 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() + finally: + if collected_chunks: + await self._handle_streaming_logging(collected_chunks) diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index ef9c662bdf5..b645483af1d 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -757,12 +757,12 @@ def get_async_streaming_response_iterator( 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 @@ async def bedrock_sse_wrapper( 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( 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..aaf6bf86518 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 @@ -233,6 +233,48 @@ async def _truncated_stream(): 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") diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index 3b8b4af78d9..92b6eb4b750 100644 --- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -171,6 +171,53 @@ async def _complete_stream(): 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."""