fix(agno): clean up async streams on iteration errors

agentops/instrumentation/agentic/agno/instrumentor.py

@@ -143,6 +143,14 @@ async def __anext__(self):
                 self.span.end()
                 self.streaming_context_manager.remove_context(self.agent_id)
             raise
+        except Exception as e:
+            if not self._consumed:
+                self._consumed = True
+                self.span.set_status(Status(StatusCode.ERROR, str(e)))
+                self.span.record_exception(e)
+                self.span.end()
+                self.streaming_context_manager.remove_context(self.agent_id)
+            raise
         finally:
             otel_context.detach(context_token)
 

tests/unit/instrumentation/test_agno_instrumentor.py

@@ -0,0 +1,59 @@
+from unittest.mock import MagicMock
+
+import pytest
+from opentelemetry import context as otel_context
+from opentelemetry.trace import StatusCode
+
+from agentops.instrumentation.agentic.agno.instrumentor import (
+    AsyncStreamingResultWrapper,
+    StreamingContextManager,
+)
+
+
+class FailingAsyncIterator:
+    def __init__(self):
+        self._yielded = False
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        if not self._yielded:
+            self._yielded = True
+            return "first event"
+        raise RuntimeError("stream failed")
+
+
+@pytest.mark.asyncio
+async def test_async_streaming_wrapper_cleans_up_post_yield_error_once():
+    span = MagicMock()
+    agent_id = "agent-id"
+    agent_context = otel_context.get_current()
+    context_manager = StreamingContextManager()
+    context_manager.store_context(agent_id, agent_context, span)
+    wrapper = AsyncStreamingResultWrapper(
+        FailingAsyncIterator(),
+        span,
+        agent_id,
+        agent_context,
+        context_manager,
+    )
+
+    assert await wrapper.__anext__() == "first event"
+
+    with pytest.raises(RuntimeError, match="stream failed") as exc_info:
+        await wrapper.__anext__()
+
+    status = span.set_status.call_args.args[0]
+    assert status.status_code is StatusCode.ERROR
+    assert status.description == "stream failed"
+    span.record_exception.assert_called_once_with(exc_info.value)
+    span.end.assert_called_once_with()
+    assert context_manager.get_context(agent_id) is None
+
+    with pytest.raises(RuntimeError, match="stream failed"):
+        await wrapper.__anext__()
+
+    span.set_status.assert_called_once()
+    span.record_exception.assert_called_once()
+    span.end.assert_called_once()