From ae2f2dac76c19d10584f21e931bb195ff711baa7 Mon Sep 17 00:00:00 2001
From: kzahiri1 <kzahiri@dons.usfca.edu>
Date: Thu, 1 Oct 2026 12:56:32 -0700
Subject: [PATCH] fix(google-genai): retain URL context metadata in streams

---
 .../google_genai/test_google_genai.py         | 45 +++++++++++++++++++
 .../integrations/google_genai/tracing.py      |  2 +
 2 files changed, 47 insertions(+)

diff --git a/py/src/braintrust/integrations/google_genai/test_google_genai.py b/py/src/braintrust/integrations/google_genai/test_google_genai.py
index a760c34c3..2ddb0efa4 100644
--- a/py/src/braintrust/integrations/google_genai/test_google_genai.py
+++ b/py/src/braintrust/integrations/google_genai/test_google_genai.py
@@ -8,6 +8,7 @@
 import pytest
 from braintrust import logger
 from braintrust.integrations.google_genai import setup_genai
+from braintrust.integrations.google_genai.tracing import _aggregate_generate_content_chunks
 from braintrust.integrations.test_utils import verify_autoinstrument_script
 from braintrust.integrations.versioning import detect_module_version, version_satisfies
 from braintrust.logger import Attachment
@@ -1244,6 +1245,50 @@ def test_serialize_content_item_with_content_and_binary_part():
 )
 
 
+@pytest.mark.parametrize("include_url_metadata", [True, False])
+def test_stream_preserves_url_context_metadata(include_url_metadata):
+    url_metadata = (
+        types.UrlContextMetadata(
+            url_metadata=[
+                types.UrlMetadata(
+                    retrieved_url="https://example.com/article",
+                    url_retrieval_status="URL_RETRIEVAL_STATUS_SUCCESS",
+                ),
+                types.UrlMetadata(
+                    retrieved_url="https://example.com/unavailable",
+                    url_retrieval_status="URL_RETRIEVAL_STATUS_ERROR",
+                ),
+            ]
+        )
+        if include_url_metadata
+        else None
+    )
+    chunks = [
+        types.GenerateContentResponse(
+            candidates=[types.Candidate(content=types.Content(parts=[types.Part(text="First ")]))]
+        ),
+        types.GenerateContentResponse(
+            candidates=[
+                types.Candidate(
+                    content=types.Content(parts=[types.Part(text="second")]),
+                    url_context_metadata=url_metadata,
+                    finish_reason="STOP",
+                )
+            ]
+        ),
+    ]
+
+    output, _, _ = _aggregate_generate_content_chunks(chunks, time.time())
+
+    assert output["text"] == "First second"
+    candidate = output["candidates"][0]
+    assert candidate["content"]["parts"] == [{"text": "First second"}]
+    if include_url_metadata:
+        assert candidate["url_context_metadata"] == url_metadata
+    else:
+        assert "url_context_metadata" not in candidate
+
+
 def _assert_grounding_metadata(span_output):
     """Assert that grounding metadata is present and well-structured in span output."""
     # The grounding_metadata should be present on the first candidate
diff --git a/py/src/braintrust/integrations/google_genai/tracing.py b/py/src/braintrust/integrations/google_genai/tracing.py
index c29fab40d..47b3ae96d 100644
--- a/py/src/braintrust/integrations/google_genai/tracing.py
+++ b/py/src/braintrust/integrations/google_genai/tracing.py
@@ -703,6 +703,8 @@ def _aggregate_generate_content_chunks(
                 candidate_dict["safety_ratings"] = candidate.safety_ratings
             if hasattr(candidate, "grounding_metadata") and candidate.grounding_metadata:
                 candidate_dict["grounding_metadata"] = candidate.grounding_metadata
+            if hasattr(candidate, "url_context_metadata") and candidate.url_context_metadata:
+                candidate_dict["url_context_metadata"] = candidate.url_context_metadata
 
             candidates.append(candidate_dict)