[NET-1387] fix: Force commit inner stream to capture output on stream break

CHANGELOG.md

@@ -4,7 +4,7 @@ All notable changes to this project will be documented in this file.
 
 The format is based on Keep a Changelog and this project adheres to Semantic Versioning.
 
-## [0.1.97b1] - 2026-07-28
+## [0.1.97] -
 
 - **Label evaluation/simulation traces on the root span** - Root spans produced by evaluation test runs (`Netra.evaluation`) and simulation runs (`Netra.simulation`) now carry a `netra.trace.origin` attribute set to `evaluation`, letting the backend and frontend distinguish these traces from normal workflow invocations.
 
@@ -20,7 +20,7 @@ The format is based on Keep a Changelog and this project adheres to Semantic Ver
 
 - **Fix span attributes in OpenAI instrumentation** - Assistant completions no longer emit empty entries when the model returns `content: null` alongside tool calls, request messages now correctly handle non-dictionary objects (such as Pydantic ChatCompletionMessage instances) by converting them with model_as_dict() instead of skipping them, and assistant `tool_calls` arrays as well as `tool_call_id` values on tool messages are now captured and serialized as indexed prompt and completion span attributes.
 
-- **Fix set_root_output_stream handling** - `set_root_output_stream` now reliably commits output for streams, even when iteration ends early (for example, via `break` or `.close()`), and correctly handles plain iterables by setting their output immediately with a warning recommending `Netra.set_root_output()`. Only true single-pass iterators are wrapped as streams.
+- **Fix set_root_output_stream handling** – `set_root_output_stream` now forces a commit of the inner stream to capture output when a stream exits early (for example, via `break` or `.close()`). It also correctly handles plain iterables by setting their output immediately with a warning recommending `Netra.set_root_output()`. Only true single-pass iterators are wrapped as streams.
 
 - **Refactor stream wrapper architecture to use callback injection** - `stream_utils` is now a pure utility module with no Netra-internal imports. The commit logic (serialize and set attribute on root span) is injected as a callback from `SessionManager`, eliminating the circular dependency between `stream_utils` and `SessionManager`.
 
@@ -336,4 +336,4 @@ Users can be now overwrite the input and ouput attributes of spans created by in
 
 - Added utility to set input and output data for any active span in a trace
 
-[0.1.97b1]: https://github.com/KeyValueSoftwareSystems/netra-sdk-py/tree/main
+[0.1.97]: https://github.com/KeyValueSoftwareSystems/netra-sdk-py/tree/main

netra/instrumentation/stream_utils.py

