Python: include Anthropic cache tokens in input_token_count

python/packages/anthropic/agent_framework_anthropic/_chat_client.py

@@ -673,7 +673,10 @@ async def _stream_turn(
                 # Anthropic streams cumulative usage snapshots (message_start seeds it,
                 # each message_delta carries the running total), so thread a per-stream
                 # accumulator to _process_stream_event to emit increments instead.
+                # raw_usage_carry keeps last-seen input/cache components when a later
+                # snapshot omits them (cache_* often only appear on message_start).
                 emitted_usage: dict[str, int] = {}
+                raw_usage_carry: dict[str, int] = {}
                 local_function_call_started = False
                 open_local_function_call_blocks: set[int] = set()
                 message_stop_received = False
@@ -683,6 +686,7 @@ async def _stream_turn(
                             chunk,
                             emitted_usage,
                             request_state=turn_request_state,
+                            raw_usage_carry=raw_usage_carry,
                         )
                         if chunk.type == "content_block_start" and parsed_chunk is not None:
                             if any(
@@ -1487,6 +1491,7 @@ def _process_stream_event(
         event: BetaRawMessageStreamEvent,
         emitted_usage: dict[str, int] | None = None,
         request_state: _AnthropicRequestState | None = None,
+        raw_usage_carry: dict[str, int] | None = None,
     ) -> ChatResponseUpdate | None:
         """Process a streaming event from the Anthropic client.
 
@@ -1497,14 +1502,22 @@ def _process_stream_event(
                 increments (see ``_incremental_usage``). Pass ``None`` for a one-off
                 event to keep the snapshot unchanged.
             request_state: State used to parse this stream.
+            raw_usage_carry: Per-stream last-seen Anthropic input/cache token
+                components. Streaming deltas often omit ``cache_*`` fields that were
+                only present on ``message_start``; carrying them forward keeps
+                ``input_token_count`` inclusive of cache tokens across the stream.
 
         Returns:
             A ChatResponseUpdate object containing the processed update.
         """
         match event.type:
             case "message_start":
                 usage_details: list[Content] = []
-                if event.message.usage and (details := self._parse_usage_from_anthropic(event.message.usage)):
+                if event.message.usage and (
+                    details := self._parse_usage_from_anthropic(
+                        event.message.usage, raw_usage_carry=raw_usage_carry
+                    )
+                ):
                     usage_details.append(
                         Content.from_usage(
                             usage_details=self._incremental_usage(details, emitted_usage),
@@ -1529,7 +1542,7 @@ def _process_stream_event(
                     raw_representation=event,
                 )
             case "message_delta":
-                usage = self._parse_usage_from_anthropic(event.usage)
+                usage = self._parse_usage_from_anthropic(event.usage, raw_usage_carry=raw_usage_carry)
                 return ChatResponseUpdate(
                     contents=[
                         Content.from_usage(
@@ -1571,19 +1584,60 @@ def _process_stream_event(
                 logger.debug(f"Ignoring unsupported event type: {event.type}")
         return None
 
-    def _parse_usage_from_anthropic(self, usage: BetaUsage | BetaMessageDeltaUsage | None) -> UsageDetails | None:
-        """Parse usage details from the Anthropic message usage."""
+    def _parse_usage_from_anthropic(
+        self,
+        usage: BetaUsage | BetaMessageDeltaUsage | None,
+        *,
+        raw_usage_carry: dict[str, int] | None = None,
+    ) -> UsageDetails | None:
+        """Parse usage details from the Anthropic message usage.
+
+        When ``raw_usage_carry`` is provided (streaming), last-seen ``input_tokens`` /
+        ``cache_*`` components are reused if the current snapshot omits them. Anthropic
+        often reports cache tokens only on ``message_start``; without the carry,
+        a later ``message_delta`` with a higher ``input_tokens`` and omitted cache
+        fields would under-count ``input_token_count`` after incremental conversion.
+        """
         if not usage:
             return None
         usage_details = UsageDetails(output_token_count=usage.output_tokens)
-        if usage.input_tokens is not None:
-            usage_details["input_token_count"] = usage.input_tokens
-        if usage.cache_creation_input_tokens is not None:
-            usage_details["anthropic.cache_creation_input_tokens"] = usage.cache_creation_input_tokens
-            usage_details["cache_creation_input_token_count"] = usage.cache_creation_input_tokens
-        if usage.cache_read_input_tokens is not None:
-            usage_details["anthropic.cache_read_input_tokens"] = usage.cache_read_input_tokens
-            usage_details["cache_read_input_token_count"] = usage.cache_read_input_tokens
+        # Anthropic's input_tokens excludes cache_creation and cache_read; total
+        # input (and gen_ai.usage.input_tokens) must include them, matching other
+        # connectors and the GenAI semconv.
+        input_tokens = usage.input_tokens
+        cache_creation = usage.cache_creation_input_tokens
+        cache_read = usage.cache_read_input_tokens
+        if raw_usage_carry is not None:
+            if input_tokens is not None:
+                raw_usage_carry["input_tokens"] = input_tokens
+            else:
+                input_tokens = raw_usage_carry.get("input_tokens")
+            if cache_creation is not None:
+                raw_usage_carry["cache_creation_input_tokens"] = cache_creation
+            else:
+                cache_creation = raw_usage_carry.get("cache_creation_input_tokens")
+            if cache_read is not None:
+                raw_usage_carry["cache_read_input_tokens"] = cache_read
+            else:
+                cache_read = raw_usage_carry.get("cache_read_input_tokens")
+
+        input_total = input_tokens if input_tokens is not None else None
+        if cache_creation is not None:
+            usage_details["anthropic.cache_creation_input_tokens"] = cache_creation
+            usage_details["cache_creation_input_token_count"] = cache_creation
+            if input_total is not None:
+                input_total += cache_creation
+            else:
+                input_total = cache_creation
+        if cache_read is not None:
+            usage_details["anthropic.cache_read_input_tokens"] = cache_read
+            usage_details["cache_read_input_token_count"] = cache_read
+            if input_total is not None:
+                input_total += cache_read
+            else:
+                input_total = cache_read
+        if input_total is not None:
+            usage_details["input_token_count"] = input_total
         return usage_details
 
     @staticmethod

python/packages/anthropic/tests/test_anthropic_client.py

@@ -3203,15 +3203,26 @@ def test_process_stream_event_message_start_sets_assistant_role(mock_anthropic_c
     assert result.role == "assistant"
 
 
-def _usage_message_start_event(*, input_tokens: int, output_tokens: int) -> MagicMock:
+def _usage_message_start_event(
+    *,
+    input_tokens: int,
+    output_tokens: int,
+    cache_creation_input_tokens: int | None = None,
+    cache_read_input_tokens: int | None = None,
+) -> MagicMock:
     event = MagicMock()
     event.type = "message_start"
     event.message.id = "msg_usage"
     event.message.role = "assistant"
     event.message.model = "claude-3-5-sonnet-20241022"
     event.message.content = []
     event.message.stop_reason = None
-    event.message.usage = BetaUsage(input_tokens=input_tokens, output_tokens=output_tokens)
+    event.message.usage = BetaUsage(
+        input_tokens=input_tokens,
+        output_tokens=output_tokens,
+        cache_creation_input_tokens=cache_creation_input_tokens,
+        cache_read_input_tokens=cache_read_input_tokens,
+    )
     return event
 
 
@@ -3282,6 +3293,48 @@ def test_streaming_usage_delta_input_not_double_counted(mock_anthropic_client: M
     assert response.usage_details["input_token_count"] == 12
 
 
+def test_streaming_usage_preserves_cache_tokens_across_delta(mock_anthropic_client: MagicMock) -> None:
+    """message_delta omitting cache_* must keep cache tokens in input_token_count.
+
+    Anthropic often reports cache_creation/cache_read only on message_start. A later
+    message_delta with a higher input_tokens and omitted cache fields must still include
+    the last-seen cache components so incremental conversion reconstructs
+    input + cache_creation + cache_read (e.g. 120+20+30=170), not bare input (120).
+    """
+    client = create_test_anthropic_client(mock_anthropic_client)
+    emitted: dict[str, int] = {}
+    raw_carry: dict[str, int] = {}
+    updates = [
+        u
+        for u in (
+            client._process_stream_event(
+                _usage_message_start_event(
+                    input_tokens=100,
+                    output_tokens=1,
+                    cache_creation_input_tokens=20,
+                    cache_read_input_tokens=30,
+                ),
+                emitted,
+                raw_usage_carry=raw_carry,
+            ),
+            client._process_stream_event(
+                _usage_message_delta_event(input_tokens=120, output_tokens=25),
+                emitted,
+                raw_usage_carry=raw_carry,
+            ),
+        )
+        if u is not None
+    ]
+
+    response = ChatResponse.from_updates(updates)
+
+    assert response.usage_details is not None
+    assert response.usage_details["output_token_count"] == 25
+    assert response.usage_details["input_token_count"] == 170
+    assert response.usage_details["cache_creation_input_token_count"] == 20
+    assert response.usage_details["cache_read_input_token_count"] == 30
+
+
 def test_process_stream_event_message_start_role_prevents_tool_use_collapse() -> None:
     """Regression test: tool_use blocks must not end up in a user-role message.
 
@@ -4886,7 +4939,8 @@ def test_parse_usage_with_cache_tokens(mock_anthropic_client: MagicMock) -> None
     assert result is not None
     result_dict = cast("dict[str, Any]", result)
     assert result["output_token_count"] == 50
-    assert result["input_token_count"] == 100
+    # 100 input + 20 cache_creation + 30 cache_read
+    assert result["input_token_count"] == 150
     assert result_dict["anthropic.cache_creation_input_tokens"] == 20
     assert result_dict["anthropic.cache_read_input_tokens"] == 30
     assert result["cache_creation_input_token_count"] == 20