diff --git a/python/packages/core/agent_framework/_workflows/_runner.py b/python/packages/core/agent_framework/_workflows/_runner.py index ac5558dbe1..def8e99fe6 100644 --- a/python/packages/core/agent_framework/_workflows/_runner.py +++ b/python/packages/core/agent_framework/_workflows/_runner.py @@ -139,12 +139,16 @@ 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 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(): diff --git a/python/packages/core/tests/workflow/test_workflow.py b/python/packages/core/tests/workflow/test_workflow.py index 2f672f591d..f1fb109129 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, ):