Python: Finalize abandoned graph workflow streams without OTel context leaks

python/packages/core/agent_framework/_workflows/_workflow.py

@@ -19,7 +19,12 @@
 from .._tools import ToolTypes, normalize_tools
 from .._types import Content, ResponseStream
 from ..exceptions import WorkflowException
-from ..observability import OtelAttr, capture_exception, create_workflow_span
+from ..observability import (
+    OtelAttr,
+    _activate_span,  # pyright: ignore[reportPrivateUsage]
+    capture_exception,
+    start_workflow_span,
+)
 from ._checkpoint import CheckpointStorage
 from ._const import (
     DEFAULT_MAX_ITERATIONS,
@@ -551,10 +556,8 @@ async def _run_workflow_with_tracing(
         if self.description:
             attributes[OtelAttr.WORKFLOW_DESCRIPTION] = self.description
 
-        with create_workflow_span(
-            OtelAttr.WORKFLOW_RUN_SPAN,
-            attributes,
-        ) as span:
+        span = start_workflow_span(OtelAttr.WORKFLOW_RUN_SPAN, attributes)
+        try:
             emitted_in_progress_pending = False
             try:
                 # Add workflow started event (telemetry + surface state to consumers)
@@ -620,10 +623,17 @@ async def _run_workflow_with_tracing(
 
                 # Execute initial setup if provided
                 if initial_executor_fn:
-                    await initial_executor_fn()
-
-                # All executor executions happen within workflow span
-                async for event in self._runner.run_until_convergence():
+                    with _activate_span(span):
+                        await initial_executor_fn()
+
+                # Activate the workflow span for each runner pull, then detach before yielding to callers.
+                runner_events = self._runner.run_until_convergence()
+                while True:
+                    with _activate_span(span):
+                        try:
+                            event = await anext(runner_events)
+                        except StopAsyncIteration:
+                            break
                     yield event
 
                     if event.type == "request_info" and not emitted_in_progress_pending:
@@ -671,6 +681,8 @@ async def _run_workflow_with_tracing(
                 )
                 capture_exception(span, exception=exc)
                 raise
+        finally:
+            span.end()
 
     async def _execute_with_message_or_checkpoint(
         self,

python/packages/core/tests/workflow/test_workflow.py

@@ -1014,6 +1014,49 @@ async def test_workflow_unconsumed_stream_releases_run_lock() -> None:
     assert result.get_final_state() == WorkflowRunState.IDLE
 
 
+@pytest.mark.parametrize("enable_sensitive_data", [False], indirect=True)
+async def test_workflow_abandoned_stream_finalizes_without_context_errors(
+    caplog: pytest.LogCaptureFixture, span_exporter: Any
+) -> None:
+    """Abandoning a partially consumed graph stream must clean up without context errors."""
+    executor = IncrementExecutor(id="abandoned_stream_exec", limit=3, increment=1)
+    workflow = WorkflowBuilder(start_executor=executor).build()
+    loop = asyncio.get_running_loop()
+    loop_errors: list[BaseException] = []
+    original_handler = loop.get_exception_handler()
+
+    def capture_loop_exception(_loop: asyncio.AbstractEventLoop, context: dict[str, Any]) -> None:
+        exc = context.get("exception")
+        if isinstance(exc, BaseException):
+            loop_errors.append(exc)
+        if original_handler is not None:
+            original_handler(_loop, context)
+
+    loop.set_exception_handler(capture_loop_exception)
+    try:
+        with caplog.at_level(logging.ERROR, logger="opentelemetry"):
+            stream = workflow.run(NumberMessage(data=0), stream=True)
+            async for event in stream:
+                assert event.type == "started"
+                break
+
+            del stream
+            gc.collect()
+            for _ in range(5):
+                await asyncio.sleep(0)
+    finally:
+        loop.set_exception_handler(original_handler)
+
+    assert loop_errors == [], f"Abandoned stream leaked loop exceptions: {loop_errors!r}"
+    otel_errors = [
+        rec.getMessage()
+        for rec in caplog.records
+        if "Failed to detach context" in rec.getMessage() or "was created in a different Context" in rec.getMessage()
+    ]
+    assert otel_errors == [], f"Abandoned stream leaked OpenTelemetry errors: {otel_errors!r}"
+    assert any(span.name == "workflow.run" for span in span_exporter.get_finished_spans())
+
+
 async def test_workflow_unawaited_run_coroutine_releases_run_lock() -> None:
     """An un-awaited non-streaming ``run()`` coroutine must also not leak the lock.
 

python/packages/core/tests/workflow/test_workflow_observability.py

@@ -314,8 +314,10 @@ async def test_end_to_end_workflow_tracing(span_exporter: InMemorySpanExporter)
 
     # Run workflow (this should create run spans)
     events: list[Any] = []
-    async for event in workflow.run("test input", stream=True):
-        events.append(event)
+    with trace.get_tracer(__name__).start_as_current_span("outer") as outer_span:
+        async for event in workflow.run("test input", stream=True):
+            assert trace.get_current_span() is outer_span
+            events.append(event)
 
     # Verify workflow executed correctly
     assert len(executor1.processed_messages) == 1
@@ -347,6 +349,13 @@ async def test_end_to_end_workflow_tracing(span_exporter: InMemorySpanExporter)
 
     # Verify workflow span events
     workflow_span = workflow_spans[0]
+    assert workflow_span.parent is not None
+    assert workflow_span.parent.span_id == outer_span.get_span_context().span_id
+    assert workflow_span.context is not None
+    assert all(
+        processing_span.parent is not None and processing_span.parent.span_id == workflow_span.context.span_id
+        for processing_span in processing_spans
+    ), "Executor processing spans should be children of workflow.run"
     assert workflow_span.events is not None
     event_names = [event.name for event in workflow_span.events]
     assert "workflow.started" in event_names