streaming fix

Author: AbhiPrasad

.agents/skills/sdk-integrations/SKILL.md

@@ -263,6 +263,28 @@ Assert on emitted spans (not just provider return values):
 
 For streaming, assert both the provider iterator/async-iterator still works AND the final span has aggregated `output` + stream-specific `metrics`.
 
+### Streaming lifecycle review
+
+Review every exit path, not just full consumption. Check the real provider's
+iterator, context-manager, cancellation, and garbage-collection behavior before
+choosing a wrapper. A proxy can preserve transport cleanup while still losing
+the span's final output and end time.
+
+- Cover exhaustion, provider errors, explicit close/cancel, and context-manager
+  exit where the provider supports it. Preserve exception and return semantics.
+- Cover `break` followed by dropping the last stream reference, and dropping a
+  stream before consuming any chunks. `break` alone does not close a retained
+  iterator; do not promise immediate finalization while callers still hold it.
+- Use existing recordings for partial-consumption tests. Drop the proxy, force
+  collection, and assert partial (or empty) output, `metrics.end`, correct
+  parentage, and no duplicate finalization after explicit close or exhaustion.
+- A GC fallback must not retain the stream through its callback or closure.
+  For async streams, do not run or schedule event-loop work from a finalizer;
+  finalize trace state and preserve the provider's own cleanup behavior.
+- Check that neither iteration nor cleanup leaves the stream span current in
+  the caller's context. GC is a best-effort fallback, not a substitute for
+  deterministic cleanup when the caller explicitly closes the stream.
+
 Cassettes live in `integrations/<provider>/cassettes/<version>/` (e.g. `cassettes/latest/`, `cassettes/0.48.0/`). Nox sets `BRAINTRUST_TEST_PACKAGE_VERSION` so cassettes land correctly. Do not add per-test `vcr_cassette_dir` / `cassette_library_dir` fixtures — `integrations/conftest.py` handles it. Re-record only when behavior intentionally changed. Sanitize binary media in both request and response bodies so checked-in cassettes do not retain large base64 payloads.
 
 Confirm the exact session name from `noxfile.py` — don't assume it matches the folder.

py/src/braintrust/integrations/google_discoveryengine/test_google_discoveryengine.py

@@ -1,8 +1,10 @@
 """Real Discovery Engine responses, recorded over REST and gRPC."""
 
+import gc
 import json
 import os
 import subprocess
+import weakref
 from contextlib import nullcontext
 from pathlib import Path
 from urllib.parse import urlsplit
@@ -398,7 +400,8 @@ def test_patch_scope():
 
 
 @pytest.mark.vcr("test_answer_query[True].yaml")
-def test_stream_close_preserves_parent(memory_logger, credentials, SERVING_CONFIG):
+@pytest.mark.parametrize("consume", ["close", "abandon", "unstarted"])
+def test_stream_lifecycle(memory_logger, credentials, SERVING_CONFIG, consume):
     from braintrust import current_span, start_span
     from braintrust.integrations.google_discoveryengine import wrap_google_discoveryengine
 
@@ -415,10 +418,20 @@ def test_stream_close_preserves_parent(memory_logger, credentials, SERVING_CONFI
             retry=None,
         )
         assert current_span() is parent
-        next(stream)
+        chunks = []
+        if consume != "unstarted":
+            for chunk in stream:
+                chunks.append(chunk)
+                if chunk.answer.answer_text:
+                    break
         assert current_span() is parent
-        stream.close()
-        stream.close()
+        if consume == "close":
+            stream.close()
+            stream.close()
+        stream_ref = weakref.ref(stream)
+        del stream
+        gc.collect()
+        assert stream_ref() is None
         assert current_span() is parent
     spans = memory_logger.pop()
     assert len(spans) == 2
@@ -427,7 +440,10 @@ def test_stream_close_preserves_parent(memory_logger, credentials, SERVING_CONFI
     )
     parent_row = next(span for span in spans if span["span_attributes"]["name"] == "caller")
     assert child["span_parents"] == [parent_row["span_id"]]
+    assert child["output"][0]["message"]["content"] == "".join(chunk.answer.answer_text for chunk in chunks)
     assert "end" in child["metrics"]
+    gc.collect()
+    assert memory_logger.pop() == []
 
 
 def test_auto_instrument_subprocess():
@@ -486,7 +502,7 @@ def broken(*args):
 
 
 @pytest.mark.asyncio
-@pytest.mark.parametrize("consume", ["read", "cancel", "aclose"])
+@pytest.mark.parametrize("consume", ["read", "cancel", "aclose", "abandon", "unstarted"])
 async def test_async_stream_lifecycle(memory_logger, credentials, vcr_cassette_dir, consume, SERVING_CONFIG):
     from braintrust import current_span, start_span
     from braintrust.integrations.google_discoveryengine import wrap_google_discoveryengine
@@ -521,13 +537,21 @@ async def test_async_stream_lifecycle(memory_logger, credentials, vcr_cassette_d
                     chunks.append(chunk)
                     assert current_span() is parent
                 assert await stream.read() is EOF
-            else:
-                chunks.append(await stream.__anext__())
+            elif consume != "unstarted":
+                async for chunk in stream:
+                    chunks.append(chunk)
+                    if chunk.answer.answer_text:
+                        break
                 if consume == "cancel":
                     assert stream.cancel()
-                else:
+                elif consume == "aclose":
                     await stream.aclose()
-                assert stream.cancelled()
+                if consume in ("cancel", "aclose"):
+                    assert stream.cancelled()
+            stream_ref = weakref.ref(stream)
+            del stream
+            gc.collect()
+            assert stream_ref() is None
             assert current_span() is parent
         spans = memory_logger.pop()
         assert len(spans) == 2
@@ -543,6 +567,8 @@ async def test_async_stream_lifecycle(memory_logger, credentials, vcr_cassette_d
         )
         assert child["output"][0]["message"]["content"] == expected_text
         assert "end" in child["metrics"]
+        gc.collect()
+        assert memory_logger.pop() == []
     finally:
         await client.transport.close()
 

py/src/braintrust/integrations/google_discoveryengine/tracing.py

@@ -7,6 +7,7 @@
 
 import logging
 import time
+import weakref
 from collections.abc import Mapping
 from itertools import islice
 
@@ -195,6 +196,9 @@ class _AnswerStream(ObjectProxy):
     def __init__(self, stream, state):
         super().__init__(stream)
         self._self_state = state
+        # Retain only trace state, not the proxy/provider stream. GC can finalize
+        # partial output while the provider handles its own transport cleanup.
+        weakref.finalize(self, state.finish)
         self._self_iterator = iter(stream)
 
     def __iter__(self):
@@ -232,6 +236,9 @@ class _AsyncAnswerStream(ObjectProxy):
     def __init__(self, stream, state):
         super().__init__(stream)
         self._self_state = state
+        # Retain only trace state, not the proxy/provider stream. GC can finalize
+        # partial output while the provider handles its own transport cleanup.
+        weakref.finalize(self, state.finish)
         self._self_iterator = None
 
     def __aiter__(self):