@@ -53,6 +53,109 @@ def _generic_extractor(wrapper: Union["RootOutputSyncStreamWrapper", "RootOutput
     return "".join(wrapper._chunks)
 
 
+def _finalize_via_method(stream: Any) -> None:
+    """Try to finalize a Netra wrapper by calling ``_finalize`` or ``_finalize_span``.
+
+    ``_finalize`` is tried first because wrappers that use it (e.g. Agno)
+    have an idempotency guard (``_finalized`` flag), making them safe to
+    call unconditionally.  ``_finalize_span`` is the fallback for LLM
+    wrappers (openai, groq, cerebras, litellm, google_genai) that lack
+    such a guard.
+    """
+    if stream is None or not getattr(stream, "_netra_stream_wrapper", False):
+        return
+    for method_name in ("_finalize", "_finalize_span"):
+        fn = getattr(stream, method_name, None)
+        if fn is not None and callable(fn):
+            try:
+                fn()
+            except Exception:
+                logger.warning("_force_finalize_inner_stream: %s() failed", method_name, exc_info=True)
+            return
+
+
+def _force_finalize_inner_stream(iterator: Any, stream: Any) -> None:
+    """Force-finalize a **sync** inner stream so ``_netra_output`` is populated
+    before the outer wrapper's extractor reads it.
+
+    On early ``break`` the outer wrapper's ``_commit`` fires before the
+    inner wrapper has finalized.  This helper covers every inner iterator
+    variant without requiring changes to any instrumentation wrapper:
+
+    1. **Generator-based iterators** — the inner ``__iter__`` returns a
+       generator (either a plain non-Netra generator, or a Netra wrapper
+       that delegates to an internal plain generator).  Calling
+       ``iterator.close()`` throws ``GeneratorExit`` into the generator;
+       if its ``finally`` block triggers the wrapper's finalization,
+       ``_netra_output`` is set and we are done.
+
+    2. **Netra return-self iterators** — the inner wrapper's ``__iter__``
+       returns ``self`` and has no ``close()`` method.  We fall back to
+       calling the finalization method directly on *stream* (the original
+       inner wrapper).  All Netra instrumentation wrappers use either
+       ``_finalize()`` or ``_finalize_span()``; we try both.
+
+    Path 2 also serves as a fallback when path 1 closes a generator but
+    ``_netra_output`` is still unset (e.g. the generator catches
+    ``GeneratorExit`` without calling finalization).
+    """
+    if iterator is None:
+        return
+
+    # Path 1: generator-based inner iterator (sync only).
+    # Async iterators use ``aclose()`` (a coroutine) instead of ``close()``.
+    # Calling ``close()`` on an ObjectProxy-based async wrapper would close
+    # the underlying transport without triggering wrapper finalization.
+    if hasattr(iterator, "close") and not hasattr(iterator, "__anext__"):
+        try:
+            iterator.close()
+        except Exception:
+            logger.debug("_force_finalize_inner_stream: failed to close iterator", exc_info=True)
+
+    # If the inner wrapper already set _netra_output (either because
+    # path 1 triggered finalization, or because __next__'s StopIteration
+    # handler already called _finalize/_finalize_span on full exhaustion),
+    # there is nothing left to do.  Skipping path 2 here is essential
+    # because several instrumentation wrappers (openai, groq, cerebras,
+    # litellm, google_genai, elevenlabs) have no idempotency guard in
+    # _finalize_span — calling it twice would double-end the span.
+    if getattr(stream, "_netra_output", None) is not None:
+        return
+
+    # Path 2: return-self wrappers — call the finalization method directly.
+    # Only reached when _netra_output is still unset (i.e. early break).
+    _finalize_via_method(stream)
+
+
+async def _aforce_finalize_inner_stream(iterator: Any, stream: Any) -> None:
+    """Async variant of :func:`_force_finalize_inner_stream`.
+
+    Handles async generators via ``aclose()`` (path 1a) in addition to
+    sync generators (path 1b) and direct finalization (path 2).
+    """
+    if iterator is None:
+        return
+
+    # Path 1a: async generator — ``aclose()`` is a coroutine.
+    if hasattr(iterator, "aclose"):
+        try:
+            await iterator.aclose()
+        except Exception:
+            logger.debug("_aforce_finalize_inner_stream: failed to aclose iterator", exc_info=True)
+    # Path 1b: sync generator attached to an async wrapper (unusual but possible).
+    elif hasattr(iterator, "close") and not hasattr(iterator, "__anext__"):
+        try:
+            iterator.close()
+        except Exception:
+            logger.debug("_aforce_finalize_inner_stream: failed to close iterator", exc_info=True)
+
+    if getattr(stream, "_netra_output", None) is not None:
+        return
+
+    # Path 2: return-self wrappers — call the finalization method directly.
+    _finalize_via_method(stream)
+
+
 # Sync wrapper
 class RootOutputSyncStreamWrapper:
     """Wraps a **single-pass** sync iterator; on exhaustion commits the output
@@ -124,6 +227,7 @@ def _commit(self) -> None:
             return
         self._committed = True
         try:
+            _force_finalize_inner_stream(self._iterator, self._stream)
             self._commit_fn(self._extractor(self))
         except Exception:
             logger.debug("RootOutputSyncWrapper: failed to commit output", exc_info=True)
@@ -191,7 +295,7 @@ async def _aiter_gen(self) -> Any:
                     self._chunks.append(str(chunk))
                 yield chunk
         finally:
-            self._commit()
+            await self._acommit()
 
     def __aiter__(self) -> Any:
         return self._aiter_gen()
@@ -203,7 +307,7 @@ async def __anext__(self) -> Any:
                 self._chunks.append(str(chunk))
             return chunk
         except StopAsyncIteration:
-            self._commit()
+            await self._acommit()
             raise
 
     async def __aenter__(self) -> "RootOutputAsyncStreamWrapper":
@@ -215,7 +319,7 @@ async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
         if hasattr(self._stream, "__aexit__"):
             await self._stream.__aexit__(exc_type, exc_val, exc_tb)
         if exc_type is None:
-            self._commit()
+            await self._acommit()
 
     def __getattr__(self, name: str) -> Any:
         return getattr(self._stream, name)
@@ -224,11 +328,24 @@ def __del__(self) -> None:
         if not self._committed:
             self._commit()
 
+    async def _acommit(self) -> None:
+        """Async commit — uses ``aclose()`` to finalize async generator inner streams."""
+        if self._committed:
+            return
+        self._committed = True
+        try:
+            await _aforce_finalize_inner_stream(self._aiterator, self._stream)
+            self._commit_fn(self._extractor(self))
+        except Exception:
+            logger.debug("RootOutputAsyncWrapper: failed to commit output", exc_info=True)
+
     def _commit(self) -> None:
+        """Sync fallback for ``__del__`` and other non-async contexts."""
         if self._committed:
             return
         self._committed = True
         try:
+            _force_finalize_inner_stream(self._aiterator, self._stream)
             self._commit_fn(self._extractor(self))
         except Exception:
             logger.debug("RootOutputAsyncWrapper: failed to commit output", exc_info=True)

netra/version.py

@@ -1 +1 @@
-__version__ = "0.1.97b1"
+__version__ = "0.1.97dev2"

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [project]
 name = "netra-sdk"
-version = "0.1.97b1"
+version = "0.1.97dev2"
 description = "A Python SDK for AI application observability that provides OpenTelemetry-based monitoring, tracing, and PII protection for LLM and vector database applications. Enables easy instrumentation, session tracking, and privacy-focused data collection for AI systems in production environments."
 authors = [
     {name = "Sooraj Thomas",email = "sooraj@keyvalue.systems"}