From eb9beced715cc071914e2248984fb4f43cf482c0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 02:49:09 -0700 Subject: [PATCH] fix(anthropic): price recovered tokens when a /v1/messages client disconnects mid-stream --- .../anthropic_passthrough_logging_handler.py | 50 ++++++++--- .../messages/test_streaming_iterator.py | 84 +++++++++++++++++++ ...t_anthropic_passthrough_logging_handler.py | 82 +++++++++++++++++- 3 files changed, 202 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py index 30b75a7b482..f1387606f4e 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py @@ -216,11 +216,16 @@ class AnthropicPassthroughLoggingHandler: model=model, speed=AnthropicPassthroughLoggingHandler._cost_relevant_speed(request_body), ) - if response is None: - return None - AnthropicPassthroughLoggingHandler._recover_interrupted_stream_output_tokens( + if not isinstance(response, ModelResponse): + return response + recovered_usage: Final = AnthropicPassthroughLoggingHandler._recover_interrupted_stream_output_tokens( response=response, all_chunks=all_chunks, model=model ) + if recovered_usage is None: + return response + AnthropicPassthroughLoggingHandler._reprice_recovered_stream( + response=response, usage=recovered_usage, model=model, logging_obj=litellm_logging_obj + ) return response @staticmethod @@ -259,7 +264,9 @@ class AnthropicPassthroughLoggingHandler: ) except Exception as e: # noqa: BLE001 # an uncostable partial stream still bills its tokens, at zero cost verbose_proxy_logger.warning( - "Anthropic passthrough: could not cost the partial usage of a failed stream (model=%s): %s", model, e + "Anthropic passthrough: could not cost the partial usage of an interrupted stream (model=%s): %s", + model, + e, ) return 0.0 @@ -359,7 +366,7 @@ class AnthropicPassthroughLoggingHandler: response: ModelResponse | TextCompletionResponse, all_chunks: Sequence[str | bytes], model: str, - ) -> None: + ) -> Usage | None: """ An Anthropic stream interrupted before its terminal ``message_delta`` (client disconnect) carries only the ``message_start`` ``output_tokens`` @@ -369,24 +376,24 @@ class AnthropicPassthroughLoggingHandler: untouched because their terminal ``message_delta`` short-circuits here. """ if not isinstance(response, ModelResponse): - return + return None if not AnthropicPassthroughLoggingHandler._stream_was_interrupted(all_chunks): - return + return None usage: Final = getattr(response, "usage", None) - if usage is None: - return + if not isinstance(usage, Usage): + return None output_text: Final = get_content_from_model_response(response) if not output_text: - return + return None try: recovered_output_tokens = litellm.token_counter(model=model, text=output_text, count_response_tokens=True) except Exception: verbose_proxy_logger.warning( "Could not re-tokenize interrupted stream output; keeping placeholder completion token count." ) - return + return None if recovered_output_tokens <= (usage.completion_tokens or 0): - return + return None usage.completion_tokens = recovered_output_tokens usage.total_tokens = (usage.prompt_tokens or 0) + recovered_output_tokens # Anthropic costing reads completion_tokens_details.text_tokens, so the @@ -395,6 +402,25 @@ class AnthropicPassthroughLoggingHandler: details: Final = getattr(usage, "completion_tokens_details", None) if details is not None and getattr(details, "text_tokens", None) is not None: details.text_tokens = recovered_output_tokens + return usage + + @staticmethod + def _reprice_recovered_stream( + response: ModelResponse, + usage: Usage, + model: str, + logging_obj: LiteLLMLoggingObj, + ) -> None: + hidden_params: Final = response._hidden_params # pyright: ignore[reportPrivateUsage] # no public accessor + usage.cost = None + hidden_params.pop("response_cost", None) + recovered_cost: Final = AnthropicPassthroughLoggingHandler._cost_partial_stream_or_zero( + partial_response=response, model=model, logging_obj=logging_obj + ) + if recovered_cost <= 0: + return + usage.cost = recovered_cost + hidden_params["response_cost"] = recovered_cost @staticmethod def _create_anthropic_response_logging_payload( 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 be33b2ee3b1..89193adca1f 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,6 +1,7 @@ import asyncio import json from datetime import datetime +from unittest.mock import patch import pytest @@ -824,6 +825,89 @@ async def test_async_sse_wrapper_bills_partial_when_detached_drains_disabled(mon assert len(streaming_iterator_module._DETACHED_STREAM_DRAINS) == 0 +class _SuccessRecorder(CustomLogger): + def __init__(self): + super().__init__() + self.success_kwargs: list = [] + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.success_kwargs.append(kwargs) + + +@pytest.mark.asyncio +async def test_client_disconnect_partial_billing_prices_recovered_tokens(monkeypatch): + """ + Regression (LIT-6872): a client disconnect that lands on partial billing + re-tokenizes the buffered text into completion_tokens, but the logged cost + stayed priced at the message_start placeholder (1 output token). The success + row's response_cost must match its recovered completion_tokens. + """ + import litellm + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS", 0) + monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", 4) + model = "claude-sonnet-5" + recorder = _SuccessRecorder() + logging_obj = LiteLLMLoggingObj( + model=model, + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="anthropic_messages", + start_time=datetime.now(), + litellm_call_id="disconnect_partial_cost", + function_id="disconnect_partial_cost", + dynamic_async_success_callbacks=[recorder], + ) + logging_obj.update_environment_variables( + model=model, + user="", + optional_params={}, + litellm_params={"custom_llm_provider": "anthropic"}, + custom_llm_provider="anthropic", + ) + iterator = BaseAnthropicMessagesStreamingIterator( + litellm_logging_obj=logging_obj, request_body={"model": model, "stream": True} + ) + sentence = "The history of computing spans centuries of mechanical and electronic invention. " + + async def _stream(): + yield {"type": "message_start", "message": {"id": "msg_1", "usage": {"input_tokens": 29, "output_tokens": 1}}} + yield {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}} + for _ in range(100): + yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": sentence}} + yield {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 1500}} + yield {"type": "message_stop"} + + enqueued: list = [] + + def _capture(async_coroutine): + enqueued.append(async_coroutine) + + with patch.object( # test-quality-ok: GLOBAL_LOGGING_WORKER is a process-global singleton with no injection seam + GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue", side_effect=_capture + ): + gen = iterator.async_sse_wrapper(_stream()) + for _ in range(4): + await gen.__anext__() + await gen.aclose() + for _ in range(500): + if enqueued: + break + await asyncio.sleep(0.01) + + assert len(enqueued) == 1, "client disconnect never reached partial billing" + await enqueued[0] + + assert len(recorder.success_kwargs) == 1 + logged = recorder.success_kwargs[0]["standard_logging_object"] + assert 1 < logged["completion_tokens"] < 1500 + prompt_cost, completion_cost = litellm.cost_per_token( + model=model, prompt_tokens=29, completion_tokens=logged["completion_tokens"] + ) + assert logged["response_cost"] == pytest.approx(prompt_cost + completion_cost) + + @pytest.mark.asyncio async def test_async_sse_wrapper_aborts_upstream_when_detached_drain_cap_reached(monkeypatch): """ diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py index d721be62efe..dc3e94fcff2 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py @@ -7,6 +7,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest +from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import ( AnthropicPassthroughLoggingHandler, @@ -1551,6 +1552,7 @@ class TestInterruptedStreamOutputTokenRecovery: return f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() _MODEL = "claude-3-5-haiku-20241022" + _PRICED_MODEL = "claude-sonnet-5" _OUTPUT_TEXT = ( "The history of computing spans centuries, beginning with mechanical " "calculators and the abacus, advancing through Charles Babbage's " @@ -1559,7 +1561,7 @@ class TestInterruptedStreamOutputTokenRecovery: "century that gave rise to the modern information age." ) - def _interrupted_chunks(self, *, placeholder_output_tokens: int = 2): + def _interrupted_chunks(self, *, placeholder_output_tokens: int = 2, model: str | None = None): from litellm.proxy.pass_through_endpoints.streaming_handler import ( PassThroughStreamingHandler, ) @@ -1574,7 +1576,7 @@ class TestInterruptedStreamOutputTokenRecovery: "id": "msg_interrupted", "type": "message", "role": "assistant", - "model": self._MODEL, + "model": model or self._MODEL, "content": [], "stop_reason": None, "stop_sequence": None, @@ -1676,6 +1678,82 @@ class TestInterruptedStreamOutputTokenRecovery: # provider count is preserved verbatim. assert usage.completion_tokens == final + @pytest.mark.asyncio + async def test_interrupted_stream_logs_cost_of_recovered_tokens(self): + """ + Regression (LIT-6872): stream_chunk_builder stamps usage.cost and + _hidden_params["response_cost"] from the message_start placeholder before + the interrupted stream is re-tokenized, and the success handler prefers + that hidden cost over the recomputed one. The logged cost must price the + recovered completion tokens, not the placeholder. + """ + import litellm + + class _SuccessRecorder(CustomLogger): + def __init__(self): + super().__init__() + self.success_kwargs: list = [] + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.success_kwargs.append(kwargs) + + recorder = _SuccessRecorder() + logging_obj = LiteLLMLoggingObj( + model=self._PRICED_MODEL, + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="pass_through_endpoint", + start_time=datetime.now(), + litellm_call_id="lit-6872", + function_id="lit-6872", + dynamic_async_success_callbacks=[recorder], + ) + logging_obj.update_environment_variables( + model=self._PRICED_MODEL, + user="", + optional_params={}, + litellm_params={"custom_llm_provider": "anthropic"}, + custom_llm_provider="anthropic", + ) + placeholder = 1 + handled = AnthropicPassthroughLoggingHandler._handle_logging_anthropic_collected_chunks( + litellm_logging_obj=logging_obj, + passthrough_success_handler_obj=MagicMock(), + url_route="/anthropic/v1/messages", + request_body={"model": self._PRICED_MODEL, "stream": True}, + endpoint_type="messages", + start_time=datetime.now(), + all_chunks=self._interrupted_chunks(placeholder_output_tokens=placeholder, model=self._PRICED_MODEL), + end_time=datetime.now(), + ) + await logging_obj.dispatch_success_handlers( + result=handled["result"], + start_time=datetime.now(), + end_time=datetime.now(), + cache_hit=False, + prefer_async_handlers=True, + **handled["kwargs"], + ) + for _ in range(300): + if recorder.success_kwargs: + break + await asyncio.sleep(0.01) + + assert len(recorder.success_kwargs) == 1 + logged = recorder.success_kwargs[0]["standard_logging_object"] + recovered_tokens = handled["result"].usage.completion_tokens + assert recovered_tokens > placeholder + assert logged["completion_tokens"] == recovered_tokens + prompt_cost, completion_cost = litellm.cost_per_token( + model=self._PRICED_MODEL, prompt_tokens=29, completion_tokens=recovered_tokens + ) + _, placeholder_completion_cost = litellm.cost_per_token( + model=self._PRICED_MODEL, prompt_tokens=29, completion_tokens=placeholder + ) + assert logged["response_cost"] == pytest.approx(prompt_cost + completion_cost) + assert logged["response_cost"] > prompt_cost + placeholder_completion_cost + assert handled["result"].usage.cost == pytest.approx(prompt_cost + completion_cost) + class TestStreamFalseDeduplication: """