fix(anthropic): price recovered tokens when a /v1/messages client disconnects mid-stream

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(

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):
     """

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:
     """