Python: fix(workflow): discard pending state after failed supersteps

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

@@ -282,8 +282,25 @@ async def send_to_edge(edge: Edge) -> bool:
                 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)
+            tasks = [asyncio.create_task(send_to_edge(edge)) for edge in deliverable_edges]
+            if not tasks:
+                return False  # asyncio.wait() requires a non-empty iterable
+            try:
+                done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_EXCEPTION)
+            except asyncio.CancelledError:
+                # If the wait() call itself is cancelled, cancel all child tasks
+                # before propagating to avoid orphaned work
+                for t in tasks:
+                    t.cancel()
+                await asyncio.gather(*tasks, return_exceptions=True)
+                raise
+            exceptions = [t.exception() for t in done if t.exception() is not None]
+            if exceptions:
+                for t in pending:
+                    t.cancel()
+                await asyncio.gather(*tasks, return_exceptions=True)
+                raise exceptions[0]
+            results = [t.result() for t in tasks]
             return any(results)
 
         # If we get here, it's a broadcast message with no deliverable edges

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

@@ -145,6 +145,8 @@ async def run_until_convergence(self) -> AsyncGenerator[WorkflowEvent, None]:
             try:
                 await iteration_task
             except Exception:
+                # Discard pending state writes from the failed superstep
+                self._state.discard()
                 # Make sure failure-related events (like ExecutorFailedEvent) are surfaced
                 if await self._ctx.has_events():
                     for event in await self._ctx.drain_events():
@@ -218,15 +220,56 @@ async def _deliver_messages_for_edge_runner(edge_runner: EdgeRunner) -> None:
                 logger.debug(f"No outgoing edges found for executor {source_executor_id}; dropping messages.")
                 return
 
-            tasks = [_deliver_messages_for_edge_runner(edge_runner) for edge_runner in associated_edge_runners]
-            await asyncio.gather(*tasks)
+            tasks = [asyncio.create_task(_deliver_messages_for_edge_runner(edge_runner)) for edge_runner in associated_edge_runners]
+            if not tasks:
+                return  # asyncio.wait() requires a non-empty iterable
+            # Use FIRST_EXCEPTION to cancel pending siblings when one fails
+            try:
+                done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_EXCEPTION)
+            except asyncio.CancelledError:
+                # If the wait() call itself is cancelled, cancel all child tasks
+                # before propagating to avoid orphaned work
+                for t in tasks:
+                    t.cancel()
+                await asyncio.gather(*tasks, return_exceptions=True)
+                raise
+            exceptions = [t.exception() for t in done if t.exception() is not None]
+            if exceptions:
+                for t in pending:
+                    t.cancel()
+                # Await all tasks to ensure cancelled tasks stop before propagating failure
+                await asyncio.gather(*tasks, return_exceptions=True)
+                raise exceptions[0]
 
         message_batches = await self._ctx.drain_messages()
