From ee3f5a8bdfb52486ca954fe4feaaad590d904108 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 22 Sep 2026 11:50:19 -0700 Subject: [PATCH] fix(streaming): let a later usage event zero out stale cache counts (#40736) (#42330) * fix(streaming): let a later usage event zero out stale cache counts (#40736) The usage merger only replaced cache_creation_input_tokens and cache_read_input_tokens with a positive value, so when Anthropic's message_delta restated the cache block with a 0 write, the 58k write from message_start survived next to the 58k read. prompt_tokens minus both cache counts then went negative and the write was billed twice A usage event that reports any prompt-side count is now authoritative for both cache fields, zeros included. An event with no prompt-side counts, such as an output-only message_delta, still leaves the earlier values alone * test(streaming): cover an input-only message_delta and the derived uncached input in the cache count merger test --------- Co-authored-by: muhammadwaqar12 Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Co-authored-by: shrey kharbanda --- .../streaming_chunk_builder_utils.py | 12 +++- .../test_streaming_chunk_builder_utils.py | 67 +++++++++++++++++++ 2 files changed, 77 insertions(+), 2 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 025db65a7ce..aa4e0cf5495 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/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, @@ -888,11 +896,11 @@ class ChunkProcessor: if usage_chunk_dict["completion_tokens"] is not None and usage_chunk_dict["completion_tokens"] > 0: 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: diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py index 5d75c6699cf..af763da2d87 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -1,5 +1,6 @@ from collections.abc import Mapping, Sequence from typing import Final +from unittest.mock import MagicMock import pytest @@ -7,6 +8,7 @@ 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, @@ -1650,6 +1652,71 @@ 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 + ) + + _ZERO_USAGE_TEXT_CHUNKS: Final = ( _openai_chunk(choices=[{"index": 0, "delta": {"role": "assistant", "content": "Hi"}, "finish_reason": None}]), _openai_chunk(choices=[{"index": 0, "delta": {"content": " there"}, "finish_reason": None}]),