fix: preserve stream error status

agentops/instrumentation/common/streaming.py

@@ -4,6 +4,7 @@
 in a consistent way across different providers.
 """
 
+import asyncio
 from typing import Optional, Any, Dict, Callable
 from abc import ABC
 import time
@@ -67,8 +68,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 +92,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 +106,41 @@ 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 asyncio.CancelledError as e:
+            failed = True
+            self.span.set_status(Status(StatusCode.ERROR, str(e)))
+            self.span.record_exception(e)
+            raise
         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

@@ -1,7 +1,11 @@
+import asyncio
+
 import pytest
 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 +252,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,10 +299,32 @@ 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()
 
+    @pytest.mark.asyncio
+    async def test_async_iteration_with_cancellation(self):
+        """Test cancellation is recorded without overwriting it with success."""
+
+        async def cancelled_async_stream():
+            yield "chunk1"
+            raise asyncio.CancelledError("stream cancelled")
+
+        mock_span = Mock()
+        extract_content = lambda x: x
+        wrapper = AsyncStreamWrapper(cancelled_async_stream(), mock_span, extract_content)
+
+        with pytest.raises(asyncio.CancelledError, match="stream cancelled") as exc_info:
+            async for chunk in wrapper:
+                pass
+
+        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_with(exc_info.value)
+        mock_span.end.assert_called_once()
+
 
 class TestCreateStreamWrapperFactory:
     """Test the create_stream_wrapper_factory function."""