Python: Finalize abandoned functional workflow streams without ContextVar leaks

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

@@ -6,8 +6,8 @@
 import sys
 import traceback as _traceback
 import warnings
-from collections.abc import Generator, Mapping
-from contextlib import contextmanager
+from collections.abc import Callable, Generator, Mapping
+from contextlib import contextmanager, suppress
 from contextvars import ContextVar
 from dataclasses import dataclass
 from enum import Enum
@@ -47,12 +47,33 @@ def _current_event_origin() -> WorkflowEventSource:
 
 @contextmanager
 def _framework_event_origin() -> Generator[None]:  # pyright: ignore[reportUnusedFunction]
-    """Temporarily mark subsequently created events as originating from the framework (internal)."""
+    """Temporarily mark subsequently created events as originating from the framework (internal).
+
+    Callers must not ``yield`` from an async generator while this manager is active.
+    Async-generator finalization can inject ``GeneratorExit`` from a different
+    ``Context`` than the one that created the token (for example when an abandoned
+    ``ResponseStream`` is garbage-collected), and ``ContextVar.reset`` then raises.
+    """
     token = _event_origin_context.set(WorkflowEventSource.FRAMEWORK)
     try:
         yield
     finally:
-        _event_origin_context.reset(token)
+        with suppress(ValueError):
+            # Token may have been created in a different Context when an
+            # abandoned ResponseStream is garbage-collected. Leave the var
+            # as-is rather than raising during generator/GC cleanup.
+            _event_origin_context.reset(token)
+
+
+def _framework_event(factory: Callable[..., WorkflowEvent[Any]], *args: Any, **kwargs: Any) -> WorkflowEvent[Any]:
+    """Build a framework-origin event and return it after resetting the origin token.
+
+    Callers can ``yield`` the result without holding ``_framework_event_origin()``
+    across an async-generator yield, which would leak the ContextVar token if the
+    stream is abandoned and finalized from a different Context.
+    """
+    with _framework_event_origin():
+        return factory(*args, **kwargs)
 
 
 class WorkflowRunState(str, Enum):

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

@@ -52,13 +52,18 @@
 from .._feature_stage import ExperimentalFeature, experimental
 from .._serialization import make_json_safe
 from .._types import AgentResponse, AgentResponseUpdate, ResponseStream
-from ..observability import OtelAttr, capture_exception, create_workflow_span
+from ..observability import (
+    OtelAttr,
+    _activate_span,
+    capture_exception,
+    start_workflow_span,
+)
 from ._checkpoint import CheckpointStorage, WorkflowCheckpoint
 from ._events import (
     WorkflowErrorDetails,
     WorkflowEvent,
     WorkflowRunState,
-    _framework_event_origin,
+    _framework_event,
 )
 from ._workflow import WorkflowRunResult
 
@@ -1029,111 +1034,114 @@ async def _on_step_completed() -> None:
 
             ctx._on_step_completed = _on_step_completed
 
-        # Tracing
+        # Tracing: start the run span without attaching it. Attaching with
+        # create_workflow_span() across a yield leaves OpenTelemetry's context
+        # token set when this generator is later closed on GC from a different
+        # Context. Activate the span only around non-yielding work.
         attributes: dict[str, Any] = {OtelAttr.WORKFLOW_NAME: self.name}
         if self.description:
             attributes[OtelAttr.WORKFLOW_DESCRIPTION] = self.description
 
-        with create_workflow_span(OtelAttr.WORKFLOW_RUN_SPAN, attributes) as span:
-            saw_request = False
-            try:
-                span.add_event(OtelAttr.WORKFLOW_STARTED)
+        span = start_workflow_span(OtelAttr.WORKFLOW_RUN_SPAN, attributes)
+        saw_request = False
+        try:
+            span.add_event(OtelAttr.WORKFLOW_STARTED)
 
