From 8922aaab9578d2ed22d24d4cebf4601f1469cbd0 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 19 Aug 2026 18:43:46 -0700 Subject: [PATCH] fix(anthropic): log partial stream spend when a /v1/messages client disconnects mid-stream (#37558) --- .../messages/streaming_iterator.py | 30 ++++++--- .../anthropic_claude3_transformation.py | 11 ++- .../messages/test_streaming_iterator.py | 67 +++++++++++++++++++ .../test_anthropic_claude3_transformation.py | 35 ++++++++++ 4 files changed, 132 insertions(+), 11 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..922769dbbfd 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py +++ b/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() 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 0161f4fadc9..f74a290d773 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/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( 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..5c1cd88835f 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 @@ -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") 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 ce81edf3101..1a3fdb9f652 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 @@ -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