fix(loongsuite): align advice and stream lifecycle

Author: sipercai

instrumentation-loongsuite/loongsuite-instrumentation-litellm/CHANGELOG.md

@@ -9,8 +9,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
-- Isolate streamed chunk accumulation and finalization callback failures so
-  instrumentation errors do not interrupt application stream consumption.
+- Route streamed chunk accumulation and finalization callbacks through shared
+  fail-open advice while keeping stream iteration and lifecycle handling in the
+  LiteLLM wrapper.
 
 ## Version 0.7.0 (2026-07-03)
 

instrumentation-loongsuite/loongsuite-instrumentation-litellm/src/opentelemetry/instrumentation/litellm/_stream_wrapper.py

@@ -44,16 +44,10 @@ def _debug_safely(message: str, *args: Any, exc_info: bool = False) -> None:
         pass
 
 
+@hook_advice("litellm", "stream_chunk")
 def _record_stream_chunk(accumulator: Any, chunk: Any) -> None:
     """Record probe-owned chunk state without affecting stream delivery."""
-    try:
-        accumulator.record_chunk(chunk)
-    except Exception as exc:  # pylint: disable=broad-exception-caught
-        _debug_safely(
-            "Error recording LiteLLM stream chunk: %s",
-            exc,
-            exc_info=True,
-        )
+    accumulator.record_chunk(chunk)
 
 
 @hook_advice("litellm", "stream_finalize")
@@ -232,7 +226,7 @@ class StreamWrapper:
 
     def __init__(
         self,
-        stream: Iterator,
+        stream: Iterator[Any],
         span: Any,
         callback: callable,
         invocation: Any = None,
@@ -331,26 +325,30 @@ def _finalize(self, error: Optional[BaseException] = None):
             return
 
         self._finalized = True
-        self._close_stream()
         try:
-            # Call the callback with only the last chunk
-            # Note: The callback is responsible for calling handler.stop_llm()
-            # or handler.fail_llm().
-            # which will end the span. We no longer call span.end() here.
-            if self.callback:
-                _invoke_stream_callback(
-                    self.callback,
-                    self.span,
-                    self.last_chunk,
-                    error,
-                )
+            self._close_stream()
+        finally:
+            try:
+                # Call the callback with only the last chunk
+                # Note: The callback is responsible for calling handler.stop_llm()
+                # or handler.fail_llm().
+                # which will end the span. We no longer call span.end() here.
+                if self.callback:
+                    _invoke_stream_callback(
+                        self.callback,
+                        self.span,
+                        self.last_chunk,
+                        error,
+                    )
 
-            # Clear reference to avoid holding memory
-            self.last_chunk = None
-        except Exception as e:
-            _debug_safely(
-                "Error finalizing LiteLLM stream: %s", e, exc_info=True
-            )
+                # Clear reference to avoid holding memory
+                self.last_chunk = None
+            except Exception as e:
+                _debug_safely(
+                    "Error finalizing LiteLLM stream: %s",
+                    e,
+                    exc_info=True,
+                )
 
     def get_output_messages(self) -> list[OutputMessage]:
         return self._accumulator.get_output_messages()
@@ -468,10 +466,12 @@ def _close_stream(self) -> None:
         if self._stream_closed:
             return
 
-        self._stream_closed = True
+        self._stream_closed = self._close_sync_stream()
+
+    def _close_sync_stream(self) -> bool:
         close = getattr(self.stream, "close", None)
         if not callable(close):
-            return
+            return False
 
         try:
             close()
@@ -481,36 +481,28 @@ def _close_stream(self) -> None:
                 exc,
                 exc_info=True,
             )
+            return False
+        return True
 
     async def _aclose_stream(self) -> None:
         if self._stream_closed:
             return
 
-        self._stream_closed = True
         aclose = getattr(self.stream, "aclose", None)
         if callable(aclose):
             try:
                 await aclose()
-                return
             except Exception as exc:
                 _debug_safely(
                     "Error closing LiteLLM async stream: %s",
                     exc,
                     exc_info=True,
                 )
+            else:
+                self._stream_closed = True
+                return
 
