fix(streaming): let a later usage event zero out stale cache counts (#40736)

litellm/litellm_core_utils/streaming_chunk_builder_utils.py

@@ -165,6 +165,14 @@ class _UsageSummary(TypedDict):
     cost: float | None
 
 
+def _reports_prompt_side_usage(usage_summary: "_UsageSummary") -> bool:
+    return (
+        (usage_summary["prompt_tokens"] or 0) > 0
+        or (usage_summary["cache_creation_input_tokens"] or 0) > 0
+        or (usage_summary["cache_read_input_tokens"] or 0) > 0
+    )
+
+
 def capture_cache_creation_token_details(
     prompt_tokens_details: PromptTokensDetailsWrapper | None,
     current: CacheCreationTokenDetails | None,
@@ -881,11 +889,11 @@ def _calculate_usage_per_chunk(
                     completion_tokens = usage_chunk_dict["completion_tokens"]
                     completion_usage_updates += 1
                 if usage_chunk_dict["cache_creation_input_tokens"] is not None and (
-                    usage_chunk_dict["cache_creation_input_tokens"] > 0 or cache_creation_input_tokens is None
+                    _reports_prompt_side_usage(usage_chunk_dict) or cache_creation_input_tokens is None
                 ):
                     cache_creation_input_tokens = usage_chunk_dict["cache_creation_input_tokens"]
                 if usage_chunk_dict["cache_read_input_tokens"] is not None and (
-                    usage_chunk_dict["cache_read_input_tokens"] > 0 or cache_read_input_tokens is None
+                    _reports_prompt_side_usage(usage_chunk_dict) or cache_read_input_tokens is None
                 ):
                     cache_read_input_tokens = usage_chunk_dict["cache_read_input_tokens"]
                 if usage_chunk_dict["completion_tokens_details"] is not None:

tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py

@@ -1,12 +1,14 @@
 from collections.abc import Mapping, Sequence
 from typing import Final
+from unittest.mock import MagicMock
 
 import pytest
 
 
 from litellm import ChatCompletionUsageBlock, stream_chunk_builder
 from litellm.types.utils import GenericStreamingChunk
 from litellm.litellm_core_utils.streaming_chunk_builder_utils import ChunkProcessor
+from litellm.llms.anthropic.chat.handler import ModelResponseIterator
 from litellm.types.utils import (
     ChatCompletionDeltaToolCall,
     ChatCompletionMessageToolCall,
@@ -1648,3 +1650,68 @@ def test_calculate_usage_falls_back_to_prompt_counter_when_mock_stream_has_no_ad
     )
 
     assert usage.prompt_tokens == 77
+
+
+@pytest.mark.parametrize(
+    ("message_delta_usage", "expected_cache_creation", "expected_cache_read"),
+    [
+        (
+            {
+                "input_tokens": 2,
+                "cache_creation_input_tokens": 0,
+                "cache_read_input_tokens": 58352,
+                "output_tokens": 408,
+            },
+            0,
+            58352,
+        ),
+        ({"output_tokens": 408}, 58352, 0),
+        ({"input_tokens": 2, "output_tokens": 408}, 58352, 0),
+    ],
+    ids=["delta_restates_cache_counts", "delta_reports_output_only", "delta_reports_input_and_output_only"],
+)
+def test_anthropic_stream_usage_takes_cache_counts_from_last_event_that_reports_them(
+    message_delta_usage: Mapping[str, int], expected_cache_creation: int, expected_cache_read: int
+) -> None:
+    iterator: Final = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False)
+    events: Final = (
+        {
+            "type": "message_start",
+            "message": {
+                "id": "msg_1",
+                "type": "message",
+                "role": "assistant",
+                "model": "claude-sonnet-4-5",
+                "content": [],
+                "stop_reason": None,
+                "usage": {
+                    "input_tokens": 2,
+                    "cache_creation_input_tokens": 58352,
+                    "cache_read_input_tokens": 0,
+                    "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": "hi"}},
+        {"type": "content_block_stop", "index": 0},
+        {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": dict(message_delta_usage)},
+        {"type": "message_stop"},
+    )
+
+    response: Final = stream_chunk_builder(
+        chunks=[iterator.chunk_parser(event) for event in events],
+        messages=[{"role": "user", "content": "hi"}],
+    )
+
+    assert response.usage.cache_creation_input_tokens == expected_cache_creation
+    assert response.usage.cache_read_input_tokens == expected_cache_read
+    assert response.usage.prompt_tokens == 58354
+    assert response.usage.prompt_tokens_details.cache_creation_tokens == expected_cache_creation
+    assert response.usage.prompt_tokens_details.cached_tokens == expected_cache_read
+    assert (
+        response.usage.prompt_tokens
+        - response.usage.cache_read_input_tokens
+        - response.usage.cache_creation_input_tokens
+        == 2
+    )