diff --git a/python/packages/core/agent_framework/_workflows/_events.py b/python/packages/core/agent_framework/_workflows/_events.py index 669887a3c2..6676e9a827 100644 --- a/python/packages/core/agent_framework/_workflows/_events.py +++ b/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 @@ -46,13 +46,36 @@ 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).""" +def _framework_event_origin() -> Generator[None]: + """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( # pyright: ignore[reportUnusedFunction] + 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): diff --git a/python/packages/core/agent_framework/_workflows/_functional.py b/python/packages/core/agent_framework/_workflows/_functional.py index 9e4deb340b..e83fa9e0e5 100644 --- a/python/packages/core/agent_framework/_workflows/_functional.py +++ b/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 diff --git a/python/packages/core/agent_framework/observability.py b/python/packages/core/agent_framework/observability.py index fd23410eee..d546b0c434 100644 --- a/python/packages/core/agent_framework/observability.py +++ b/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( diff --git a/python/packages/core/tests/core/test_observability.py b/python/packages/core/tests/core/test_observability.py index 65b254b60c..ba723743b3 100644 --- a/python/packages/core/tests/core/test_observability.py +++ b/python/packages/core/tests/core/test_observability.py @@ -2169,7 +2169,25 @@ def test_create_workflow_span(span_exporter): spans = span_exporter.get_finished_spans() # type: ignore[attr-defined] assert len(spans) == 1 assert spans[0].name == "test_workflow" - assert spans[0].attributes["key"] == "value" + assert spans[0].attributes["key"] == "value" # type: ignore[index] # pyrefly: ignore[unsupported-operation] # ty: ignore[not-subscriptable] + + +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" # type: ignore[index] # pyrefly: ignore[unsupported-operation] # ty: ignore[not-subscriptable] def test_create_workflow_span_uses_scoped_conversation_id(span_exporter: InMemorySpanExporter) -> None: diff --git a/python/packages/core/tests/workflow/test_functional_workflow.py b/python/packages/core/tests/workflow/test_functional_workflow.py index 96b7e6edf6..5ed7ae0e8a 100644 --- a/python/packages/core/tests/workflow/test_functional_workflow.py +++ b/python/packages/core/tests/workflow/test_functional_workflow.py @@ -5,6 +5,7 @@ from __future__ import annotations import asyncio +import gc import json import logging from collections.abc import Awaitable, Callable, Iterator @@ -25,6 +26,7 @@ RunContext, StepWrapper, WorkflowEvent, + WorkflowEventSource, WorkflowRunResult, WorkflowRunState, get_run_context, @@ -533,6 +535,95 @@ async def wf(x: int, ctx: RunContext) -> int: await wf.run(1) assert streaming_flag is False + async def test_abandoned_stream_finalizes_without_event_loop_error(self, caplog: pytest.LogCaptureFixture) -> None: + """Breaking out of a streaming run must not leak ContextVar tokens on GC. + + Regression for https://github.com/microsoft/agent-framework/issues/7787: + the run span and ``_framework_event_origin()`` used to stay open across + event yields, so abandoning the stream and letting it be garbage-collected + reset those tokens from a different Context. + """ + 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) + + @built_workflow + async def pipeline(x: int) -> int: + return await add_one(x) + + try: + with caplog.at_level(logging.ERROR, logger="opentelemetry"): + stream = pipeline.run(5, stream=True) + async for _event in stream: + 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}" + + follow_up = await pipeline.run(6) + assert follow_up.get_outputs() == [7] + + async def test_nested_processing_span_parents_under_workflow_run(self, span_exporter: Any) -> None: + """Spans opened during ``_execute`` must parent under the unattached ``workflow.run`` span.""" + from agent_framework.observability import OtelAttr, create_processing_span + + @step + async def traced_add(x: int) -> int: + with create_processing_span("traced_add", "StepWrapper", "int", "int"): + return x + 1 + + @built_workflow + async def pipeline(x: int) -> int: + return await traced_add(x) + + span_exporter.clear() # type: ignore[attr-defined] + result = await pipeline.run(5) + assert result.get_outputs() == [6] + + spans = span_exporter.get_finished_spans() # type: ignore[attr-defined] + run_spans = [s for s in spans if s.name == OtelAttr.WORKFLOW_RUN_SPAN] + process_spans = [s for s in spans if s.name == f"{OtelAttr.EXECUTOR_PROCESS_SPAN} traced_add"] + assert len(run_spans) == 1 + assert len(process_spans) == 1 + process_parent = process_spans[0].parent + assert process_parent is not None + assert process_parent.span_id == run_spans[0].context.span_id + + async def test_started_event_origin_is_framework_after_origin_manager_closes(self) -> None: + """Framework lifecycle events stay tagged FRAMEWORK even when yielded outside the origin CM.""" + + @built_workflow + async def pipeline(x: int) -> int: + return await add_one(x) + + stream = pipeline.run(5, stream=True) + started = await anext(aiter(stream)) + assert started.type == "started" + assert started.origin == WorkflowEventSource.FRAMEWORK + await stream.get_final_response() + # --------------------------------------------------------------------------- # Step passthrough outside workflow