-        close = getattr(self.stream, "close", None)
-        if not callable(close):
-            return
-
-        try:
-            close()
-        except Exception as exc:
-            _debug_safely(
-                "Error closing LiteLLM async stream: %s",
-                exc,
-                exc_info=True,
-            )
+        self._stream_closed = self._close_sync_stream()
 
     def _finalize(self, error: Optional[BaseException] = None):
         """Finalize the span with data from last chunk."""

instrumentation-loongsuite/loongsuite-instrumentation-litellm/tests/test_stream_isolation.py

@@ -16,11 +16,11 @@
 
 import pytest
 
-from opentelemetry.instrumentation.litellm import _stream_wrapper
 from opentelemetry.instrumentation.litellm._stream_wrapper import (
     AsyncStreamWrapper,
     StreamWrapper,
 )
+from opentelemetry.util.genai import extended_advice
 
 
 def test_stream_chunk_advice_failure_preserves_chunk(monkeypatch):
@@ -60,7 +60,7 @@ def test_stream_chunk_logging_failure_also_preserves_chunk(monkeypatch):
         ),
     )
     monkeypatch.setattr(
-        _stream_wrapper.logger,
+        extended_advice._logger,
         "debug",
         lambda *_args, **_kwargs: (_ for _ in ()).throw(
             RuntimeError("logger boom")
@@ -110,6 +110,34 @@ def callback(*_args):
     assert calls == [True]
 
 
+def test_stream_close_base_exception_still_finalizes():
+    calls = []
+    expected = KeyboardInterrupt("business interruption")
+
+    class Source:
+        def __iter__(self):
+            return self
+
+        def __next__(self):
+            raise StopIteration
+
+        def close(self):
+            raise expected
+
+    stream = StreamWrapper(
+        stream=Source(),
+        span=None,
+        callback=lambda *_args: calls.append(_args),
+    )
+
+    with pytest.raises(KeyboardInterrupt) as caught:
+        stream.close()
+
+    assert caught.value is expected
+    assert len(calls) == 1
+    assert stream._finalized is True
+
+
 @pytest.mark.asyncio
 async def test_async_stream_chunk_advice_failure_preserves_chunk(monkeypatch):
     expected = object()
@@ -209,6 +237,29 @@ async def aclose(self):
     assert stream._finalized is True
 
 
+@pytest.mark.asyncio
+async def test_async_only_stream_can_aclose_after_sync_close():
+    close_calls = []
+
+    class Source:
+        async def aclose(self):
+            close_calls.append(True)
+
+    stream = AsyncStreamWrapper(
+        stream=Source(),
+        span=None,
+        callback=lambda *_args: None,
+    )
+
+    stream.close()
+    assert close_calls == []
+    assert stream._stream_closed is False
+
+    await stream.aclose()
+    assert close_calls == [True]
+    assert stream._stream_closed is True
+
+
 def test_async_stream_sync_close_base_exception_still_finalizes():
     calls = []
     expected = KeyboardInterrupt("business interruption")

util/opentelemetry-util-genai/src/opentelemetry/util/genai/__init__.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# LoongSuite Extension: expose a stable advice facade for instrumentations and
+# downstream packaging adaptations.
 from opentelemetry.util.genai.extended_advice import (
     async_hook_advice,
     hook_advice,

util/opentelemetry-util-genai/src/opentelemetry/util/genai/extended_advice.py

@@ -55,6 +55,9 @@ def hook_advice(
 ) -> Callable[[_F], _F]:
     """Decorate synchronous instrumentation-only work with fail-open semantics.
 
+    Suppressed failures return ``None``, so callers must not depend on advice
+    return values unless ``throw_exception`` is enabled.
+
     Generator functions are rejected because their bodies execute during later
     iteration, after this decorator has returned the generator object. Stream
     iteration must instead be isolated by the owning instrumentation wrapper.
@@ -100,6 +103,9 @@ def async_hook_advice(
 ) -> Callable[[_F], _F]:
     """Decorate asynchronous instrumentation-only work with fail-open semantics.
 
+    Suppressed failures return ``None``, so callers must not depend on advice
+    return values unless ``throw_exception`` is enabled.
+
     Async generators are rejected because their bodies execute during later
     iteration. The owning instrumentation wrapper must isolate that lifecycle.
     """