From da6e66d21b222014ee02a3ee2b858f54fcec6e80 Mon Sep 17 00:00:00 2001 From: yashvanthange Date: Sun, 30 Aug 2026 00:08:22 +0530 Subject: [PATCH 1/5] Python: carry fan-in edge buffers through checkpoint restore A fan-in buffers messages until every source has produced one. That buffer was neither checkpointed nor reset on restore, so buffered messages were lost when resuming a rebuilt workflow, and a buffer left behind by a run that failed mid-superstep survived the restore and duplicated deliveries. Edge runner delivery state now travels with the checkpoint under the reserved shared-state key _edge_state. Both restore paths reset every edge runner before reapplying what the checkpoint held, so a checkpoint written before this change resets the buffer instead of leaking it. --- .../agent_framework/_workflows/_checkpoint.py | 9 +- .../core/agent_framework/_workflows/_const.py | 3 + .../_workflows/_edge_runner.py | 76 ++++++++++ .../agent_framework/_workflows/_runner.py | 45 +++++- .../core/tests/workflow/test_checkpoint.py | 78 ++++++++++ .../core/tests/workflow/test_runner.py | 87 ++++++++++- .../core/tests/workflow/test_sub_workflow.py | 142 ++++++++++++++++++ 7 files changed, 432 insertions(+), 8 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint.py b/python/packages/core/agent_framework/_workflows/_checkpoint.py index 62f225b04a8..d9b30d83919 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint.py @@ -52,10 +52,11 @@ class WorkflowCheckpoint: allows chaining checkpoints together to form a history of workflow states. timestamp: ISO 8601 timestamp when checkpoint was created messages: Messages exchanged between executors - state: Committed workflow state including user data and executor states. - This contains only committed state; pending state changes are not - included in checkpoints. Executor states are stored under the - reserved key '_executor_state'. + state: Committed workflow state including user data, executor states, and + edge runner delivery state. This contains only committed state; pending + state changes are not included in checkpoints. Executor states are stored + under the reserved key '_executor_state', and edge runner state, such as + partially filled fan-in buffers, under '_edge_state'. pending_request_info_events: Any pending request info events that have not yet been processed at the time of checkpointing. This allows the workflow to resume with the correct pending events after a restore. diff --git a/python/packages/core/agent_framework/_workflows/_const.py b/python/packages/core/agent_framework/_workflows/_const.py index e83025bbdc7..27b9c249617 100644 --- a/python/packages/core/agent_framework/_workflows/_const.py +++ b/python/packages/core/agent_framework/_workflows/_const.py @@ -6,6 +6,9 @@ # Key used to store executor state in state. EXECUTOR_STATE_KEY = "_executor_state" +# Key used to store edge runner delivery state (for example, fan-in buffers) in state. +EDGE_STATE_KEY = "_edge_state" + # Source identifier for internal workflow messages. INTERNAL_SOURCE_PREFIX = "internal" diff --git a/python/packages/core/agent_framework/_workflows/_edge_runner.py b/python/packages/core/agent_framework/_workflows/_edge_runner.py index c14582894b9..0bf83f53b5d 100644 --- a/python/packages/core/agent_framework/_workflows/_edge_runner.py +++ b/python/packages/core/agent_framework/_workflows/_edge_runner.py @@ -7,6 +7,7 @@ from collections.abc import Callable from typing import Any, cast +from ..exceptions import WorkflowCheckpointException from ..observability import EdgeGroupDeliveryStatus, OtelAttr, create_edge_group_processing_span from ._edge import ( Edge, @@ -57,6 +58,38 @@ async def send_message( """ raise NotImplementedError + @property + def state_key(self) -> str: + """Return a topology-derived key identifying this runner across workflow instances. + + ``EdgeGroup.id`` defaults to a random UUID, so it differs between two instances of the + same workflow definition and cannot key state that has to survive a rebuild. Checkpoint + compatibility is decided by graph topology, so the key is derived from topology too. + Builder validation rejects two edges that share a ``source -> target`` pair anywhere in + the workflow, so no two edge groups can produce the same key. + """ + edges = sorted(f"{edge.source_id}->{edge.target_id}" for edge in self._edge_group.edges) + return f"{self._edge_group.__class__.__name__}:{','.join(edges)}" + + def snapshot_state(self) -> dict[str, Any] | None: + """Capture in-flight delivery state so a checkpoint can restore it. + + Returns: + A serializable snapshot, or None when the runner holds no state to checkpoint. + """ + return None + + def restore_state(self, state: dict[str, Any] | None) -> None: + """Reset in-flight delivery state, then apply ``state`` when one was checkpointed. + + Called on every runner during checkpoint restoration, including with ``None``, so that + state left over from an interrupted run does not survive into the restored run. + + Args: + state: A snapshot previously produced by :meth:`snapshot_state`, or None to only reset. + """ + return + def _can_handle(self, executor_id: str, message: WorkflowMessage) -> bool: """Check if an executor can handle the given message data.""" if executor_id not in self._executors: @@ -297,6 +330,8 @@ def _validate_selection_result(self, selection_results: list[str]) -> bool: class FanInEdgeRunner(EdgeRunner): """Runner for fan-in edge groups.""" + _BUFFER_KEY = "buffer" + def __init__(self, edge_group: FanInEdgeGroup, executors: dict[str, Executor]) -> None: super().__init__(edge_group, executors) self._edges = edge_group.edges @@ -304,6 +339,47 @@ def __init__(self, edge_group: FanInEdgeGroup, executors: dict[str, Executor]) - # Key is the source executor ID, value is a list of messages self._buffer: dict[str, list[WorkflowMessage]] = defaultdict(list) + def snapshot_state(self) -> dict[str, Any] | None: + """Capture the buffered messages that are still waiting for the remaining sources. + + A fan-in only delivers once every source has produced a message, so a superstep + boundary can fall while the buffer holds a subset of them. Those messages have + already been drained from the runner context, so the checkpoint has to carry them. + """ + buffered = {source_id: list(messages) for source_id, messages in self._buffer.items() if messages} + if not buffered: + return None + return {self._BUFFER_KEY: buffered} + + def restore_state(self, state: dict[str, Any] | None) -> None: + """Reset the buffer, then refill it from the checkpointed snapshot when there is one. + + The reset matters on its own: a run that failed mid-superstep can leave messages from + a subset of sources in the buffer, and the sources are re-executed after the restore. + Without the reset those messages would be delivered twice, and the second delivery + could fire the fan-in before the restored superstep produced all of its messages. + """ + self._buffer.clear() + if state is None: + return + + buffered: Any = state.get(self._BUFFER_KEY, {}) + if not isinstance(buffered, dict): + raise WorkflowCheckpointException( + f"Fan-in buffer for edge group {self._edge_group.id} is not a dictionary. Unable to restore." + ) + + for source_id, messages in cast(dict[Any, Any], buffered).items(): + if ( + not isinstance(source_id, str) + or not isinstance(messages, list) + or not all(isinstance(message, WorkflowMessage) for message in cast(list[Any], messages)) + ): + raise WorkflowCheckpointException( + f"Fan-in buffer for edge group {self._edge_group.id} is malformed. Unable to restore." + ) + self._buffer[source_id] = list(cast(list[WorkflowMessage], messages)) + async def send_message( self, message: WorkflowMessage, diff --git a/python/packages/core/agent_framework/_workflows/_runner.py b/python/packages/core/agent_framework/_workflows/_runner.py index ac5558dbe1c..1948e12b81d 100644 --- a/python/packages/core/agent_framework/_workflows/_runner.py +++ b/python/packages/core/agent_framework/_workflows/_runner.py @@ -6,14 +6,14 @@ import warnings from collections import defaultdict from collections.abc import AsyncGenerator, Sequence -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast from ..exceptions import ( WorkflowCheckpointException, WorkflowConvergenceException, ) from ._checkpoint import CheckpointID, CheckpointStorage, WorkflowCheckpoint -from ._const import EXECUTOR_STATE_KEY +from ._const import EDGE_STATE_KEY, EXECUTOR_STATE_KEY from ._edge import EdgeGroup from ._edge_runner import EdgeRunner, create_edge_runner from ._events import WorkflowEvent @@ -70,6 +70,9 @@ def __init__( self._executors = executors self._edge_runners = [create_edge_runner(group, executors) for group in edge_groups] self._edge_runner_map = self._parse_edge_runners(self._edge_runners) + # Keyed by topology rather than EdgeGroup.id, which is a random UUID: a checkpoint has to + # restore onto a rebuilt instance of the same workflow definition. + self._edge_runner_state_keys = [runner.state_key for runner in self._edge_runners] self._ctx = ctx self._workflow_name = workflow_name self._graph_signature_hash = graph_signature_hash @@ -235,6 +238,7 @@ async def _prepare_checkpoint_state(self) -> None: state payload without necessarily writing to a checkpoint storage backend. """ await self._save_executor_states() + self._save_edge_runner_states() self._state.commit() async def create_checkpoint_if_enabled(self) -> None: @@ -325,6 +329,8 @@ async def restore_from_checkpoint( self._state.import_state(checkpoint.state) # Restore executor states using the restored state await self._restore_executor_states() + # Restore edge runner states, resetting any left over from an interrupted run + self._restore_edge_runner_states() # Apply the checkpoint to the context await self._ctx.apply_checkpoint(checkpoint) # Mark the runner as resumed @@ -382,6 +388,7 @@ async def restore_checkpoint(self, checkpoint: WorkflowCheckpoint) -> None: self._state.clear() self._state.import_state(checkpoint.state) await self._restore_executor_states() + self._restore_edge_runner_states() await self._ctx.apply_checkpoint(checkpoint) self._mark_resumed(checkpoint) except Exception as e: @@ -428,6 +435,40 @@ async def _restore_executor_states(self) -> None: except Exception as ex: # pragma: no cover - defensive raise WorkflowCheckpointException(f"Executor {executor_id} on_checkpoint_restore failed") from ex + def _save_edge_runner_states(self) -> None: + """Store edge runner delivery state, such as partially filled fan-in buffers, in state. + + A fan-in buffers messages until every source has produced one, and the runner context has + already drained those messages, so a checkpoint that omits them loses them on restore. + """ + edge_states: dict[str, dict[str, Any]] = {} + for key, runner in zip(self._edge_runner_state_keys, self._edge_runners, strict=True): + snapshot = runner.snapshot_state() + if snapshot is not None: + edge_states[key] = snapshot + + if edge_states: + self._state.set(EDGE_STATE_KEY, edge_states) + elif self._state.has(EDGE_STATE_KEY): + self._state.delete(EDGE_STATE_KEY) + + def _restore_edge_runner_states(self) -> None: + """Reset every edge runner and reapply the delivery state held by the restored checkpoint. + + Every runner is reset, including runners the checkpoint has no entry for, so that state + left behind by an interrupted run cannot survive into the restored run. + """ + stored: Any = self._state.get(EDGE_STATE_KEY, {}) + if not isinstance(stored, dict): + raise WorkflowCheckpointException("Edge states in shared state is not a dictionary. Unable to restore.") + edge_states = cast(dict[Any, Any], stored) + + for key, runner in zip(self._edge_runner_state_keys, self._edge_runners, strict=True): + state: Any = edge_states.get(key) + if state is not None and not isinstance(state, dict): + raise WorkflowCheckpointException(f"Edge state for {key} is not a dictionary. Unable to restore.") + runner.restore_state(cast("dict[str, Any] | None", state)) + def _parse_edge_runners(self, edge_runners: list[EdgeRunner]) -> dict[str, list[EdgeRunner]]: """Parse the edge runners of the workflow into a mapping where each source executor ID maps to its edge runners. diff --git a/python/packages/core/tests/workflow/test_checkpoint.py b/python/packages/core/tests/workflow/test_checkpoint.py index 5f3da78cd1d..be8de7c13b6 100644 --- a/python/packages/core/tests/workflow/test_checkpoint.py +++ b/python/packages/core/tests/workflow/test_checkpoint.py @@ -1080,6 +1080,84 @@ async def test_memory_checkpoint_storage_roundtrip_empty_collections(): assert loaded.pending_request_info_events == {} +async def test_workflow_resume_restores_partially_filled_fan_in_buffer(): + """A fan-in still waiting on a source must keep the messages it already holds across a resume. + + ``fast`` reaches the fan-in one superstep before ``slow``. At that boundary its message + lives only in the fan-in buffer - the runner context drained it on delivery - so a + checkpoint that omits the buffer strands the workflow with a fan-in that never fires. + """ + from typing_extensions import Never + + from agent_framework import WorkflowBuilder, WorkflowContext, handler + from agent_framework._workflows._const import EDGE_STATE_KEY + from agent_framework._workflows._executor import Executor + + class Dispatcher(Executor): + @handler + async def dispatch(self, message: str, ctx: WorkflowContext[str]) -> None: + await ctx.send_message(message) + + class Fast(Executor): + @handler + async def run(self, message: str, ctx: WorkflowContext[str]) -> None: + await ctx.send_message(f"{message}-fast") + + class Relay(Executor): + @handler + async def run(self, message: str, ctx: WorkflowContext[str]) -> None: + await ctx.send_message(message) + + class Slow(Executor): + @handler + async def run(self, message: str, ctx: WorkflowContext[str]) -> None: + await ctx.send_message(f"{message}-slow") + + class Joiner(Executor): + @handler + async def join(self, messages: list[str], ctx: WorkflowContext[Never, str]) -> None: # type: ignore[valid-type] + await ctx.yield_output(", ".join(sorted(messages))) + + storage = InMemoryCheckpointStorage() + + def _build_workflow() -> Any: + dispatcher = Dispatcher(id="dispatcher") + fast = Fast(id="fast") + relay = Relay(id="relay") + slow = Slow(id="slow") + joiner = Joiner(id="joiner") + return ( + WorkflowBuilder( + name="fan-in-resume-test", + max_iterations=10, + start_executor=dispatcher, + checkpoint_storage=storage, + ) + .add_edge(dispatcher, fast) + .add_edge(dispatcher, relay) + .add_edge(relay, slow) + .add_fan_in_edges([fast, slow], joiner) + .build() + ) + + workflow = _build_workflow() + workflow_name = workflow.name + outputs = (await workflow.run("seed")).get_outputs() + assert outputs == ["seed-fast, seed-slow"] + + checkpoints = await storage.list_checkpoints(workflow_name=workflow_name) + buffered = [checkpoint for checkpoint in checkpoints if EDGE_STATE_KEY in checkpoint.state] + assert len(buffered) == 1, ( + f"Exactly one superstep boundary should fall while the fan-in holds only the fast branch, got {len(buffered)}" + ) + + # Resume on a fresh instance of the same workflow definition, which has new edge group ids. + resumed_workflow = _build_workflow() + resumed_outputs = (await resumed_workflow.run(checkpoint_id=buffered[0].checkpoint_id)).get_outputs() + + assert resumed_outputs == ["seed-fast, seed-slow"] + + # endregion # region FileCheckpointStorage diff --git a/python/packages/core/tests/workflow/test_runner.py b/python/packages/core/tests/workflow/test_runner.py index 2299c382432..63d97052244 100644 --- a/python/packages/core/tests/workflow/test_runner.py +++ b/python/packages/core/tests/workflow/test_runner.py @@ -20,8 +20,8 @@ WorkflowRunState, handler, ) -from agent_framework._workflows._const import EXECUTOR_STATE_KEY -from agent_framework._workflows._edge import FanOutEdgeGroup, SingleEdgeGroup +from agent_framework._workflows._const import EDGE_STATE_KEY, EXECUTOR_STATE_KEY +from agent_framework._workflows._edge import FanInEdgeGroup, FanOutEdgeGroup, SingleEdgeGroup from agent_framework._workflows._runner import Runner from agent_framework._workflows._runner_context import ( InProcRunnerContext, @@ -566,6 +566,89 @@ async def on_checkpoint_restore(self, state: dict[str, Any]) -> None: assert runner._previous_checkpoint_id == checkpoint.checkpoint_id # pyright: ignore[reportPrivateUsage] +class CollectingExecutor(Executor): + """A fan-in target that records every aggregated batch it receives.""" + + def __init__(self, id: str) -> None: + super().__init__(id=id) + self.batches: list[list[int]] = [] + + @handler + async def collect(self, messages: list[MockMessage], ctx: WorkflowContext[Any, int]) -> None: + self.batches.append([message.data for message in messages]) + + +def _fan_in_runner(target: CollectingExecutor) -> tuple[Runner, InProcRunnerContext]: + """Build a runner for a two-source fan-in into ``target``.""" + source_a = MockExecutor(id="source_a") + source_b = MockExecutor(id="source_b") + edge_group = FanInEdgeGroup([source_a.id, source_b.id], target.id) + executors: dict[str, Executor] = {source_a.id: source_a, source_b.id: source_b, target.id: target} + ctx = InProcRunnerContext() + + return Runner([edge_group], executors, State(), ctx, "test_name", graph_signature_hash="test_hash"), ctx + + +async def test_runner_checkpoint_roundtrips_partially_filled_fan_in_buffer(): + """A fan-in holding messages from a subset of its sources must checkpoint and restore them. + + The runner context has already drained those messages, so the checkpoint is the only + place they still exist. A rebuilt workflow gets fresh ``EdgeGroup`` ids, so the restore + also has to find the buffer without relying on them. + """ + target = CollectingExecutor(id="target") + runner, ctx = _fan_in_runner(target) + + await ctx.send_message(WorkflowMessage(data=MockMessage(data=1), source_id="source_a")) + await runner._run_iteration() # pyright: ignore[reportPrivateUsage] + + # The fan-in is still waiting on source_b, so nothing has reached the target yet. + assert target.batches == [] + + checkpoint = await runner.build_checkpoint() + assert EDGE_STATE_KEY in checkpoint.state + + # Resume on a fresh instance of the same workflow definition. + restored_target = CollectingExecutor(id="target") + restored_runner, restored_ctx = _fan_in_runner(restored_target) + await restored_runner.restore_checkpoint(checkpoint) + + await restored_ctx.send_message(WorkflowMessage(data=MockMessage(data=2), source_id="source_b")) + await restored_runner._run_iteration() # pyright: ignore[reportPrivateUsage] + + assert restored_target.batches == [[1, 2]] + + +async def test_runner_restore_clears_fan_in_buffer_left_by_an_interrupted_run(): + """Messages buffered after the restored checkpoint must not survive the restore. + + Their sources are re-executed once the run resumes, so keeping them would deliver the + same message twice and could fire the fan-in before the resumed superstep produced all + of its messages. + """ + target = CollectingExecutor(id="target") + runner, ctx = _fan_in_runner(target) + + # Captured while the fan-in buffer is empty, so the checkpoint carries no edge state at all - + # the same shape as a checkpoint written before edge state was captured. + checkpoint = await runner.build_checkpoint() + assert EDGE_STATE_KEY not in checkpoint.state + + # A superstep that fails after the fan-in buffered source_a leaves that message behind. + await ctx.send_message(WorkflowMessage(data=MockMessage(data=1), source_id="source_a")) + await runner._run_iteration() # pyright: ignore[reportPrivateUsage] + assert target.batches == [] + + await runner.restore_checkpoint(checkpoint) + + # Both sources are re-executed after the resume; the target must see each message once. + await ctx.send_message(WorkflowMessage(data=MockMessage(data=1), source_id="source_a")) + await ctx.send_message(WorkflowMessage(data=MockMessage(data=2), source_id="source_b")) + await runner._run_iteration() # pyright: ignore[reportPrivateUsage] + + assert target.batches == [[1, 2]] + + async def test_runner_build_checkpoint_includes_in_flight_messages(): """build_checkpoint() must snapshot in-flight messages non-destructively.""" executor = MockExecutor(id="executor_a") diff --git a/python/packages/core/tests/workflow/test_sub_workflow.py b/python/packages/core/tests/workflow/test_sub_workflow.py index f320146e3f5..76756a56d76 100644 --- a/python/packages/core/tests/workflow/test_sub_workflow.py +++ b/python/packages/core/tests/workflow/test_sub_workflow.py @@ -731,6 +731,148 @@ async def test_sub_workflow_checkpoint_restore_preserves_sub_workflow_state() -> assert result.get_final_state() == WorkflowRunState.IDLE +@dataclass +class FanInResult: + """Result yielded by the fan-in sub-workflow.""" + + value: str + + +class FanInSubWorkflowFast(Executor): + """Sub-workflow branch that reaches the fan-in immediately.""" + + def __init__(self) -> None: + super().__init__(id="fan_in_fast") + + @handler + async def run(self, message: str, ctx: WorkflowContext[str]) -> None: + await ctx.send_message(f"{message}-fast") + + +class FanInSubWorkflowAsk(Executor): + """Sub-workflow branch that pauses on a request before reaching the fan-in.""" + + def __init__(self) -> None: + super().__init__(id="fan_in_ask") + + @handler + async def run(self, message: str, ctx: WorkflowContext) -> None: + await ctx.request_info(request_data=CheckpointRequest(prompt=message), response_type=str) + + @response_handler + async def handle_response( + self, + original_request: CheckpointRequest, + response: str, + ctx: WorkflowContext[str], + ) -> None: + await ctx.send_message(response) + + +class FanInSubWorkflowJoiner(Executor): + """Sub-workflow fan-in target.""" + + def __init__(self) -> None: + super().__init__(id="fan_in_joiner") + + @handler + async def join(self, messages: list[str], ctx: WorkflowContext[Never, FanInResult]) -> None: # type: ignore[valid-type] + await ctx.yield_output(FanInResult(value=", ".join(sorted(messages)))) + + +class FanInCheckpointCoordinator(Executor): + """Parent coordinator that forwards the sub-workflow's request and yields its result.""" + + def __init__(self) -> None: + super().__init__(id="fan_in_coordinator") + self._pending_requests: dict[str, SubWorkflowRequestMessage] = {} + + @handler + async def start(self, value: str, ctx: WorkflowContext[str]) -> None: + await ctx.send_message(value) + + @handler + async def handle_sub_workflow_request(self, request: SubWorkflowRequestMessage, ctx: WorkflowContext) -> None: + data = request.source_event.data + if isinstance(data, CheckpointRequest): + self._pending_requests[data.id] = request + await ctx.request_info(data, str) + + @handler + async def handle_sub_workflow_result(self, result: FanInResult, ctx: WorkflowContext[Never, str]) -> None: # type: ignore[valid-type] + await ctx.yield_output(result.value) + + @response_handler + async def handle_response( + self, + original_request: CheckpointRequest, + response: str, + ctx: WorkflowContext[SubWorkflowResponseMessage], + ) -> None: + sub_request = self._pending_requests.pop(original_request.id, None) + if sub_request is None: + raise ValueError(f"No pending request for ID: {original_request.id}") + await ctx.send_message(sub_request.create_response(response)) + + async def on_checkpoint_save(self) -> dict[str, Any]: + return {"pending_requests": self._pending_requests} + + async def on_checkpoint_restore(self, state: dict[str, Any]) -> None: + self._pending_requests = state.get("pending_requests", {}) + + +def _build_fan_in_sub_workflow(storage: InMemoryCheckpointStorage) -> Workflow: + """Build a parent workflow whose sub-workflow pauses with a half-filled fan-in.""" + start = FanInSubWorkflowFast() + ask = FanInSubWorkflowAsk() + joiner = FanInSubWorkflowJoiner() + sub_workflow = ( + WorkflowBuilder(start_executor=start).add_edge(start, ask).add_fan_in_edges([start, ask], joiner).build() + ) + sub_workflow_executor = WorkflowExecutor(sub_workflow, id="fan_in_sub_workflow_executor") + + coordinator = FanInCheckpointCoordinator() + return ( + WorkflowBuilder(start_executor=coordinator, checkpoint_storage=storage) + .add_edge(coordinator, sub_workflow_executor) + .add_edge(sub_workflow_executor, coordinator) + .build() + ) + + +async def test_sub_workflow_checkpoint_restore_preserves_partially_filled_fan_in() -> None: + """A sub-workflow that pauses with a half-filled fan-in must keep those messages on resume. + + The sub-workflow's own checkpoint is embedded in the parent checkpoint, so the fast + branch's message - already drained out of the sub-workflow's runner context and into the + fan-in buffer - only survives if the buffer travels with that checkpoint. + """ + storage = InMemoryCheckpointStorage() + + workflow1 = _build_fan_in_sub_workflow(storage) + request_id: str | None = None + async for event in workflow1.run("seed", stream=True): + if event.type == "request_info": + request_id = event.request_id + assert request_id is not None + + checkpoints = await storage.list_checkpoints(workflow_name=workflow1.name) + checkpoint_id = max(checkpoints, key=lambda cp: cp.iteration_count).checkpoint_id + + # Resume on a fresh instance, whose sub-workflow has new edge group ids, then answer the + # request so the second branch finally reaches the fan-in. + workflow2 = _build_fan_in_sub_workflow(storage) + resumed_request_id: str | None = None + async for event in workflow2.run(checkpoint_id=checkpoint_id, stream=True): + if event.type == "request_info": + resumed_request_id = event.request_id + assert resumed_request_id is not None + + result = await workflow2.run(responses={resumed_request_id: "answered"}) + + assert result.get_outputs() == ["answered, seed-fast"] + + async def test_sub_workflow_intermediate_outputs_propagate_to_parent() -> None: """A child workflow's intermediate emissions must bubble up through the parent. From b5c6a5d53f5beb4b9a6f97ca1566111b506bb5c2 Mon Sep 17 00:00:00 2001 From: yashvanthange Date: Sun, 30 Aug 2026 00:33:01 +0530 Subject: [PATCH 2/5] Python: encode edge state keys structurally Executor ids are only required to be non-empty, so joining source and target ids with "->" and "," let two different fan-in groups produce the same key: sources ['a', 'b->t,c'] and ['a->t,b', 'c'] into the same target collide. One group's buffer snapshot would then overwrite the other's, and the restore would hand both runners the same messages. JSON-encode the sorted pairs instead. --- .../agent_framework/_workflows/_edge_runner.py | 9 ++++++--- .../packages/core/tests/workflow/test_runner.py | 16 ++++++++++++++++ 2 files changed, 22 insertions(+), 3 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_edge_runner.py b/python/packages/core/agent_framework/_workflows/_edge_runner.py index 0bf83f53b5d..d3994875287 100644 --- a/python/packages/core/agent_framework/_workflows/_edge_runner.py +++ b/python/packages/core/agent_framework/_workflows/_edge_runner.py @@ -1,6 +1,7 @@ # Copyright (c) Microsoft. All rights reserved. import asyncio +import json import logging from abc import ABC, abstractmethod from collections import defaultdict @@ -66,10 +67,12 @@ def state_key(self) -> str: same workflow definition and cannot key state that has to survive a rebuild. Checkpoint compatibility is decided by graph topology, so the key is derived from topology too. Builder validation rejects two edges that share a ``source -> target`` pair anywhere in - the workflow, so no two edge groups can produce the same key. + the workflow, so no two edge groups can produce the same key. Executor ids are only + required to be non-empty, so the pairs are JSON-encoded rather than joined with + separators an id could itself contain. """ - edges = sorted(f"{edge.source_id}->{edge.target_id}" for edge in self._edge_group.edges) - return f"{self._edge_group.__class__.__name__}:{','.join(edges)}" + edges = sorted([edge.source_id, edge.target_id] for edge in self._edge_group.edges) + return f"{self._edge_group.__class__.__name__}:{json.dumps(edges, separators=(',', ':'))}" def snapshot_state(self) -> dict[str, Any] | None: """Capture in-flight delivery state so a checkpoint can restore it. diff --git a/python/packages/core/tests/workflow/test_runner.py b/python/packages/core/tests/workflow/test_runner.py index 63d97052244..e5f7e340f6a 100644 --- a/python/packages/core/tests/workflow/test_runner.py +++ b/python/packages/core/tests/workflow/test_runner.py @@ -22,6 +22,7 @@ ) from agent_framework._workflows._const import EDGE_STATE_KEY, EXECUTOR_STATE_KEY from agent_framework._workflows._edge import FanInEdgeGroup, FanOutEdgeGroup, SingleEdgeGroup +from agent_framework._workflows._edge_runner import FanInEdgeRunner from agent_framework._workflows._runner import Runner from agent_framework._workflows._runner_context import ( InProcRunnerContext, @@ -589,6 +590,21 @@ def _fan_in_runner(target: CollectingExecutor) -> tuple[Runner, InProcRunnerCont return Runner([edge_group], executors, State(), ctx, "test_name", graph_signature_hash="test_hash"), ctx +def test_edge_runner_state_key_distinguishes_ids_containing_separators(): + """Executor ids are only required to be non-empty, so the key cannot join them with separators. + + Both groups below hold three edges into ``t`` and differ only in where the source ids place + ``->`` and ``,``. Sharing a key would make one group's buffer overwrite the other's. + """ + group_a = FanInEdgeGroup(["a", "b->t,c"], "t") + group_b = FanInEdgeGroup(["a->t,b", "c"], "t") + + key_a = FanInEdgeRunner(group_a, {}).state_key + key_b = FanInEdgeRunner(group_b, {}).state_key + + assert key_a != key_b + + async def test_runner_checkpoint_roundtrips_partially_filled_fan_in_buffer(): """A fan-in holding messages from a subset of its sources must checkpoint and restore them. From 285595713060d8711b6741497d4619abebf40420 Mon Sep 17 00:00:00 2001 From: yashvanthange Date: Mon, 31 Aug 2026 10:25:32 +0530 Subject: [PATCH 3/5] Python: cancel sibling deliveries when one fails mid-superstep @moonbox3 on PR #7948: _run_iteration's asyncio.gather calls did not cancel their other tasks when one raised, so a still-in-flight fan-in delivery could resume after restore_checkpoint/restore_from_checkpoint had already reset the FanInEdgeRunner's buffer on the same runner instance, appending a message from the failed, never-checkpointed superstep into the freshly restored buffer. Confirmed the race is real before fixing it: reproduced it with a test that fails a sibling source's delivery while another source's fan-in delivery is parked mid-flight, restores from a checkpoint taken before either message was sent, then releases the parked delivery. Without the fix it appends the stale message into the just-restored buffer, and the fan-in later fires using that stale value instead of the fresh resumed delivery. Added _gather_cancelling_siblings_on_error to replace both gather() call sites in _run_iteration, so a failing delivery cancels and awaits every other one before the exception propagates. --- .../agent_framework/_workflows/_runner.py | 27 ++++- .../core/tests/workflow/test_runner.py | 100 +++++++++++++++++- 2 files changed, 123 insertions(+), 4 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_runner.py b/python/packages/core/agent_framework/_workflows/_runner.py index 1948e12b81d..2a5fe0e05d4 100644 --- a/python/packages/core/agent_framework/_workflows/_runner.py +++ b/python/packages/core/agent_framework/_workflows/_runner.py @@ -5,7 +5,7 @@ import logging import warnings from collections import defaultdict -from collections.abc import AsyncGenerator, Sequence +from collections.abc import AsyncGenerator, Coroutine, Sequence from typing import TYPE_CHECKING, Any, cast from ..exceptions import ( @@ -27,6 +27,27 @@ logger = logging.getLogger(__name__) +async def _gather_cancelling_siblings_on_error(*coroutines: Coroutine[Any, Any, Any]) -> None: + """Run coroutines concurrently; on any failure, cancel and await every other one before raising. + + Plain ``asyncio.gather()`` does not cancel its other tasks when one raises - by default they keep + running as orphaned background tasks even though the caller has already moved on with the raised + exception. For fan-in edge delivery this is a real race with checkpoint restoration: a delivery + that is still in-flight when a sibling delivery fails can append into a ``FanInEdgeRunner``'s + buffer *after* ``restore_checkpoint``/``restore_from_checkpoint`` has already cleared it on the + same runner instance, corrupting the freshly restored state and producing a duplicate-aggregated + fan-in batch once the resumed superstep redelivers the same source. + """ + tasks = [asyncio.ensure_future(coro) for coro in coroutines] + try: + await asyncio.gather(*tasks) + except BaseException: + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + raise + + def warn_runner_deprecated() -> None: """Emit a deprecation warning when ``Runner`` is accessed from the public API. @@ -222,14 +243,14 @@ async def _deliver_messages_for_edge_runner(edge_runner: EdgeRunner) -> None: return tasks = [_deliver_messages_for_edge_runner(edge_runner) for edge_runner in associated_edge_runners] - await asyncio.gather(*tasks) + await _gather_cancelling_siblings_on_error(*tasks) message_batches = await self._ctx.drain_messages() tasks = [ _deliver_messages(source_executor_id, source_messages) for source_executor_id, source_messages in message_batches.items() ] - await asyncio.gather(*tasks) + await _gather_cancelling_siblings_on_error(*tasks) async def _prepare_checkpoint_state(self) -> None: """Persist executor snapshots into committed shared state. diff --git a/python/packages/core/tests/workflow/test_runner.py b/python/packages/core/tests/workflow/test_runner.py index e5f7e340f6a..16ef867d93c 100644 --- a/python/packages/core/tests/workflow/test_runner.py +++ b/python/packages/core/tests/workflow/test_runner.py @@ -2,7 +2,7 @@ import asyncio from dataclasses import dataclass -from typing import Any +from typing import Any, cast from unittest.mock import AsyncMock, MagicMock import pytest @@ -665,6 +665,104 @@ async def test_runner_restore_clears_fan_in_buffer_left_by_an_interrupted_run(): assert target.batches == [[1, 2]] +async def test_runner_orphaned_delivery_cannot_repopulate_a_restored_fan_in_buffer(): + """A sibling delivery still in flight when another source fails must not survive into a restore. + + Reported on PR #7953 (@moonbox3): ``_run_iteration``'s ``asyncio.gather`` calls do not cancel their + other tasks when one raises - the others keep running as orphaned background tasks. If a fan-in + delivery is one of those orphans, it can resume after ``restore_checkpoint`` has already reset the + ``FanInEdgeRunner``'s buffer for the *next* run and append into it, so the buffer ends up holding a + message from the failed, never-checkpointed superstep. Once the caller redelivers that same source + as part of the normal resume flow, the fan-in can fire using the stale message instead of - or + alongside - the fresh redelivery. + + The pre-restore (stale) and post-restore (fresh) ``source_b`` messages deliberately carry + different payloads (``99`` vs ``2``) so a fan-in that fires on the stale message is + distinguishable from one that correctly waits for the fresh redelivery; using the same payload + for both would let a buggy run and a correct run produce an identical-looking batch by + coincidence. + """ + target = CollectingExecutor(id="target") + source_a = MockExecutor(id="source_a") + source_b = MockExecutor(id="source_b") + source_c = MockExecutor(id="source_c") + fan_in_group = FanInEdgeGroup([source_a.id, source_b.id], target.id) + executors: dict[str, Executor] = { + source_a.id: source_a, + source_b.id: source_b, + source_c.id: source_c, + target.id: target, + } + ctx = InProcRunnerContext() + runner = Runner([fan_in_group], executors, State(), ctx, "test_name", graph_signature_hash="test_hash") + + real_fan_in_runner = cast(FanInEdgeRunner, runner._edge_runner_map["source_a"][0]) # pyright: ignore[reportPrivateUsage] + assert runner._edge_runner_map["source_b"][0] is real_fan_in_runner # pyright: ignore[reportPrivateUsage] + + entered_delay = asyncio.Event() + release = asyncio.Event() + + class DelayedFanInDelivery: + """Wraps the real fan-in runner's delivery for source_b, pausing before it actually appends.""" + + async def send_message(self, message: WorkflowMessage, state: State, ctx: RunnerContext) -> bool: + entered_delay.set() + await release.wait() + return await real_fan_in_runner.send_message(message, state, ctx) + + class FailingDelivery: + async def send_message(self, message: WorkflowMessage, state: State, ctx: RunnerContext) -> bool: + raise RuntimeError("source_c delivery failed") + + runner._edge_runner_map["source_b"] = [DelayedFanInDelivery()] # type: ignore[assignment, list-item] # pyright: ignore[reportPrivateUsage] # ty: ignore[invalid-assignment] + runner._edge_runner_map["source_c"] = [FailingDelivery()] # type: ignore[assignment, list-item] # pyright: ignore[reportPrivateUsage] # ty: ignore[invalid-assignment] + + # A checkpoint from before this superstep started: empty fan-in buffer, matching the last + # known-good superstep boundary that a real resume flow would restore to. + checkpoint = await runner.build_checkpoint() + assert EDGE_STATE_KEY not in checkpoint.state + + await ctx.send_message(WorkflowMessage(data=MockMessage(data=1), source_id="source_a")) + await ctx.send_message(WorkflowMessage(data=MockMessage(data=99), source_id="source_b")) + await ctx.send_message(WorkflowMessage(data=MockMessage(data=-1), source_id="source_c")) + + iteration_task = asyncio.create_task(runner._run_iteration()) # pyright: ignore[reportPrivateUsage] + + # source_a's delivery completes immediately; source_b's is parked mid-flight, about to append + # into the *same* FanInEdgeRunner instance, when source_c's delivery fails. + await entered_delay.wait() + + with pytest.raises(RuntimeError, match="source_c delivery failed"): + await iteration_task + + # The app's resume flow restores the same runner from the last good checkpoint - this is exactly + # what a caller does after catching the failure above, per Workflow._execute_with_message_or_checkpoint. + await runner.restore_checkpoint(checkpoint) + assert real_fan_in_runner._buffer == {} # pyright: ignore[reportPrivateUsage] + + # Release the delayed delivery. If _run_iteration cancelled it when source_c failed (the fix), + # `release.wait()` raises CancelledError here and the real fan-in runner's send_message for the + # stale data=99 message is never reached at all; releasing an already-cancelled waiter is a no-op. + release.set() + await asyncio.sleep(0) # yield so the orphaned task either appends or finishes cancelling + + # The stale data=99 message from the failed, never-checkpointed superstep must not have reached + # the buffer the restore just reset for the next run. + assert real_fan_in_runner._buffer.get("source_b", []) == [] # pyright: ignore[reportPrivateUsage] + + # The resume flow redelivers the failed superstep's sources with the correct payloads, as the + # existing test_runner_restore_clears_fan_in_buffer_left_by_an_interrupted_run models. source_b's + # fresh message (data=2) differs from the stale one (data=99) precisely so the assertion below can + # tell which one the fan-in actually used. + await ctx.send_message(WorkflowMessage(data=MockMessage(data=1), source_id="source_a")) + await ctx.send_message(WorkflowMessage(data=MockMessage(data=2), source_id="source_b")) + await runner._run_iteration() # pyright: ignore[reportPrivateUsage] + + # The target must see the fresh redelivery (source_a=1, source_b=2), not the stale pre-restore + # message (99) that the orphaned delivery left behind, and not both. + assert target.batches == [[1, 2]] + + async def test_runner_build_checkpoint_includes_in_flight_messages(): """build_checkpoint() must snapshot in-flight messages non-destructively.""" executor = MockExecutor(id="executor_a") From da23d14afb129f6587816559594dc212f6558874 Mon Sep 17 00:00:00 2001 From: yashvanthange Date: Mon, 31 Aug 2026 14:39:01 +0530 Subject: [PATCH 4/5] Python: cancel fan-out sibling deliveries too, not just _run_iteration's @moonbox3 follow-up on PR #7948: the previous fix wrapped _run_iteration's two gather() sites (across sources, and across edge runners for one source), but FanOutEdgeRunner.send_message has its own separate asyncio.gather() across one message's fan-out targets. When a fan-out has a single edge runner for its source - the common case - that inner gather is the only level where one target's failure has a sibling to race against, so the outer fix never reached it. Confirmed the race is real before fixing it: reproduced a fan-out target still executing when its sibling fails, calling ctx.send_message after restore_checkpoint has already cleared RunnerContext._messages, landing its stale output in the queue the restore reset for the resumed run. The test fails on the pre-fix code and passes after. Moved the cancel-and-await helper (renamed gather_cancelling_siblings_on_error, no longer private) into _edge_runner.py, since _runner.py already imports from there and the reverse would be circular, and used it at the fan-out gather site too. --- .../_workflows/_edge_runner.py | 36 ++++++++-- .../agent_framework/_workflows/_runner.py | 29 ++------ .../core/tests/workflow/test_runner.py | 71 +++++++++++++++++++ 3 files changed, 105 insertions(+), 31 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_edge_runner.py b/python/packages/core/agent_framework/_workflows/_edge_runner.py index d3994875287..ec8820886dc 100644 --- a/python/packages/core/agent_framework/_workflows/_edge_runner.py +++ b/python/packages/core/agent_framework/_workflows/_edge_runner.py @@ -5,7 +5,7 @@ import logging from abc import ABC, abstractmethod from collections import defaultdict -from collections.abc import Callable +from collections.abc import Callable, Coroutine from typing import Any, cast from ..exceptions import WorkflowCheckpointException @@ -26,6 +26,32 @@ logger = logging.getLogger(__name__) +async def gather_cancelling_siblings_on_error(*coroutines: Coroutine[Any, Any, Any]) -> None: + """Run coroutines concurrently; on any failure, cancel and await every other one before raising. + + Plain ``asyncio.gather()`` does not cancel its other tasks when one raises - by default they keep + running as orphaned background tasks even though the caller has already moved on with the raised + exception. For edge delivery this is a real race with checkpoint restoration: work that is still + in-flight when a sibling fails can mutate runner/context state - a ``FanInEdgeRunner``'s buffer, or + a ``RunnerContext``'s pending-message queue via a target executor's own ``ctx.send_message`` call - + *after* ``restore_checkpoint``/``restore_from_checkpoint`` has already reset that same state for the + resumed run, corrupting it with output from the failed, never-checkpointed superstep. + + Shared by :class:`RunnerImpl._run_iteration` (across sources and across edge runners for a source) + and :class:`FanOutEdgeRunner.send_message` (across one message's fan-out targets) - every point in + the delivery path where sibling coroutines run concurrently and one can fail while another is still + executing. + """ + tasks = [asyncio.ensure_future(coro) for coro in coroutines] + try: + await asyncio.gather(*tasks) + except BaseException: + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + raise + + class EdgeRunner(ABC): """Abstract base class for edge runners that handle message delivery.""" @@ -314,13 +340,11 @@ async def send_message( if deliverable_edges: - async def send_to_edge(edge: Edge) -> bool: + async def send_to_edge(edge: Edge) -> None: await self._execute_on_target(edge.target_id, [edge.source_id], message, state, ctx) - return True - tasks = [send_to_edge(edge) for edge in deliverable_edges] - results = await asyncio.gather(*tasks) - return any(results) + await gather_cancelling_siblings_on_error(*(send_to_edge(edge) for edge in deliverable_edges)) + return True # If we get here, it's a broadcast message with no deliverable edges return False diff --git a/python/packages/core/agent_framework/_workflows/_runner.py b/python/packages/core/agent_framework/_workflows/_runner.py index 2a5fe0e05d4..ba8fdf188ca 100644 --- a/python/packages/core/agent_framework/_workflows/_runner.py +++ b/python/packages/core/agent_framework/_workflows/_runner.py @@ -5,7 +5,7 @@ import logging import warnings from collections import defaultdict -from collections.abc import AsyncGenerator, Coroutine, Sequence +from collections.abc import AsyncGenerator, Sequence from typing import TYPE_CHECKING, Any, cast from ..exceptions import ( @@ -15,7 +15,7 @@ from ._checkpoint import CheckpointID, CheckpointStorage, WorkflowCheckpoint from ._const import EDGE_STATE_KEY, EXECUTOR_STATE_KEY from ._edge import EdgeGroup -from ._edge_runner import EdgeRunner, create_edge_runner +from ._edge_runner import EdgeRunner, create_edge_runner, gather_cancelling_siblings_on_error from ._events import WorkflowEvent from ._executor import Executor from ._runner_context import ( @@ -27,27 +27,6 @@ logger = logging.getLogger(__name__) -async def _gather_cancelling_siblings_on_error(*coroutines: Coroutine[Any, Any, Any]) -> None: - """Run coroutines concurrently; on any failure, cancel and await every other one before raising. - - Plain ``asyncio.gather()`` does not cancel its other tasks when one raises - by default they keep - running as orphaned background tasks even though the caller has already moved on with the raised - exception. For fan-in edge delivery this is a real race with checkpoint restoration: a delivery - that is still in-flight when a sibling delivery fails can append into a ``FanInEdgeRunner``'s - buffer *after* ``restore_checkpoint``/``restore_from_checkpoint`` has already cleared it on the - same runner instance, corrupting the freshly restored state and producing a duplicate-aggregated - fan-in batch once the resumed superstep redelivers the same source. - """ - tasks = [asyncio.ensure_future(coro) for coro in coroutines] - try: - await asyncio.gather(*tasks) - except BaseException: - for task in tasks: - task.cancel() - await asyncio.gather(*tasks, return_exceptions=True) - raise - - def warn_runner_deprecated() -> None: """Emit a deprecation warning when ``Runner`` is accessed from the public API. @@ -243,14 +222,14 @@ async def _deliver_messages_for_edge_runner(edge_runner: EdgeRunner) -> None: return tasks = [_deliver_messages_for_edge_runner(edge_runner) for edge_runner in associated_edge_runners] - await _gather_cancelling_siblings_on_error(*tasks) + await gather_cancelling_siblings_on_error(*tasks) message_batches = await self._ctx.drain_messages() tasks = [ _deliver_messages(source_executor_id, source_messages) for source_executor_id, source_messages in message_batches.items() ] - await _gather_cancelling_siblings_on_error(*tasks) + await gather_cancelling_siblings_on_error(*tasks) async def _prepare_checkpoint_state(self) -> None: """Persist executor snapshots into committed shared state. diff --git a/python/packages/core/tests/workflow/test_runner.py b/python/packages/core/tests/workflow/test_runner.py index 16ef867d93c..b6f9bc531ef 100644 --- a/python/packages/core/tests/workflow/test_runner.py +++ b/python/packages/core/tests/workflow/test_runner.py @@ -763,6 +763,77 @@ async def send_message(self, message: WorkflowMessage, state: State, ctx: Runner assert target.batches == [[1, 2]] +async def test_runner_orphaned_fan_out_target_cannot_repopulate_a_restored_message_queue(): + """The same orphaned-task race applies one level deeper, inside FanOutEdgeRunner.send_message. + + Follow-up from @moonbox3 on PR #7948: ``_gather_cancelling_siblings_on_error`` wraps the two + ``gather()`` call sites in ``_run_iteration``, but ``FanOutEdgeRunner.send_message`` has its own, + separate ``asyncio.gather()`` across a single message's fan-out targets (``_edge_runner.py:322``). + When a fan-out has only one edge runner for its source (the common case), that inner gather is + the *only* level at which one target's failure has a sibling target to race against - the outer + per-source/per-edge-runner levels my fix wraps see just one task each, so there is nothing for + the fix to cancel there. A target still executing when a sibling target fails can call + ``WorkflowContext.send_message`` (via its own handler) after ``restore_checkpoint`` has already + cleared ``RunnerContext._messages`` for the resumed run, repopulating it with output from the + failed, never-checkpointed superstep. + """ + entered_delay = asyncio.Event() + release = asyncio.Event() + + class FailingTarget(Executor): + @handler + async def handle(self, message: MockMessage, ctx: WorkflowContext[Any, int]) -> None: + raise RuntimeError("target1 failed") + + class BlockingTarget(Executor): + """Still executing when its fan-out sibling fails; emits a message once released.""" + + @handler + async def handle(self, message: MockMessage, ctx: WorkflowContext[Any, int]) -> None: + entered_delay.set() + await release.wait() + # A stale output from the failed superstep, sent after the sibling's failure surfaced. + await ctx.send_message(MockMessage(data=999)) + + source = MockExecutor(id="source") + target1 = FailingTarget(id="target1") + target2 = BlockingTarget(id="target2") + edge_group = FanOutEdgeGroup(source_id=source.id, target_ids=[target1.id, target2.id]) + executors: dict[str, Executor] = {source.id: source, target1.id: target1, target2.id: target2} + ctx = InProcRunnerContext() + runner = Runner([edge_group], executors, State(), ctx, "test_name", graph_signature_hash="test_hash") + + # A checkpoint from before this superstep started: no in-flight messages, matching the last + # known-good superstep boundary a real resume flow would restore to. + checkpoint = await runner.build_checkpoint() + assert not checkpoint.messages + + await ctx.send_message(WorkflowMessage(data=MockMessage(data=1), source_id=source.id)) + + iteration_task = asyncio.create_task(runner._run_iteration()) # pyright: ignore[reportPrivateUsage] + + # target2 is now parked mid-handler, about to call ctx.send_message once released, when target1 + # raises inside the same FanOutEdgeRunner.send_message call's own internal gather. + await entered_delay.wait() + + with pytest.raises(RuntimeError, match="target1 failed"): + await iteration_task + + # The app's resume flow restores the same runner from the last good checkpoint. + await runner.restore_checkpoint(checkpoint) + assert not await ctx.has_messages() + + # Release the parked target. If _run_iteration's fix reached this inner gather, target2's task + # would already be cancelled and this is a no-op; if it does not, target2 proceeds to call + # ctx.send_message with its stale output. + release.set() + await asyncio.sleep(0) # yield so the orphaned target's send_message actually runs before we continue + + # The message queue the restore just reset for the next run must not have picked up output from + # the failed, never-checkpointed superstep. + assert not await ctx.has_messages() + + async def test_runner_build_checkpoint_includes_in_flight_messages(): """build_checkpoint() must snapshot in-flight messages non-destructively.""" executor = MockExecutor(id="executor_a") From c7453aff49185923f8c90703d6c2e4b1877b045d Mon Sep 17 00:00:00 2001 From: yashvanthange Date: Mon, 31 Aug 2026 15:27:58 +0530 Subject: [PATCH 5/5] Python: cancel WorkflowExecutor's output-forwarding siblings too Retroactive check on our own claim, per mistake #12: this PR's review reply said the other gather() sites in _workflow.py/_workflow_executor.py "hold no cross-invocation buffered state a restore can race with." Tested rather than re-argued it. _workflow.py:1034 (send_request_info_response) survives: it pops a still- pending event before writing anything, and a restore clears pending events too, so an orphaned write fails closed (raises) instead of landing. Reproduced directly against InProcRunnerContext to confirm. _workflow_executor.py's four gather() sites do not survive. ctx.send_message and ctx.yield_output/add_event write into RunnerContext state (_messages, the event queue) with no equivalent precondition check. Reproduced the same shape as the fan-in/fan-out bugs: a sub-workflow output still being forwarded when a sibling output's send fails calls ctx.send_message after a checkpoint restore has already cleared the parent's message queue, landing stale output in it. Fixed all four sites with the same gather_cancelling_siblings_on_error helper already used in _run_iteration and FanOutEdgeRunner - moved to _edge_runner.py in the prior commit specifically so _workflow_executor.py could import it without a cycle. --- .../_workflows/_workflow_executor.py | 20 +++-- .../core/tests/workflow/test_sub_workflow.py | 78 +++++++++++++++++++ 2 files changed, 90 insertions(+), 8 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_workflow_executor.py b/python/packages/core/agent_framework/_workflows/_workflow_executor.py index 1a8f988d19b..774a7ce02b3 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow_executor.py +++ b/python/packages/core/agent_framework/_workflows/_workflow_executor.py @@ -1,6 +1,5 @@ # Copyright (c) Microsoft. All rights reserved. -import asyncio import logging import sys import types @@ -11,6 +10,7 @@ from ._workflow import Workflow from ._const import GLOBAL_KWARGS_KEY, WORKFLOW_RUN_KWARGS_KEY +from ._edge_runner import gather_cancelling_siblings_on_error from ._events import ( WorkflowEvent, WorkflowRunState, @@ -490,10 +490,12 @@ async def on_checkpoint_restore(self, state: dict[str, Any]) -> None: self.id, execution_context.execution_id, ) - await asyncio.gather(*[ - self.workflow._runner_context.add_request_info_event(event) # pyright: ignore[reportPrivateUsage] - for event in request_info_events - ]) + await gather_cancelling_siblings_on_error( + *( + self.workflow._runner_context.add_request_info_event(event) # pyright: ignore[reportPrivateUsage] + for event in request_info_events + ) + ) async def _process_workflow_result( self, @@ -524,9 +526,9 @@ async def _process_workflow_result( # Process outputs if self.allow_direct_output: # Note that the executor is allowed to continue its own execution after yielding outputs. - await asyncio.gather(*[ctx.yield_output(output) for output in outputs]) + await gather_cancelling_siblings_on_error(*(ctx.yield_output(output) for output in outputs)) else: - await asyncio.gather(*[ctx.send_message(output) for output in outputs]) + await gather_cancelling_siblings_on_error(*(ctx.send_message(output) for output in outputs)) # Pipe sub-workflow intermediate emissions up through the parent's event stream. # Bypasses the parent's yield-output classifier so the 'intermediate' label is preserved @@ -539,7 +541,9 @@ async def _forward_intermediate_output(output: Any) -> None: event = WorkflowEvent("intermediate", executor_id=self.id, data=output) await ctx.add_event(event) - await asyncio.gather(*[_forward_intermediate_output(output) for output in intermediate_outputs]) + await gather_cancelling_siblings_on_error( + *(_forward_intermediate_output(output) for output in intermediate_outputs) + ) # Process request info events for event in request_info_events: diff --git a/python/packages/core/tests/workflow/test_sub_workflow.py b/python/packages/core/tests/workflow/test_sub_workflow.py index 76756a56d76..f4ca5b38190 100644 --- a/python/packages/core/tests/workflow/test_sub_workflow.py +++ b/python/packages/core/tests/workflow/test_sub_workflow.py @@ -1,5 +1,6 @@ # Copyright (c) Microsoft. All rights reserved. +import asyncio import logging from dataclasses import dataclass, field from typing import Any @@ -14,14 +15,18 @@ SubWorkflowResponseMessage, Workflow, WorkflowBuilder, + WorkflowCheckpoint, WorkflowContext, WorkflowEvent, WorkflowExecutor, + WorkflowRunResult, WorkflowRunState, handler, response_handler, ) from agent_framework._workflows._checkpoint import InMemoryCheckpointStorage +from agent_framework._workflows._runner_context import InProcRunnerContext, WorkflowMessage +from agent_framework._workflows._state import State # Test message types @@ -943,3 +948,76 @@ async def run(self, message: str, ctx: WorkflowContext[Never, str]) -> None: # # The parent's own terminal output is unaffected. assert any(e.executor_id == "parent_sink" and e.data == "final: hello" for e in output_events) + + +async def test_workflow_executor_orphaned_output_forward_cannot_repopulate_a_restored_message_queue() -> None: + """The orphaned-task race from PR #7948 applies to WorkflowExecutor's own output-forwarding gather too. + + That PR's review excluded ``_workflow_executor.py``'s ``gather()`` sites with "none of them hold + cross-invocation buffered state a restore can race with" - that reasoning was wrong. When + ``allow_direct_output`` is false (the default), ``_process_workflow_result`` forwards every + sub-workflow output to the parent by gathering ``ctx.send_message(output)`` per output. + ``WorkflowContext.send_message`` reaches straight into ``RunnerContext._messages`` with no + precondition check (unlike, say, ``send_request_info_response``'s pop-and-validate, which turned + out to protect a sibling gather site in ``_workflow.py``). A sibling output still being sent when + another fails is therefore not cancelled by plain ``asyncio.gather``, and can append into the + parent's message queue after a checkpoint restore has already cleared it for the resumed run. + """ + + class _Inner(Executor): + @handler + async def handle(self, message: str, ctx: WorkflowContext[str]) -> None: + pass + + inner_workflow = WorkflowBuilder(start_executor=_Inner(id="inner")).build() + workflow_executor = WorkflowExecutor(inner_workflow, id="wrapped") + + ctx_impl = InProcRunnerContext() + state = State() + wctx: WorkflowContext[Any] = WorkflowContext(workflow_executor, ["parent_source"], state, ctx_impl) + + entered_delay = asyncio.Event() + release = asyncio.Event() + real_send_message = ctx_impl.send_message + + async def patched_send_message(message: WorkflowMessage) -> None: + if message.data == "bad-output": + raise RuntimeError("output send failed") + if message.data == "stale-output-from-failed-run": + entered_delay.set() + await release.wait() + await real_send_message(message) + + ctx_impl.send_message = patched_send_message # type: ignore[method-assign] # ty: ignore[invalid-assignment] + + # allow_direct_output defaults to False, so outputs go through the ctx.send_message branch. + # A status event is required for get_final_state(), which _process_workflow_result reads before + # the gather; without it the method raises immediately and entered_delay is never set. + result = WorkflowRunResult( + [ + WorkflowEvent("output", data="bad-output"), + WorkflowEvent("output", data="stale-output-from-failed-run"), + ], + status_events=[WorkflowEvent.status(WorkflowRunState.IDLE)], + ) + + task = asyncio.create_task(workflow_executor._process_workflow_result(result, wctx)) # pyright: ignore[reportPrivateUsage] + + # The second output's send is now parked mid-flight when the first one raises. + await asyncio.wait_for(entered_delay.wait(), timeout=5) + + with pytest.raises(RuntimeError, match="output send failed"): + await task + + # The parent restores the same runner context from its last good checkpoint - what a caller does + # after catching a failed superstep, per Workflow._execute_with_message_or_checkpoint. + checkpoint = WorkflowCheckpoint(workflow_name="test_name", graph_signature_hash="test_hash") + await ctx_impl.apply_checkpoint(checkpoint) + assert not await ctx_impl.has_messages() + + # Release the parked send. If the fix reached this gather, its task was already cancelled and + # this is a no-op; if it did not, the stale output lands in the just-restored queue. + release.set() + await asyncio.sleep(0) + + assert not await ctx_impl.has_messages()