-                with _framework_event_origin():
-                    yield WorkflowEvent.started()
-                with _framework_event_origin():
-                    yield WorkflowEvent.status(WorkflowRunState.IN_PROGRESS)
+            yield _framework_event(WorkflowEvent.started)
+            yield _framework_event(WorkflowEvent.status, WorkflowRunState.IN_PROGRESS)
 
-                # Execute the user function
+            # Execute the user function with the run span current so nested
+            # executor/processing spans parent correctly.
+            with _activate_span(span):
                 return_value = await self._execute(ctx, message)
 
                 # Emit the return value as the workflow output.
                 if return_value is not None:
-                    with _framework_event_origin():
-                        await ctx.add_event(WorkflowEvent("output", executor_id=self.name, data=return_value))
+                    await ctx.add_event(
+                        _framework_event(WorkflowEvent, "output", executor_id=self.name, data=return_value)
+                    )
 
                 # Persist step cache for response-only replay
                 self._last_step_cache = dict(ctx._step_cache)
                 self._last_step_cache_auto_request_info_counts = dict(ctx._step_cache_auto_request_info_counts)
 
-                # Yield collected events.
-                # NOTE: Events are buffered during _execute() and yielded after
-                # the user function completes.  This is *not* true streaming —
-                # all events have already been produced by this point.  True
-                # per-token streaming from inner agent calls is a future
-                # enhancement.
-                for event in ctx._get_events():
-                    if event.type == "request_info":
-                        saw_request = True
-                    yield event
-                    if event.type == "request_info":
-                        with _framework_event_origin():
-                            yield WorkflowEvent.status(WorkflowRunState.IN_PROGRESS_PENDING_REQUESTS)
-
-                # Save final checkpoint if storage is available
-                if storage is not None:
-                    await self._save_checkpoint(ctx, storage, ckpt_chain[0])
-
-                # Final status
-                if saw_request:
-                    self._last_pending_request_ids = set(ctx._pending_requests)
-                    with _framework_event_origin():
-                        yield WorkflowEvent.status(WorkflowRunState.IDLE_WITH_PENDING_REQUESTS)
-                else:
-                    # Clean completion — drop cross-run replay state.
-                    self._last_message = None
-                    self._last_step_cache = {}
-                    self._last_step_cache_auto_request_info_counts = {}
-                    self._last_pending_request_ids = set()
-                    with _framework_event_origin():
-                        yield WorkflowEvent.status(WorkflowRunState.IDLE)
+            # Yield collected events.
+            # NOTE: Events are buffered during _execute() and yielded after
+            # the user function completes.  This is *not* true streaming —
+            # all events have already been produced by this point.  True
+            # per-token streaming from inner agent calls is a future
+            # enhancement.
+            for event in ctx._get_events():
+                if event.type == "request_info":
+                    saw_request = True
+                yield event
+                if event.type == "request_info":
+                    yield _framework_event(WorkflowEvent.status, WorkflowRunState.IN_PROGRESS_PENDING_REQUESTS)
+
+            # Save final checkpoint if storage is available
+            if storage is not None:
+                await self._save_checkpoint(ctx, storage, ckpt_chain[0])
+
+            # Final status
+            if saw_request:
+                self._last_pending_request_ids = set(ctx._pending_requests)
+                yield _framework_event(WorkflowEvent.status, WorkflowRunState.IDLE_WITH_PENDING_REQUESTS)
+            else:
+                # Clean completion — drop cross-run replay state.
+                self._last_message = None
+                self._last_step_cache = {}
+                self._last_step_cache_auto_request_info_counts = {}
+                self._last_pending_request_ids = set()
+                yield _framework_event(WorkflowEvent.status, WorkflowRunState.IDLE)
 
-                span.add_event(OtelAttr.WORKFLOW_COMPLETED)
+            span.add_event(OtelAttr.WORKFLOW_COMPLETED)
 
