From 2dabac186abc7329222985b295545b38f60dfcde Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 03:53:40 -0700 Subject: [PATCH] fix(anthropic): clear the placeholder cost once and let logging price recovered tokens Drop the second pricing pass on interrupted /v1/messages streams: clearing the stale usage.cost and hidden response_cost is enough for the existing success and failure logging to price the recovered usage. Add an iterator test for the upstream-close path the proxy takes on a client disconnect. --- .../anthropic_passthrough_logging_handler.py | 21 +-- .../messages/test_streaming_iterator.py | 121 +++++++++++++++--- ...t_anthropic_passthrough_logging_handler.py | 1 - 3 files changed, 106 insertions(+), 37 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 f1387606f4e..7cec3bac207 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 @@ -223,9 +223,7 @@ class AnthropicPassthroughLoggingHandler: ) if recovered_usage is None: return response - AnthropicPassthroughLoggingHandler._reprice_recovered_stream( - response=response, usage=recovered_usage, model=model, logging_obj=litellm_logging_obj - ) + AnthropicPassthroughLoggingHandler._clear_placeholder_cost(response=response, usage=recovered_usage) return response @staticmethod @@ -405,22 +403,9 @@ class AnthropicPassthroughLoggingHandler: 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 + def _clear_placeholder_cost(response: ModelResponse, usage: Usage) -> None: 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 + response._hidden_params.pop("response_cost", None) # pyright: ignore[reportPrivateUsage] # no public accessor @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 89193adca1f..8043496f299 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 @@ -834,6 +834,50 @@ class _SuccessRecorder(CustomLogger): self.success_kwargs.append(kwargs) +def _make_priced_logging_obj(call_id: str, recorder: _SuccessRecorder, model: str) -> LiteLLMLoggingObj: + logging_obj = LiteLLMLoggingObj( + model=model, + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="anthropic_messages", + start_time=datetime.now(), + litellm_call_id=call_id, + function_id=call_id, + 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", + ) + return logging_obj + + +class _UpstreamClosedOnDetach: + """Upstream that yields its events and then, like a socket read, waits until it is closed.""" + + def __init__(self, events: tuple[dict, ...]): + self._events = iter(events) + self._closed = asyncio.Event() + + def __aiter__(self): + return self + + async def __anext__(self) -> dict: + if self._closed.is_set(): + raise StopAsyncIteration + try: + return next(self._events) + except StopIteration: + await self._closed.wait() + raise StopAsyncIteration + + async def aclose(self) -> None: + self._closed.set() + + @pytest.mark.asyncio async def test_client_disconnect_partial_billing_prices_recovered_tokens(monkeypatch): """ @@ -849,25 +893,9 @@ async def test_client_disconnect_partial_billing_prices_recovered_tokens(monkeyp 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} + litellm_logging_obj=_make_priced_logging_obj("disconnect_partial_cost", recorder, model), + request_body={"model": model, "stream": True}, ) sentence = "The history of computing spans centuries of mechanical and electronic invention. " @@ -908,6 +936,63 @@ async def test_client_disconnect_partial_billing_prices_recovered_tokens(monkeyp assert logged["response_cost"] == pytest.approx(prompt_cost + completion_cost) +@pytest.mark.asyncio +async def test_proxy_disconnect_closing_upstream_prices_recovered_tokens(): + """ + Regression (LIT-6872), proxy path: after a client disconnect the proxy's + shielded cleanup closes the upstream stream while the pump is still reading + it, so the pump bills the chunks collected so far without ever seeing + message_delta. That row's response_cost must be priced from its recovered + completion_tokens, not from the message_start placeholder. + """ + import litellm + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + model = "claude-sonnet-5" + recorder = _SuccessRecorder() + iterator = BaseAnthropicMessagesStreamingIterator( + litellm_logging_obj=_make_priced_logging_obj("disconnect_upstream_closed", recorder, model), + request_body={"model": model, "stream": True}, + ) + sentence = "The history of computing spans centuries of mechanical and electronic invention. " + upstream = _UpstreamClosedOnDetach( + ( + {"type": "message_start", "message": {"id": "msg_1", "usage": {"input_tokens": 29, "output_tokens": 1}}}, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + *({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": sentence}} for _ in range(6)), + ) + ) + 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(upstream) + for _ in range(4): + await gen.__anext__() + await gen.aclose() + assert not enqueued, "billing must wait for the upstream read to end, not the client detach" + await upstream.aclose() + for _ in range(500): + if enqueued: + break + await asyncio.sleep(0.01) + + assert len(enqueued) == 1, "closing the upstream never reached partial billing" + await enqueued[0] + + assert len(recorder.success_kwargs) == 1 + logged = recorder.success_kwargs[0]["standard_logging_object"] + assert logged["completion_tokens"] > 1 + 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 dc3e94fcff2..2d7397594aa 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 @@ -1752,7 +1752,6 @@ class TestInterruptedStreamOutputTokenRecovery: ) 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: