fix(ai): keep Anthropic stream tool arguments after server tool blocks

.sampo/changesets/anthropic-stream-block-index.md

@@ -0,0 +1,5 @@
+---
+pypi/posthog: patch
+---
+
+Keep streamed Anthropic tool call arguments when the response also contains server tool blocks, such as web search.

posthog/ai/anthropic/_anthropic_stream.py

@@ -21,6 +21,10 @@ def __init__(self) -> None:
         self.tools_in_progress: Dict[str, ToolInProgress] = {}
         self.current_text_block: Optional[StreamingContentBlock] = None
         self.stop_reason: Optional[str] = None
+        # Delta and stop events address blocks by their index in the message. That
+        # index also counts blocks content_blocks leaves out (server_tool_use,
+        # web_search_tool_result, ...), so keep a list that lines up with it.
+        self._blocks_by_index: List[StreamingContentBlock] = []
 
     def consume(self, event: Any) -> None:
         event_usage = extract_anthropic_usage_from_event(event)
@@ -29,6 +33,8 @@ def consume(self, event: Any) -> None:
         if getattr(event, "type", None) == "content_block_start":
             block, tool = handle_anthropic_content_block_start(event)
 
+            self._blocks_by_index.append(block or {})
+
             if block:
                 self.content_blocks.append(block)
                 if block.get("type") in ("text", "thinking"):
@@ -45,12 +51,14 @@ def consume(self, event: Any) -> None:
         if delta_text:
             self.accumulated_content += delta_text
 
-        handle_anthropic_tool_delta(event, self.content_blocks, self.tools_in_progress)
+        handle_anthropic_tool_delta(
+            event, self._blocks_by_index, self.tools_in_progress
+        )
 
         if getattr(event, "type", None) == "content_block_stop":
             self.current_text_block = None
             finalize_anthropic_tool_input(
-                event, self.content_blocks, self.tools_in_progress
+                event, self._blocks_by_index, self.tools_in_progress
             )
 
         if getattr(event, "type", None) == "message_delta":

posthog/test/ai/anthropic/test_anthropic.py

@@ -1212,6 +1212,88 @@ def test_streaming_with_tool_calls(mock_client, mock_anthropic_stream_with_tools
         assert "output_tokens" in props["$ai_usage"]
 
 
+def test_streaming_tool_call_after_server_tool_blocks(mock_client):
+    """Stream events address content blocks by their index in the message. Blocks the
+    wrapper does not format (server_tool_use, web_search_tool_result) still take an
+    index, so a later client tool call must keep its streamed arguments."""
+    from anthropic.types import (
+        InputJSONDelta,
+        ServerToolUseBlock,
+        ToolUseBlock,
+        WebSearchToolResultBlock,
+    )
+
+    def content_block(index, block, deltas=()):
+        yield RawContentBlockStartEvent(
+            type="content_block_start", index=index, content_block=block
+        )
+        for delta in deltas:
+            yield RawContentBlockDeltaEvent(
+                type="content_block_delta", index=index, delta=delta
+            )
+        yield RawContentBlockStopEvent(type="content_block_stop", index=index)
+
+    def stream_generator():
+        yield from content_block(
+            0,
+            ServerToolUseBlock(
+                type="server_tool_use", id="srvtoolu_1", name="web_search", input={}
+            ),
+            [
+                InputJSONDelta(
+                    type="input_json_delta", partial_json='{"query": "weather sf"}'
+                )
+            ],
+        )
+        yield from content_block(
+            1,
+            WebSearchToolResultBlock(
+                type="web_search_tool_result", tool_use_id="srvtoolu_1", content=[]
+            ),
+        )
+        yield from content_block(
+            2,
+            TextBlock(type="text", text=""),
+            [TextDelta(type="text_delta", text="Booking a table.")],
+        )
+        yield from content_block(
+            3,
+            ToolUseBlock(type="tool_use", id="toolu_1", name="book_table", input={}),
+            [
+                InputJSONDelta(type="input_json_delta", partial_json='{"city": '),
+                InputJSONDelta(type="input_json_delta", partial_json='"SF"}'),
+            ],
+        )
+
+    with patch(
+        "anthropic.resources.Messages.create",
+        return_value=stream_generator(),
+    ):
+        client = Anthropic(api_key="test-key", posthog_client=mock_client)
+        response = client.messages.create(
+            model="claude-sonnet-4-20250514",
+            messages=[{"role": "user", "content": "Book dinner if it's sunny"}],
+            stream=True,
+            posthog_distinct_id="test-id",
+        )
+        list(response)
+
+    props = mock_client.capture.call_args[1]["properties"]
+    assert props["$ai_output_choices"] == [
+        {
+            "role": "assistant",
+            "content": [
+                {"type": "text", "text": "Booking a table."},
+                {
+                    "type": "function",
+                    "id": "toolu_1",
+                    "function": {"name": "book_table", "arguments": {"city": "SF"}},
+                },
+            ],
+        }
+    ]
+
+
 def test_async_streaming_with_tool_calls(mock_client, mock_anthropic_stream_with_tools):
     """Test that tool calls are properly captured in async streaming mode."""
     import asyncio