-            except WorkflowInterrupted:
-                # Persist step cache for response-only replay
-                self._last_step_cache = dict(ctx._step_cache)
-                self._last_step_cache_auto_request_info_counts = dict(ctx._step_cache_auto_request_info_counts)
-                self._last_pending_request_ids = set(ctx._pending_requests)
+        except WorkflowInterrupted:
+            # Persist step cache for response-only replay
+            self._last_step_cache = dict(ctx._step_cache)
+            self._last_step_cache_auto_request_info_counts = dict(ctx._step_cache_auto_request_info_counts)
+            self._last_pending_request_ids = set(ctx._pending_requests)
 
-                # HITL interruption — yield events collected so far
-                for event in ctx._get_events():
-                    if event.type == "request_info":
-                        saw_request = True
-                    yield event
-                    if event.type == "request_info":
-                        with _framework_event_origin():
-                            yield WorkflowEvent.status(WorkflowRunState.IN_PROGRESS_PENDING_REQUESTS)
+            # HITL interruption — yield events collected so far
+            for event in ctx._get_events():
+                if event.type == "request_info":
+                    saw_request = True
+                yield event
+                if event.type == "request_info":
+                    yield _framework_event(WorkflowEvent.status, WorkflowRunState.IN_PROGRESS_PENDING_REQUESTS)
 
-                # Save checkpoint
-                if storage is not None:
-                    await self._save_checkpoint(ctx, storage, ckpt_chain[0])
+            # Save checkpoint
+            if storage is not None:
+                await self._save_checkpoint(ctx, storage, ckpt_chain[0])
 
-                with _framework_event_origin():
-                    yield WorkflowEvent.status(WorkflowRunState.IDLE_WITH_PENDING_REQUESTS)
+            yield _framework_event(WorkflowEvent.status, WorkflowRunState.IDLE_WITH_PENDING_REQUESTS)
 
-                span.add_event(OtelAttr.WORKFLOW_COMPLETED)
+            span.add_event(OtelAttr.WORKFLOW_COMPLETED)
 
-            except Exception as exc:
-                # Yield any events collected before the failure
-                for event in ctx._get_events():
-                    yield event
-
-                details = WorkflowErrorDetails.from_exception(exc)
-                with _framework_event_origin():
-                    yield WorkflowEvent.failed(details)
-                with _framework_event_origin():
-                    yield WorkflowEvent.status(WorkflowRunState.FAILED)
-
-                span.add_event(
-                    name=OtelAttr.WORKFLOW_ERROR,
-                    attributes={
-                        "error.message": str(exc),
-                        "error.type": type(exc).__name__,
-                    },
-                )
-                capture_exception(span, exception=exc)
-                raise
+        except Exception as exc:
+            # Yield any events collected before the failure
+            for event in ctx._get_events():
+                yield event
+
+            details = WorkflowErrorDetails.from_exception(exc)
+            yield _framework_event(WorkflowEvent.failed, details)
+            yield _framework_event(WorkflowEvent.status, WorkflowRunState.FAILED)
+
+            span.add_event(
+                name=OtelAttr.WORKFLOW_ERROR,
+                attributes={
+                    "error.message": str(exc),
+                    "error.type": type(exc).__name__,
+                },
+            )
+            capture_exception(span, exception=exc)
+            raise
+        finally:
+            # ResponseStream cleanup_hooks do not run when the generator is
+            # closed by GC. Release the run lock here so a follow-up run
+            # after an abandoned stream is not rejected as concurrent.
+            self._release_run_guard()
+            span.end()
 
     async def _execute(self, ctx: RunContext, message: Any) -> Any:
         """Run the user's async function with the active context."""
@@ -1296,9 +1304,12 @@ def _ensure_not_running(self) -> None:
             raise RuntimeError("Workflow is already running. Concurrent executions are not allowed.")
         self._is_running = True
 
-    async def _run_cleanup(self) -> None:
+    def _release_run_guard(self) -> None:
         self._is_running = False
 
