fix(ai): take cumulative usage and web searches from the Claude Agent SDK message_delta

.sampo/changesets/claude-agent-sdk-message-delta-usage.md

@@ -0,0 +1,5 @@
+---
+pypi/posthog: patch
+---
+
+The Claude Agent SDK integration now takes the cumulative input and cache token counts from the final `message_delta` usage, which grow past the `message_start` values when server tools such as web search run, and sets `$ai_web_search_count` from `server_tool_use.web_search_requests`. Streamed generations used to under-report input tokens and cache writes and never recorded the web search count.

posthog/ai/claude_agent_sdk/processor.py

@@ -51,6 +51,7 @@ class _GenerationData:
     output_tokens: Optional[int] = None
     cache_read_input_tokens: Optional[int] = None
     cache_creation_input_tokens: Optional[int] = None
+    web_search_count: Optional[int] = None
     raw_usage: Optional[Dict[str, Any]] = None
     start_time: float = 0.0
     end_time: float = 0.0
@@ -92,9 +93,25 @@ def process_stream_event(self, event: "StreamEvent") -> None:
         elif event_type == "message_delta" and self._current is not None:
             usage = raw.get("usage", {})
             self._current.raw_usage = {**(self._current.raw_usage or {}), **usage}
-            # message_delta usage reports cumulative output tokens
+            # message_delta usage is cumulative: output tokens, and with server
+            # tools (web search, ...) the input and cache counts grow past the
+            # message_start values. Keep the larger count; a field the delta
+            # omits keeps the message_start value.
             if usage.get("output_tokens"):
                 self._current.output_tokens = usage["output_tokens"]
+            for field_name in (
+                "input_tokens",
+                "cache_read_input_tokens",
+                "cache_creation_input_tokens",
+            ):
+                delta_value = usage.get(field_name)
+                current_value = getattr(self._current, field_name)
+                if isinstance(delta_value, int) and delta_value > (current_value or 0):
+                    setattr(self._current, field_name, delta_value)
+            server_tool_use = usage.get("server_tool_use") or {}
+            web_search_requests = server_tool_use.get("web_search_requests")
+            if isinstance(web_search_requests, int) and web_search_requests > 0:
+                self._current.web_search_count = web_search_requests
             # Extract stop reason from message_delta
             delta_stop_reason = raw.get("delta", {}).get("stop_reason")
             if delta_stop_reason is not None:
@@ -443,6 +460,8 @@ def _emit_generation(
             properties["$ai_cache_creation_input_tokens"] = (
                 gen.cache_creation_input_tokens
             )
+        if gen.web_search_count:
+            properties["$ai_web_search_count"] = gen.web_search_count
         if gen.raw_usage:
             properties["$ai_usage"] = gen.raw_usage
 

posthog/test/ai/claude_agent_sdk/test_processor.py

@@ -221,6 +221,81 @@ async def test_emits_generation_from_stream_events(self, processor, mock_client)
         assert props["$ai_output_tokens"] == 50
         assert props["$ai_cache_read_input_tokens"] == 20
 
+    @pytest.mark.asyncio
+    async def test_message_delta_usage_is_cumulative_with_server_tools(
+        self, processor, mock_client
+    ):
+        # With server tools the final message_delta carries the cumulative input
+        # and cache counts and the number of web searches.
+        message_delta = _make_stream_event(
+            "message_delta",
+            {
+                "delta": {"stop_reason": "end_turn"},
+                "usage": {
+                    "input_tokens": 5120,
+                    "output_tokens": 300,
+                    "cache_read_input_tokens": 64,
+                    "cache_creation_input_tokens": 32,
+                    "server_tool_use": {"web_search_requests": 2},
+                },
+            },
+        )
+        messages = [
+            _make_message_start(input_tokens=120, cache_read=0, cache_creation=0),
+            message_delta,
+            _make_message_stop(),
+            _make_assistant_message(),
+            _make_result_message(),
+        ]
+
+        with patch(
+            "posthog.ai.claude_agent_sdk.processor.original_query",
+            side_effect=lambda **kw: _fake_query(messages),
+        ):
+            async for _ in processor.query(prompt="Hi", options=ClaudeAgentOptions()):
+                pass
+
+        gen_call = next(
+            c
+            for c in mock_client.capture.call_args_list
+            if (c.kwargs.get("event") or c[1].get("event")) == "$ai_generation"
+        )
+        props = gen_call.kwargs.get("properties") or gen_call[1].get("properties")
+        assert props["$ai_input_tokens"] == 5120
+        assert props["$ai_output_tokens"] == 300
+        assert props["$ai_cache_read_input_tokens"] == 64
+        assert props["$ai_cache_creation_input_tokens"] == 32
+        assert props["$ai_web_search_count"] == 2
+
+    @pytest.mark.asyncio
+    async def test_message_delta_without_input_counts_keeps_message_start(
+        self, processor, mock_client
+    ):
+        messages = [
+            _make_message_start(input_tokens=100, cache_read=20),
+            _make_message_delta(output_tokens=50),
+            _make_message_stop(),
+            _make_assistant_message(),
+            _make_result_message(),
+        ]
+
+        with patch(
+            "posthog.ai.claude_agent_sdk.processor.original_query",
+            side_effect=lambda **kw: _fake_query(messages),
+        ):
+            async for _ in processor.query(prompt="Hi", options=ClaudeAgentOptions()):
+                pass
+
+        gen_call = next(
+            c
+            for c in mock_client.capture.call_args_list
+            if (c.kwargs.get("event") or c[1].get("event")) == "$ai_generation"
+        )
+        props = gen_call.kwargs.get("properties") or gen_call[1].get("properties")
+        assert props["$ai_input_tokens"] == 100
+        assert props["$ai_cache_read_input_tokens"] == 20
+        assert "$ai_web_search_count" not in props
+
     @pytest.mark.asyncio
     async def test_emits_cache_creation_ttl_breakdown_from_stream_events(
         self, processor, mock_client