fix(passthrough): recover output tokens for interrupted anthropic streams (#30787)

(cherry picked from commit bd74c62ff1)

litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py

@@ -8,6 +8,9 @@ import litellm
 from litellm._logging import verbose_proxy_logger
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 from litellm.litellm_core_utils.litellm_logging import use_custom_pricing_for_model
+from litellm.litellm_core_utils.prompt_templates.common_utils import (
+    get_content_from_model_response,
+)
 from litellm.llms.anthropic import get_anthropic_config
 from litellm.llms.anthropic.chat.handler import (
     ModelResponseIterator as AnthropicModelResponseIterator,
@@ -136,6 +139,84 @@ class AnthropicPassthroughLoggingHandler:
                         return model
         return None
 
+    @staticmethod
+    def _stream_was_interrupted(
+        all_chunks: Sequence[Union[str, bytes]],
+    ) -> bool:
+        """
+        Anthropic ends a stream with ``content_block_stop`` -> ``message_delta``
+        -> ``message_stop``; a client disconnect leaves the last event mid
+        ``content_block_delta``. Scan from the tail and decide on the first
+        terminal-region event, so the common completed case is O(1) rather than
+        re-deserializing every line of the stream.
+        """
+        for raw in reversed(all_chunks):
+            text = raw.decode("utf-8") if isinstance(raw, bytes) else raw
+            for line in reversed(text.splitlines()):
+                if not line.startswith("data:"):
+                    continue
+                try:
+                    data = json.loads(line[len("data:") :].strip())
+                except (json.JSONDecodeError, ValueError):
+                    continue
+                if not isinstance(data, dict):
+                    continue
+                etype = data.get("type")
+                if etype == "message_delta":
+                    return False
+                if etype in (
+                    "content_block_delta",
+                    "content_block_stop",
+                    "message_start",
+                ):
+                    return True
+        return True
+
+    @staticmethod
+    def _recover_interrupted_stream_output_tokens(
+        response: Union[ModelResponse, TextCompletionResponse],
+        all_chunks: Sequence[Union[str, bytes]],
+        model: str,
+    ) -> None:
+        """
+        An Anthropic stream interrupted before its terminal ``message_delta``
+        (client disconnect) carries only the ``message_start`` ``output_tokens``
+        placeholder (typically 1-3), so completion tokens and spend are
+        undercounted ~20x. Re-tokenize the buffered output text to recover a
+        realistic ``output_tokens`` for usage/cost. Completed streams are
+        untouched because their terminal ``message_delta`` short-circuits here.
+        """
+        if not isinstance(response, ModelResponse):
+            return
+        if not AnthropicPassthroughLoggingHandler._stream_was_interrupted(all_chunks):
+            return
+        usage = getattr(response, "usage", None)
+        if usage is None:
+            return
+        output_text = get_content_from_model_response(response)
+        if not output_text:
+            return
+        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
+        if recovered_output_tokens <= (usage.completion_tokens or 0):
+            return
+        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
+        # stale message_start placeholder there must be corrected too or spend
+        # stays undercounted even after completion_tokens is fixed.
+        details = 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
+
     @staticmethod
     def _create_anthropic_response_logging_payload(
         litellm_model_response: Union[ModelResponse, TextCompletionResponse],
@@ -277,6 +358,11 @@ class AnthropicPassthroughLoggingHandler:
                 "result": None,
                 "kwargs": {},
             }
+        AnthropicPassthroughLoggingHandler._recover_interrupted_stream_output_tokens(
+            response=complete_streaming_response,
+            all_chunks=all_chunks,
+            model=model,
+        )
         kwargs = AnthropicPassthroughLoggingHandler._create_anthropic_response_logging_payload(
             litellm_model_response=complete_streaming_response,
             model=model,

tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py

@@ -1015,6 +1015,149 @@ class TestBuildCompleteStreamingResponseRobustness:
         result = self._build(chunks)
         assert result is not None
         assert result.choices[0].message.content == "The stream ends with [DONE]"
+
+
+class TestInterruptedStreamOutputTokenRecovery:
+    """
+    When an Anthropic pass-through stream is interrupted (client disconnect)
+    before the terminal ``message_delta``, the only usage signal is the
+    ``message_start`` ``output_tokens`` placeholder (typically 1-3), so
+    completion tokens and spend are undercounted ~20x. The handler must
+    re-tokenize the buffered ``content_block_delta`` text to recover a
+    realistic ``output_tokens``; completed streams must stay untouched.
+    """
+
+    @staticmethod
+    def _sse(event, data):
+        return f"event: {event}\ndata: {json.dumps(data)}\n\n".encode()
+
+    _MODEL = "claude-3-5-haiku-20241022"
+    _OUTPUT_TEXT = (
+        "The history of computing spans centuries, beginning with mechanical "
+        "calculators and the abacus, advancing through Charles Babbage's "
+        "analytical engine, Ada Lovelace's first algorithm, Alan Turing's "
+        "theoretical machine, and the electronic computers of the twentieth "
+        "century that gave rise to the modern information age."
+    )
+
+    def _interrupted_chunks(self, *, placeholder_output_tokens: int = 2):
+        from litellm.proxy.pass_through_endpoints.streaming_handler import (
+            PassThroughStreamingHandler,
+        )
+
+        words = self._OUTPUT_TEXT.split(" ")
+        frames = [
+            self._sse(
+                "message_start",
+                {
+                    "type": "message_start",
+                    "message": {
+                        "id": "msg_interrupted",
+                        "type": "message",
+                        "role": "assistant",
+                        "model": self._MODEL,
+                        "content": [],
+                        "stop_reason": None,
+                        "stop_sequence": None,
+                        "usage": {
+                            "input_tokens": 29,
+                            "output_tokens": placeholder_output_tokens,
+                        },
+                    },
+                },
+            ),
+            self._sse(
+                "content_block_start",
+                {
+                    "type": "content_block_start",
+                    "index": 0,
+                    "content_block": {"type": "text", "text": ""},
+                },
+            ),
+        ]
+        for i, word in enumerate(words):
+            text = word if i == 0 else " " + word
+            frames.append(
+                self._sse(
+                    "content_block_delta",
+                    {
+                        "type": "content_block_delta",
+                        "index": 0,
+                        "delta": {"type": "text_delta", "text": text},
+                    },
+                )
+            )
+        # Client disconnects here: no content_block_stop / message_delta /
+        # message_stop are ever received.
+        return list(PassThroughStreamingHandler._convert_raw_bytes_to_str_lines(frames))
+
+    def _completed_chunks(self, *, final_output_tokens: int = 80):
+        chunks = self._interrupted_chunks()
+        chunks.append(
+            "data: "
+            + json.dumps(
+                {
+                    "type": "message_delta",
+                    "delta": {"stop_reason": "end_turn", "stop_sequence": None},
+                    "usage": {"output_tokens": final_output_tokens},
+                }
+            )
+        )
+        chunks.append('data: {"type": "message_stop"}')
+        return chunks
+
+    def _run(self, all_chunks):
+        logging_obj = MagicMock()
+        logging_obj.model_call_details = {"model": self._MODEL, "stream": True}
+        logging_obj.litellm_call_id = "test-call-id"
+        logging_obj.litellm_params = {}
+        logging_obj.get_router_model_id.return_value = None
+
+        return 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._MODEL, "stream": True},
+            endpoint_type="messages",
+            start_time=datetime.now(),
+            all_chunks=all_chunks,
+            end_time=datetime.now(),
+        )
+
+    def test_interrupted_stream_retokenizes_buffered_output(self):
+        import litellm
+
+        placeholder = 2
+        result = self._run(
+            self._interrupted_chunks(placeholder_output_tokens=placeholder)
+        )
+        usage = result["result"].usage
+
+        expected = litellm.token_counter(
+            model=self._MODEL,
+            text=self._OUTPUT_TEXT,
+            count_response_tokens=True,
+        )
+
+        assert expected > placeholder * 5
+        assert usage.completion_tokens == expected
+        assert usage.completion_tokens > placeholder
+        assert usage.total_tokens == usage.prompt_tokens + expected
+        # Anthropic spend is priced off completion_tokens_details.text_tokens; if the
+        # placeholder leaks through here, cost stays undercounted even though
+        # completion_tokens looks right.
+        assert usage.completion_tokens_details.text_tokens == expected
+
+    def test_completed_stream_keeps_message_delta_tokens(self):
+        final = 80
+        result = self._run(self._completed_chunks(final_output_tokens=final))
+        usage = result["result"].usage
+
+        # Terminal message_delta present: recovery must not fire; the authoritative
+        # provider count is preserved verbatim.
+        assert usage.completion_tokens == final
+
+
 class TestStreamFalseDeduplication:
     """
     Regression tests for the duplicate-callback bug where a streaming pass-through