+    async def _run_cleanup(self) -> None:
+        self._release_run_guard()
+
 
 # ---------------------------------------------------------------------------
 # @workflow decorator

python/packages/core/agent_framework/observability.py

@@ -2939,17 +2939,48 @@ def workflow_tracer() -> Tracer:
     return get_tracer() if OBSERVABILITY_SETTINGS.ENABLED else trace.NoOpTracer()
 
 
-def create_workflow_span(
+def _workflow_span_attributes(
     name: str,
     attributes: Mapping[str, str | int] | None = None,
-    kind: trace.SpanKind = trace.SpanKind.INTERNAL,
-) -> _AgnosticContextManager[trace.Span]:
-    """Create a generic workflow span."""
-    span_attributes = dict(attributes) if attributes is not None else {}
+) -> dict[str, str | int] | None:
+    span_attributes: dict[str, str | int] = dict(attributes) if attributes is not None else {}
     conversation_id = _TELEMETRY_CONVERSATION_ID.get()
     if name == OtelAttr.WORKFLOW_RUN_SPAN and conversation_id is not None:
         span_attributes.setdefault(OtelAttr.CONVERSATION_ID, conversation_id)
-    return workflow_tracer().start_as_current_span(name, kind=kind, attributes=span_attributes or None)
+    return span_attributes or None
+
+
+def create_workflow_span(
+    name: str,
+    attributes: Mapping[str, str | int] | None = None,
+    kind: trace.SpanKind = trace.SpanKind.INTERNAL,
+) -> _AgnosticContextManager[trace.Span]:
+    """Create a generic workflow span attached as the current span.
+
+    Do not use this from an async generator that yields to callers: attaching
+    across a ``yield`` leaves OpenTelemetry's context token set when the
+    generator is later closed on GC from a different ``Context``. Streaming
+    workflow runs should use :func:`start_workflow_span` instead.
+    """
+    return workflow_tracer().start_as_current_span(
+        name, kind=kind, attributes=_workflow_span_attributes(name, attributes)
+    )
+
+
+def start_workflow_span(
+    name: str,
+    attributes: Mapping[str, str | int] | None = None,
+    kind: trace.SpanKind = trace.SpanKind.INTERNAL,
+) -> trace.Span:
+    """Start a workflow span without attaching it to the current context.
+
+    Streaming generators that yield to callers must start the run span this
+    way and activate it only around non-yielding work (see
+    :func:`_activate_span`). Attaching with :func:`create_workflow_span`
+    across a yield causes ``ValueError: Token was created in a different
+    Context`` when the generator is garbage-collected.
+    """
+    return workflow_tracer().start_span(name, kind=kind, attributes=_workflow_span_attributes(name, attributes))
 
 
 def create_processing_span(

python/packages/core/tests/core/test_observability.py

@@ -2172,6 +2172,24 @@ def test_create_workflow_span(span_exporter):
     assert spans[0].attributes["key"] == "value"
 
 
+def test_start_workflow_span_does_not_attach_as_current(span_exporter: InMemorySpanExporter) -> None:
+    """start_workflow_span must not attach, so callers can yield without a dangling OTel token."""
+    from opentelemetry import trace
+
+    from agent_framework.observability import start_workflow_span
+
+    span_exporter.clear()  # type: ignore[attr-defined]
+    before = trace.get_current_span()
+    span = start_workflow_span("test_workflow_unattached", attributes={"key": "value"})
+    assert trace.get_current_span() is before
+    span.end()
+
+    spans = span_exporter.get_finished_spans()  # type: ignore[attr-defined]
+    assert len(spans) == 1
+    assert spans[0].name == "test_workflow_unattached"
+    assert spans[0].attributes["key"] == "value"
+
+
 def test_create_workflow_span_uses_scoped_conversation_id(span_exporter: InMemorySpanExporter) -> None:
     """An ambient conversation id is applied only within its workflow execution scope."""
     from agent_framework.observability import (