From 48feaeec1b0445aa4e575ce9e6c524acb7b2e975 Mon Sep 17 00:00:00 2001 From: Alex Castillo Date: Wed, 21 Oct 2026 14:33:29 -0400 Subject: [PATCH] fix(ai): take cumulative usage and web searches from the Claude Agent SDK message_delta The tracker only read output_tokens from a message_delta. Its usage is cumulative, and with server tools (web search, ...) the input and cache counts exceed the message_start values, so streamed generations under-reported them, and the number of web searches was never recorded. Keep the larger input and cache counts, as the Anthropic wrapper does, and set $ai_web_search_count from server_tool_use.web_search_requests. Co-Authored-By: Claude Sonnet 5.5 --- .../claude-agent-sdk-message-delta-usage.md | 5 ++ posthog/ai/claude_agent_sdk/processor.py | 21 +++++- .../ai/claude_agent_sdk/test_processor.py | 75 +++++++++++++++++++ 3 files changed, 100 insertions(+), 1 deletion(-) create mode 100644 .sampo/changesets/claude-agent-sdk-message-delta-usage.md diff --git a/.sampo/changesets/claude-agent-sdk-message-delta-usage.md b/.sampo/changesets/claude-agent-sdk-message-delta-usage.md new file mode 100644 index 000000000..3db867892 --- /dev/null +++ b/.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. diff --git a/posthog/ai/claude_agent_sdk/processor.py b/posthog/ai/claude_agent_sdk/processor.py index 1d45b61c2..d5405ec66 100644 --- a/posthog/ai/claude_agent_sdk/processor.py +++ b/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 diff --git a/posthog/test/ai/claude_agent_sdk/test_processor.py b/posthog/test/ai/claude_agent_sdk/test_processor.py index 3dc776002..548fb8244 100644 --- a/posthog/test/ai/claude_agent_sdk/test_processor.py +++ b/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