From b17b98268a228bbaa5e60313f0d401fbd8a84459 Mon Sep 17 00:00:00 2001 From: KXH Date: Sat, 25 Jul 2026 20:24:38 +0800 Subject: [PATCH 1/2] fix: clean up failed Agno async streams --- .../agentic/agno/instrumentor.py | 8 +++ .../instrumentation/test_agno_instrumentor.py | 59 +++++++++++++++++++ 2 files changed, 67 insertions(+) create mode 100644 tests/unit/instrumentation/test_agno_instrumentor.py diff --git a/agentops/instrumentation/agentic/agno/instrumentor.py b/agentops/instrumentation/agentic/agno/instrumentor.py index b8d41cbfe..cb1eeca3b 100644 --- a/agentops/instrumentation/agentic/agno/instrumentor.py +++ b/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) diff --git a/tests/unit/instrumentation/test_agno_instrumentor.py b/tests/unit/instrumentation/test_agno_instrumentor.py new file mode 100644 index 000000000..da7702ccd --- /dev/null +++ b/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() From 57d256d7f22fde62cc1f778d154b295ef71ac0fb Mon Sep 17 00:00:00 2001 From: KXH Date: Sun, 30 Aug 2026 04:40:25 +0000 Subject: [PATCH 2/2] fix: clean up cancelled Agno streams --- .../agentic/agno/instrumentor.py | 9 ++++ .../instrumentation/test_agno_instrumentor.py | 44 +++++++++++++++++++ 2 files changed, 53 insertions(+) diff --git a/agentops/instrumentation/agentic/agno/instrumentor.py b/agentops/instrumentation/agentic/agno/instrumentor.py index cb1eeca3b..3db6f0c51 100644 --- a/agentops/instrumentation/agentic/agno/instrumentor.py +++ b/agentops/instrumentation/agentic/agno/instrumentor.py @@ -7,6 +7,7 @@ parent-child span relationships. """ +import asyncio from typing import List, Any, Optional, Dict from opentelemetry import trace, context as otel_context from opentelemetry.trace import Status, StatusCode @@ -143,6 +144,14 @@ async def __anext__(self): self.span.end() self.streaming_context_manager.remove_context(self.agent_id) raise + except asyncio.CancelledError 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 except Exception as e: if not self._consumed: self._consumed = True diff --git a/tests/unit/instrumentation/test_agno_instrumentor.py b/tests/unit/instrumentation/test_agno_instrumentor.py index da7702ccd..39e333886 100644 --- a/tests/unit/instrumentation/test_agno_instrumentor.py +++ b/tests/unit/instrumentation/test_agno_instrumentor.py @@ -1,3 +1,4 @@ +import asyncio from unittest.mock import MagicMock import pytest @@ -24,6 +25,14 @@ async def __anext__(self): raise RuntimeError("stream failed") +class CancelledAsyncIterator(FailingAsyncIterator): + async def __anext__(self): + if not self._yielded: + self._yielded = True + return "first event" + raise asyncio.CancelledError("stream cancelled") + + @pytest.mark.asyncio async def test_async_streaming_wrapper_cleans_up_post_yield_error_once(): span = MagicMock() @@ -57,3 +66,38 @@ async def test_async_streaming_wrapper_cleans_up_post_yield_error_once(): span.set_status.assert_called_once() span.record_exception.assert_called_once() span.end.assert_called_once() + + +@pytest.mark.asyncio +async def test_async_streaming_wrapper_cleans_up_cancellation_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( + CancelledAsyncIterator(), + span, + agent_id, + agent_context, + context_manager, + ) + + assert await wrapper.__anext__() == "first event" + + with pytest.raises(asyncio.CancelledError, match="stream cancelled") 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 cancelled" + 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(asyncio.CancelledError, match="stream cancelled"): + await wrapper.__anext__() + + span.set_status.assert_called_once() + span.record_exception.assert_called_once() + span.end.assert_called_once()