diff --git a/agentops/sdk/decorators/utility.py b/agentops/sdk/decorators/utility.py index e6894ad15..924316a4f 100644 --- a/agentops/sdk/decorators/utility.py +++ b/agentops/sdk/decorators/utility.py @@ -40,11 +40,11 @@ def _process_sync_generator(span: trace.Span, generator: types.GeneratorType): # Ensure span context is attached to the generator context context_api.attach(trace.set_span_in_context(span)) - # Yield from the generator while maintaining span context - yield from generator - - # End the span when generator is exhausted - span.end() + try: + yield from generator + finally: + # End the span on exhaustion, errors, and explicit generator.close(). + span.end() # No detach because of OpenTelemetry issue #2606 # Context will be detached during garbage collection diff --git a/tests/unit/sdk/test_decorators.py b/tests/unit/sdk/test_decorators.py index 96824d0fe..0f192ee78 100644 --- a/tests/unit/sdk/test_decorators.py +++ b/tests/unit/sdk/test_decorators.py @@ -1,6 +1,7 @@ from typing import AsyncGenerator import asyncio import pytest +from opentelemetry import context as context_api from agentops.sdk.decorators import agent, operation, session, workflow, task, tool from agentops.semconv import SpanKind @@ -9,6 +10,45 @@ from agentops.sdk.decorators.factory import create_entity_decorator +@pytest.fixture +def isolated_generator_context(): + token = context_api.attach(context_api.get_current()) + yield + context_api.detach(token) + + +@pytest.mark.usefixtures("isolated_generator_context") +@pytest.mark.parametrize("entity_kind", [SpanKind.TASK, SpanKind.SESSION]) +@pytest.mark.parametrize("termination", ["error", "close"]) +def test_sync_generator_span_ends_on_early_exit(instrumentation: InstrumentationTester, entity_kind, termination): + error = ValueError("stream failed") + cleanup = [] + + @create_entity_decorator(entity_kind) + def stream(): + try: + yield "first" + raise error + finally: + cleanup.append("closed") + + result = stream() + assert next(result) == "first" + assert instrumentation.get_finished_spans() == [] + + if termination == "error": + with pytest.raises(ValueError) as caught: + next(result) + assert caught.value is error + else: + result.close() + + assert cleanup == ["closed"] + spans = instrumentation.get_finished_spans() + assert len(spans) == 1 + assert spans[0].name == f"stream.{entity_kind}" + + class TestSpanNesting: """Tests for proper nesting of spans in the tracing hierarchy."""