fix(ai): carry Gemini URL context flag across stream chunks

URL context metadata can arrive on a different chunk than the final
usage_metadata, so the tool-use prompt tokens were dropped in streaming.
Track the flag across chunks in the sync and async clients.

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>

Author: pcbeingused333

posthog/ai/gemini/gemini.py

@@ -25,6 +25,7 @@
     extract_gemini_embedding_token_count as extract_gemini_embedding_token_count,
     extract_gemini_stop_reason_from_chunk,
     extract_gemini_usage_from_chunk,
+    has_gemini_url_context,
     format_gemini_streaming_output as format_gemini_streaming_output,
 )
 
@@ -267,10 +268,16 @@ def generator():
             nonlocal accumulated_content
             nonlocal stop_reason
             stream_error: Optional[BaseException] = None
+            url_context_seen = False
             try:
                 for chunk in response:
                     # Extract usage stats from chunk
-                    chunk_usage = extract_gemini_usage_from_chunk(chunk)
+                    url_context_seen = url_context_seen or has_gemini_url_context(
+                        chunk
+                    )
+                    chunk_usage = extract_gemini_usage_from_chunk(
+                        chunk, url_context_seen=url_context_seen
+                    )
 
                     if chunk_usage:
                         # Gemini reports cumulative totals, not incremental values

posthog/ai/gemini/gemini_async.py

@@ -25,6 +25,7 @@
     extract_gemini_embedding_token_count as extract_gemini_embedding_token_count,
     extract_gemini_stop_reason_from_chunk,
     extract_gemini_usage_from_chunk,
+    has_gemini_url_context,
     format_gemini_streaming_output as format_gemini_streaming_output,
 )
 
@@ -254,10 +255,16 @@ async def async_generator():
             nonlocal stop_reason
 
             stream_error: Optional[BaseException] = None
+            url_context_seen = False
             try:
                 async for chunk in response:
                     # Extract usage stats from chunk
-                    chunk_usage = extract_gemini_usage_from_chunk(chunk)
+                    url_context_seen = url_context_seen or has_gemini_url_context(
+                        chunk
+                    )
+                    chunk_usage = extract_gemini_usage_from_chunk(
+                        chunk, url_context_seen=url_context_seen
+                    )
 
                     if chunk_usage:
                         # Gemini reports cumulative totals, not incremental values

posthog/ai/gemini/gemini_converter.py

@@ -517,7 +517,7 @@ def extract_gemini_web_search_count(response: Any) -> int:
     return 0
 
 
-def _has_url_context(response: Any) -> bool:
+def has_gemini_url_context(response: Any) -> bool:
     """Whether a response or chunk carries URL context retrieval metadata."""
     candidates = getattr(response, "candidates", None)
     if not isinstance(candidates, (list, tuple)):
@@ -597,7 +597,7 @@ def extract_gemini_usage_from_response(response: Any) -> TokenUsage:
         return TokenUsage()
 
     usage = _extract_usage_from_metadata(
-        response.usage_metadata, bill_tool_use_tokens=_has_url_context(response)
+        response.usage_metadata, bill_tool_use_tokens=has_gemini_url_context(response)
     )
 
     # Add web search count if present
@@ -608,12 +608,17 @@ def extract_gemini_usage_from_response(response: Any) -> TokenUsage:
     return usage
 
 
-def extract_gemini_usage_from_chunk(chunk: Any) -> TokenUsage:
+def extract_gemini_usage_from_chunk(
+    chunk: Any, url_context_seen: bool = False
+) -> TokenUsage:
     """
     Extract usage statistics from a Gemini streaming chunk.
 
     Args:
         chunk: Streaming chunk from Gemini API
+        url_context_seen: whether an earlier chunk of the same stream carried
+            URL context metadata. Tool indicators can arrive on a different
+            chunk than the final usage_metadata, so callers carry this flag.
 
     Returns:
         TokenUsage with standardized usage statistics
@@ -631,7 +636,8 @@ def extract_gemini_usage_from_chunk(chunk: Any) -> TokenUsage:
         return usage
 
     usage_from_metadata = _extract_usage_from_metadata(
-        chunk.usage_metadata, bill_tool_use_tokens=_has_url_context(chunk)
+        chunk.usage_metadata,
+        bill_tool_use_tokens=url_context_seen or has_gemini_url_context(chunk),
     )
 
     # Merge the usage from metadata with any web search count we found

posthog/test/ai/gemini/test_gemini.py

@@ -1741,3 +1741,40 @@ def test_ai_lane_client_routes_through_capture_ai(
     mock_client.capture.assert_not_called()
     assert mock_client.capture_ai.call_count == 1
     assert mock_client.capture_ai.call_args[1]["event"] == "$ai_generation"
+
+
+def _url_context_stream_chunks():
+    from google.genai import types
+
+    # URL context metadata arrives on an early chunk; usage on the last one.
+    first = types.GenerateContentResponse(
+        candidates=[types.Candidate(url_context_metadata=types.UrlContextMetadata())]
+    )
+    last = types.GenerateContentResponse(
+        usage_metadata=types.GenerateContentResponseUsageMetadata(
+            prompt_token_count=20,
+            tool_use_prompt_token_count=300,
+            candidates_token_count=10,
+            total_token_count=330,
+        )
+    )
+    return first, last
+
+
+def test_streaming_url_context_on_earlier_chunk_bills_tool_use_tokens(
+    mock_client, mock_google_genai_client
+):
+    first, last = _url_context_stream_chunks()
+    mock_google_genai_client.models.generate_content_stream.return_value = iter(
+        [first, last]
+    )
+
+    client = Client(api_key="test-key", posthog_client=mock_client)
+    response = client.models.generate_content_stream(
+        model="gemini-2.0-flash", contents=["Summarise this page"]
+    )
+    list(response)
+
+    props = mock_client.capture.call_args[1]["properties"]
+    assert props["$ai_input_tokens"] == 320
+    assert props["$ai_output_tokens"] == 10

posthog/test/ai/gemini/test_gemini_async.py

@@ -1178,3 +1178,40 @@ async def mock_streaming_response():
     call_args = mock_client.capture.call_args[1]
     assert call_args["event"] == "$ai_generation"
     assert call_args["properties"]["$ai_provider"] == "gemini"
+
+
+async def test_async_streaming_url_context_on_earlier_chunk_bills_tool_use_tokens(
+    mock_client, mock_google_genai_client
+):
+    from google.genai import types
+
+    first = types.GenerateContentResponse(
+        candidates=[types.Candidate(url_context_metadata=types.UrlContextMetadata())]
+    )
+    last = types.GenerateContentResponse(
+        usage_metadata=types.GenerateContentResponseUsageMetadata(
+            prompt_token_count=20,
+            tool_use_prompt_token_count=300,
+            candidates_token_count=10,
+            total_token_count=330,
+        )
+    )
+
+    async def stream():
+        yield first
+        yield last
+
+    mock_google_genai_client.aio.models.generate_content_stream = AsyncMock(
+        return_value=stream()
+    )
+
+    client = AsyncClient(api_key="test-key", posthog_client=mock_client)
+    response = await client.models.generate_content_stream(
+        model="gemini-2.0-flash", contents=["Summarise this page"]
+    )
+    async for _ in response:
+        pass
+
+    props = mock_client.capture.call_args[1]["properties"]
+    assert props["$ai_input_tokens"] == 320
+    assert props["$ai_output_tokens"] == 10