-        tasks = [
-            _deliver_messages(source_executor_id, source_messages)
+        # Create actual Task objects so we can cancel them if needed
+        task_objects = [
+            asyncio.create_task(_deliver_messages(source_executor_id, source_messages))
             for source_executor_id, source_messages in message_batches.items()
         ]
-        await asyncio.gather(*tasks)
+
+        if not task_objects:
+            return  # asyncio.wait() requires a non-empty iterable
+
+        try:
+            done, pending = await asyncio.wait(
+                task_objects,
+                return_when=asyncio.FIRST_EXCEPTION
+            )
+        except asyncio.CancelledError:
+            # If the wait() call itself is cancelled, cancel all child tasks
+            # before propagating to avoid orphaned work
+            for t in task_objects:
+                t.cancel()
+            await asyncio.gather(*task_objects, return_exceptions=True)
+            raise
+
+        exceptions = [t.exception() for t in done if t.exception() is not None]
+        if exceptions:
+            for t in pending:
+                t.cancel()
+            await asyncio.gather(*task_objects, return_exceptions=True)
+            raise exceptions[0]
 
     async def _prepare_checkpoint_state(self) -> None:
         """Persist executor snapshots into committed shared state.

python/packages/core/tests/workflow/test_workflow.py

@@ -1,6 +1,7 @@
 # Copyright (c) Microsoft. All rights reserved.
 
 import asyncio
+import contextlib
 import gc
 import logging
 import tempfile
@@ -640,6 +641,151 @@ def _build():
         assert result2.get_outputs()[0] == ["run2:message2"]
 
 
+@dataclass
+class FlakyMessage:
+    """A message that can fail on demand for testing state discard behavior."""
+
+    fail: bool
+
+
+class FlakyStateExecutor(Executor):
+    """An executor that fails on demand to test state discard on failure."""
+
+    @handler
+    async def handle_message(
+        self,
+        message: FlakyMessage,
+        ctx: WorkflowContext[FlakyMessage, str],
+    ) -> None:
+        if message.fail:
+            ctx.set_state("secret", "leaked-from-failed-run")
+            raise RuntimeError("simulated transient failure")
+
+        await ctx.yield_output("ok")
+
+
+async def test_workflow_discards_pending_state_after_failed_superstep():
+    """Test that pending state from a failed superstep is discarded and not committed.
+
+    This is a regression test for GitHub issue #7859: pending state writes from
+    a failed superstep must not leak into a later successful run on the same
+    Workflow instance.
+    """
+    workflow = WorkflowBuilder(start_executor=FlakyStateExecutor(id="flaky")).build()
+
+    # First run: fails after staging a state write
+    with pytest.raises(RuntimeError, match="simulated transient failure"):
+        await workflow.run(FlakyMessage(fail=True))
+
+    # Verify the failed run did not leave the staged write pending
+    assert workflow._runner.state._pending == {}
+
+    # Second run: succeeds without touching "secret"
+    result = await workflow.run(FlakyMessage(fail=False))
+    assert result.get_final_state() == WorkflowRunState.IDLE
+    assert result.get_outputs() == ["ok"]
+
+    # Verify the leaked state from the failed run is NOT in committed state
+    committed_state = workflow._runner.state.export_state()
+    assert "secret" not in committed_state
+
+
+@dataclass
+class FanOutTestMessage:
+    """Message for fan-out state leak test."""
+    should_fail: bool
+
+
+class FanOutSourceExecutor(Executor):
+    """Source executor that sends fan-out test messages."""
+
+    @handler
+    async def handle(self, message: FanOutTestMessage, ctx: WorkflowContext[FanOutTestMessage]) -> None:
+        # Forward the message to targets
+        await ctx.send_message(message)
+
+
+class FailingTargetExecutor(Executor):
+    """Target that fails immediately on execute."""
+
+    def __init__(self, id: str, b_started: asyncio.Event | None = None) -> None:
+        super().__init__(id=id)
+        self._b_started = b_started
+
+    @handler
+    async def handle(self, message: FanOutTestMessage, ctx: WorkflowContext) -> None:
+        if message.should_fail:
+            # Wait for slow target to start before failing to ensure deterministic ordering
+            if self._b_started:
+                await self._b_started.wait()
+            raise RuntimeError("target A failed")
+
+
+class SlowStateWritingTargetExecutor(Executor):
+    """Target that writes state after being unblocked (vulnerable pattern if not cancelled)."""
+
+    def __init__(self, id: str, b_started: asyncio.Event, release_b: asyncio.Event, b_task_ref: list) -> None:
+        super().__init__(id=id)
+        self._b_started = b_started
+        self._release_b = release_b
+        self._b_task_ref = b_task_ref
+
+    @handler
+    async def handle(self, message: FanOutTestMessage, ctx: WorkflowContext) -> None:
+        self._b_task_ref.append(asyncio.current_task())
+        self._b_started.set()
+        if message.should_fail:
+            await self._release_b.wait()
+            ctx.set_state("leak_key", "leaked_value")
+
+
+async def test_workflow_discards_pending_state_after_fanout_failure():
+    """Test that pending state from a fan-out sibling target is discarded when another target fails.
+
+    Regression test for GitHub issue #7859: when a fan-out superstep has one target fail
+    while a sibling target is still running, the sibling's state writes must not leak
+    into committed state.
+    """
+    b_started = asyncio.Event()
+    release_b = asyncio.Event()
+    b_task_ref: list[asyncio.Task] = []
+
+    source = FanOutSourceExecutor(id="source")
+    failing_target = FailingTargetExecutor(id="failing_target", b_started=b_started)
+    slow_target = SlowStateWritingTargetExecutor(
+        id="slow_target", b_started=b_started, release_b=release_b, b_task_ref=b_task_ref
+    )
+
+    workflow = (
+        WorkflowBuilder(start_executor=source)
+        .add_fan_out_edges(source, [failing_target, slow_target])
+        .build()
+    )
+
+    # Verify topology: single FanOutEdgeGroup with two targets under one edge_runner
+    from agent_framework._workflows._edge import FanOutEdgeGroup
+    fan_out_groups = [eg for eg in workflow.edge_groups if isinstance(eg, FanOutEdgeGroup)]
+    assert len(fan_out_groups) == 1, f"Expected 1 FanOutEdgeGroup, got {len(fan_out_groups)}"
+    assert fan_out_groups[0].target_ids == [failing_target.id, slow_target.id], \
+        f"Expected targets [{failing_target.id}, {slow_target.id}], got {fan_out_groups[0].target_ids}"
+
+    run_task = asyncio.create_task(workflow.run(FanOutTestMessage(should_fail=True)))
+    with pytest.raises(RuntimeError, match="target A failed"):
+        await asyncio.wait_for(run_task, timeout=5.0)
+
+    release_b.set()
+    with contextlib.suppress(asyncio.CancelledError, Exception):
+        await asyncio.wait_for(b_task_ref[0], timeout=0.5)
+
+    b_task_ref.clear()
+
+    result = await workflow.run(FanOutTestMessage(should_fail=False))
+    assert result.get_final_state() == WorkflowRunState.IDLE
+
+    committed_state = workflow._runner.state.export_state()
+    assert "leak_key" not in committed_state
+
+
 async def test_workflow_checkpoint_runtime_only_configuration(
     simple_executor: Executor,
 ):