fix: preserve stream error status

agentops/instrumentation/common/streaming.py

@@ -67,8 +67,12 @@ def _process_chunk(self, chunk: Any):
                     self.token_usage.completion_tokens or 0
                 ) + chunk_usage.completion_tokens
 
-    def _finalize(self):
-        """Finalize the stream processing."""
+    def _finalize(self, *, set_success_status: bool = True):
+        """Finalize the stream processing.
+
+        Args:
+            set_success_status: Whether to mark the span as successful after setting final attributes.
+        """
         try:
             # Set final content
             final_content = "".join(self.accumulated_content)
@@ -87,7 +91,8 @@ def _finalize(self):
             for attr_name, value in self.token_usage.to_attributes().items():
                 safe_set_attribute(self.span, attr_name, value)
 
-            self.span.set_status(Status(StatusCode.OK))
+            if set_success_status:
+                self.span.set_status(Status(StatusCode.OK))
         except Exception as e:
             logger.error(f"Error finalizing stream: {e}")
             self.span.set_status(Status(StatusCode.ERROR, str(e)))
@@ -100,32 +105,36 @@ class SyncStreamWrapper(BaseStreamWrapper):
     """Wrapper for synchronous streaming responses."""
 
     def __iter__(self):
+        failed = False
         try:
             for chunk in self.stream:
                 self._process_chunk(chunk)
                 yield chunk
         except Exception as e:
+            failed = True
             self.span.set_status(Status(StatusCode.ERROR, str(e)))
             self.span.record_exception(e)
             raise
         finally:
-            self._finalize()
+            self._finalize(set_success_status=not failed)
 
 
 class AsyncStreamWrapper(BaseStreamWrapper):
     """Wrapper for asynchronous streaming responses."""
 
     async def __aiter__(self):
+        failed = False
         try:
             async for chunk in self.stream:
                 self._process_chunk(chunk)
                 yield chunk
         except Exception as e:
+            failed = True
             self.span.set_status(Status(StatusCode.ERROR, str(e)))
             self.span.record_exception(e)
             raise
         finally:
-            self._finalize()
+            self._finalize(set_success_status=not failed)
 
 
 def create_stream_wrapper_factory(

tests/unit/instrumentation/common/test_streaming.py

@@ -2,6 +2,8 @@
 from unittest.mock import Mock, patch
 from types import SimpleNamespace
 
+from opentelemetry.trace import StatusCode
+
 from agentops.instrumentation.common.streaming import (
     BaseStreamWrapper,
     SyncStreamWrapper,
@@ -248,7 +250,8 @@ def failing_stream():
         with pytest.raises(ValueError, match="Test error"):
             list(wrapper)
 
-        mock_span.set_status.assert_called()
+        statuses = [call.args[0].status_code for call in mock_span.set_status.call_args_list]
+        assert statuses == [StatusCode.ERROR]
         mock_span.record_exception.assert_called_once()
         mock_span.end.assert_called_once()
 
@@ -294,7 +297,8 @@ async def failing_async_stream():
             async for chunk in wrapper:
                 pass
 
-        mock_span.set_status.assert_called()
+        statuses = [call.args[0].status_code for call in mock_span.set_status.call_args_list]
+        assert statuses == [StatusCode.ERROR]
         mock_span.record_exception.assert_called_once()
         mock_span.end.assert_called_once()