Python: fix(python): include cache tokens in input_token_count for Bedrock and Claude

python/packages/bedrock/agent_framework_bedrock/_chat_client.py

@@ -868,17 +868,29 @@ def _parse_usage(self, usage: dict[str, Any] | None) -> UsageDetails | None:
         if not usage:
             return None
         details: UsageDetails = {}
-        if (input_tokens := usage.get("inputTokens")) is not None:
-            details["input_token_count"] = input_tokens
-        if (output_tokens := usage.get("outputTokens")) is not None:
-            details["output_token_count"] = output_tokens
-        if (total_tokens := usage.get("totalTokens")) is not None:
-            details["total_token_count"] = total_tokens
+        input_tokens = usage.get("inputTokens")
+        output_tokens = usage.get("outputTokens")
+        cache_read = usage.get("cacheReadInputTokens")
+        cache_write = usage.get("cacheWriteInputTokens")
         # Bedrock Converse reports these when prompt caching is active.
-        if (cache_read := usage.get("cacheReadInputTokens")) is not None:
+        if cache_read is not None:
             details["cache_read_input_token_count"] = cache_read
-        if (cache_write := usage.get("cacheWriteInputTokens")) is not None:
+        if cache_write is not None:
             details["cache_creation_input_token_count"] = cache_write
+        # inputTokens covers only the uncached part when prompt caching is active;
+        # total input (and gen_ai.usage.input_tokens) must include cache tokens,
+        # matching other connectors and the GenAI semconv.
+        input_total = input_tokens
+        if cache_read is not None:
+            input_total = cache_read if input_total is None else input_total + cache_read
+        if cache_write is not None:
+            input_total = cache_write if input_total is None else input_total + cache_write
+        if input_total is not None:
+            details["input_token_count"] = input_total
+        if output_tokens is not None:
+            details["output_token_count"] = output_tokens
+        if (total_tokens := usage.get("totalTokens")) is not None:
+            details["total_token_count"] = total_tokens
         return details or None
 
     def _parse_message_contents(self, content_blocks: Sequence[dict[str, Any]]) -> list[Any]:

python/packages/bedrock/tests/test_bedrock_client.py

@@ -620,7 +620,8 @@ def test_parse_usage_surfaces_cache_tokens() -> None:
     })
 
     assert details is not None
-    assert details["input_token_count"] == 10
+    # inputTokens (10) + cacheReadInputTokens (8) + cacheWriteInputTokens (3) = 21
+    assert details["input_token_count"] == 21
     assert details["cache_read_input_token_count"] == 8
     assert details["cache_creation_input_token_count"] == 3
 

python/packages/claude/agent_framework_claude/_agent.py

@@ -988,19 +988,32 @@ async def _stream_run(
                 usage = message.usage or {}
                 input_tokens = usage.get("input_tokens")
                 output_tokens = usage.get("output_tokens")
+                cache_creation_tokens = usage.get("cache_creation_input_tokens")
+                cache_read_tokens = usage.get("cache_read_input_tokens")
+                # Anthropic's input_tokens excludes cache creation/read; total input
+                # (and gen_ai.usage.input_tokens) must include them, matching other
+                # connectors and the GenAI semconv.
+                input_total = input_tokens
+                for cache_tokens in (cache_creation_tokens, cache_read_tokens):
+                    if isinstance(cache_tokens, int):
+                        input_total = (
+                            cache_tokens
+                            if input_total is None
+                            else input_total + cache_tokens
+                        )
                 total_token_count = (
-                    input_tokens + output_tokens
-                    if isinstance(input_tokens, int) and isinstance(output_tokens, int)
+                    input_total + output_tokens
+                    if isinstance(input_total, int) and isinstance(output_tokens, int)
                     else None
                 )
                 usage_details = UsageDetails(**{
                     key: value
                     for key, value in {
-                        "input_token_count": input_tokens,
+                        "input_token_count": input_total,
                         "output_token_count": output_tokens,
                         "total_token_count": total_token_count,
-                        "cache_creation_input_token_count": usage.get("cache_creation_input_tokens"),
-                        "cache_read_input_token_count": usage.get("cache_read_input_tokens"),
+                        "cache_creation_input_token_count": cache_creation_tokens,
+                        "cache_read_input_token_count": cache_read_tokens,
                     }.items()
                     if isinstance(value, int)
                 })

python/packages/claude/tests/test_claude_agent.py

@@ -353,9 +353,11 @@ async def test_run_captures_result_message_usage_and_finish_reason(
 
         assert response.finish_reason == expected_finish_reason
         assert response.usage_details == {
-            "input_token_count": 42,
+            # input_tokens (42) + cache_creation (3) + cache_read (5) = 50
+            "input_token_count": 50,
             "output_token_count": 18,
-            "total_token_count": 60,
+            # 50 + 18 = 68
+            "total_token_count": 68,
             "cache_creation_input_token_count": 3,
             "cache_read_input_token_count": 5,
         }