From 956f43f046c78479ba207b17591af9e212e38299 Mon Sep 17 00:00:00 2001 From: Namraa Patel Date: Tue, 25 Aug 2026 19:23:08 +0530 Subject: [PATCH 1/3] fix(workflow): discard pending state after failed supersteps --- .../agent_framework/_workflows/_runner.py | 4 ++ .../core/tests/workflow/test_workflow.py | 49 +++++++++++++++++++ 2 files changed, 53 insertions(+) diff --git a/python/packages/core/agent_framework/_workflows/_runner.py b/python/packages/core/agent_framework/_workflows/_runner.py index ac5558dbe1c..519a76eace5 100644 --- a/python/packages/core/agent_framework/_workflows/_runner.py +++ b/python/packages/core/agent_framework/_workflows/_runner.py @@ -139,6 +139,8 @@ async def run_until_convergence(self) -> AsyncGenerator[WorkflowEvent, None]: iteration_task.cancel() with contextlib.suppress(asyncio.CancelledError): await iteration_task + # Discard pending state writes from the cancelled superstep + self._state.discard() raise # Propagate errors from iteration, but first surface any pending events @@ -149,6 +151,8 @@ async def run_until_convergence(self) -> AsyncGenerator[WorkflowEvent, None]: if await self._ctx.has_events(): for event in await self._ctx.drain_events(): yield event + # Discard pending state writes from the failed superstep + self._state.discard() raise self._iteration += 1 diff --git a/python/packages/core/tests/workflow/test_workflow.py b/python/packages/core/tests/workflow/test_workflow.py index 2f672f591d3..f1fb109129f 100644 --- a/python/packages/core/tests/workflow/test_workflow.py +++ b/python/packages/core/tests/workflow/test_workflow.py @@ -609,6 +609,55 @@ 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 + + async def test_workflow_checkpoint_runtime_only_configuration( simple_executor: Executor, ): From da0154c748843813e83f734f660612c09d691165 Mon Sep 17 00:00:00 2001 From: Namraa Patel Date: Thu, 27 Aug 2026 14:01:28 +0530 Subject: [PATCH 2/3] fix: discard pending state before yielding failure events --- python/packages/core/agent_framework/_workflows/_runner.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_runner.py b/python/packages/core/agent_framework/_workflows/_runner.py index 519a76eace5..def8e99fe6b 100644 --- a/python/packages/core/agent_framework/_workflows/_runner.py +++ b/python/packages/core/agent_framework/_workflows/_runner.py @@ -147,12 +147,12 @@ 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(): yield event - # Discard pending state writes from the failed superstep - self._state.discard() raise self._iteration += 1 From 8bfce51295b0f62ca9eb67442673eeb8c8e15900 Mon Sep 17 00:00:00 2001 From: Namraa Patel Date: Mon, 31 Aug 2026 18:41:48 +0530 Subject: [PATCH 3/3] fix: discard pending state from failed supersteps --- .../_workflows/_edge_runner.py | 21 +++- .../agent_framework/_workflows/_runner.py | 53 ++++++++-- .../core/tests/workflow/test_workflow.py | 97 +++++++++++++++++++ 3 files changed, 162 insertions(+), 9 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_edge_runner.py b/python/packages/core/agent_framework/_workflows/_edge_runner.py index c14582894b9..bfe0334a338 100644 --- a/python/packages/core/agent_framework/_workflows/_edge_runner.py +++ b/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 diff --git a/python/packages/core/agent_framework/_workflows/_runner.py b/python/packages/core/agent_framework/_workflows/_runner.py index def8e99fe6b..d1893112747 100644 --- a/python/packages/core/agent_framework/_workflows/_runner.py +++ b/python/packages/core/agent_framework/_workflows/_runner.py @@ -139,8 +139,6 @@ async def run_until_convergence(self) -> AsyncGenerator[WorkflowEvent, None]: iteration_task.cancel() with contextlib.suppress(asyncio.CancelledError): await iteration_task - # Discard pending state writes from the cancelled superstep - self._state.discard() raise # Propagate errors from iteration, but first surface any pending events @@ -222,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. diff --git a/python/packages/core/tests/workflow/test_workflow.py b/python/packages/core/tests/workflow/test_workflow.py index f1fb109129f..95829a64232 100644 --- a/python/packages/core/tests/workflow/test_workflow.py +++ b/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 @@ -658,6 +659,102 @@ async def test_workflow_discards_pending_state_after_failed_superstep(): 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, ):