Python: carry fan-in edge buffers through checkpoint restore

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.

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"
 

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

@@ -1,12 +1,14 @@
 # Copyright (c) Microsoft. All rights reserved.
 
 import asyncio
+import json
 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
 from ..observability import EdgeGroupDeliveryStatus, OtelAttr, create_edge_group_processing_span
 from ._edge import (
     Edge,
@@ -24,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."""
 
@@ -57,6 +85,40 @@ 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. 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([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.
+
+        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:
@@ -278,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
@@ -297,13 +357,56 @@ 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
         # Buffer to hold messages before sending them to the target 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,

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

@@ -6,16 +6,16 @@
 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 ._edge_runner import EdgeRunner, create_edge_runner, gather_cancelling_siblings_on_error
 from ._events import WorkflowEvent
 from ._executor import Executor
 from ._runner_context import (
@@ -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
@@ -219,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 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.
@@ -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.
 

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: