From d8603e9ac0168485792b36e33c1d33b1718b1adf Mon Sep 17 00:00:00 2001 From: WhaleTech <304937387+ryo-whaletech@users.noreply.github.com> Date: Tue, 15 Sep 2026 23:14:29 +0900 Subject: [PATCH] Python: Avoid OTel context detach errors in graph workflow streams --- .../agent_framework/_workflows/_workflow.py | 30 +++++++++---- .../core/tests/workflow/test_workflow.py | 43 +++++++++++++++++++ .../workflow/test_workflow_observability.py | 13 +++++- 3 files changed, 75 insertions(+), 11 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_workflow.py b/python/packages/core/agent_framework/_workflows/_workflow.py index 18fb5c9a4db..afbf45af11b 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/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, diff --git a/python/packages/core/tests/workflow/test_workflow.py b/python/packages/core/tests/workflow/test_workflow.py index 762eff1d36e..497b402dee9 100644 --- a/python/packages/core/tests/workflow/test_workflow.py +++ b/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. diff --git a/python/packages/core/tests/workflow/test_workflow_observability.py b/python/packages/core/tests/workflow/test_workflow_observability.py index eb2772be124..daaea7682fc 100644 --- a/python/packages/core/tests/workflow/test_workflow_observability.py +++ b/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