Python: discard pending State on failed/cancelled superstep (#7859)

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

@@ -80,6 +80,9 @@ def __init__(
         self._iteration = 0
         self._max_iterations = max_iterations
         self._state = state
+        # When True, Workflow.run must reject successors even if the ResponseStream
+        # weakref is already gone (cleanup still in progress after a drop/cancel).
+        self._blocking_reuse_until_cleanup = False
 
         # Checkpointing related attributes
         self._previous_checkpoint_id: CheckpointID | None = None
@@ -127,51 +130,86 @@ async def run_until_convergence(self) -> AsyncGenerator[WorkflowEvent, None]:
             # Wake on either a live event or iteration completion, including silent supersteps.
             iteration_task = asyncio.create_task(self._run_iteration())
             event_task: asyncio.Task[WorkflowEvent] | None = None
+            # Track commit so cancel/abort cleanup spans the whole superstep
+            # (event wait → await iteration → drain → commit), not only the poll loop (#7859).
+            committed = False
+            # Defer failure events until after discard so dropping the ResponseStream
+            # cannot race a successor commit of stale pending writes (#7859).
+            deferred_failure_events: list[WorkflowEvent] = []
+            self._blocking_reuse_until_cleanup = True
             try:
-                while not iteration_task.done():
-                    event_task = asyncio.create_task(self._ctx.next_event())
-                    done, _ = await asyncio.wait((iteration_task, event_task), return_when=asyncio.FIRST_COMPLETED)
-                    if event_task in done:
-                        yield event_task.result()
-            finally:
-                # Cancellation and generator closure must not leave an event waiter or executor running.
-                tasks: list[asyncio.Task[Any]] = (
-                    [iteration_task] if event_task is None else [iteration_task, event_task]
-                )
-                for task in tasks:
-                    if not task.done():
-                        task.cancel()
-                await asyncio.gather(*tasks, return_exceptions=True)
+                try:
+                    while not iteration_task.done():
+                        event_task = asyncio.create_task(self._ctx.next_event())
+                        done, _ = await asyncio.wait(
+                            (iteration_task, event_task), return_when=asyncio.FIRST_COMPLETED
+                        )
+                        if event_task in done:
+                            event = event_task.result()
+                            event_task = None
+                            if event.type == "executor_failed":
+                                deferred_failure_events.append(event)
+                            else:
+                                yield event
+                finally:
+                    # Cancellation and generator closure must not leave an event waiter or executor running.
+                    tasks: list[asyncio.Task[Any]] = (
+                        [iteration_task] if event_task is None else [iteration_task, event_task]
+                    )
+                    for task in tasks:
+                        if not task.done():
+                            task.cancel()
+                    await asyncio.gather(*tasks, return_exceptions=True)
+
+                # Propagate errors from iteration, but first surface any pending events
+                try:
+                    await iteration_task
+                except Exception:
+                    # Discard staged writes immediately — before any await/yield — so a
+                    # streaming consumer that stops after executor_failed cannot leave
+                    # pending state for a later run to commit (#7859).
+                    self._state.discard()
+                    for event in deferred_failure_events:
+                        yield event
+                    deferred_failure_events.clear()
+                    if await self._ctx.has_events():
+                        for event in await self._ctx.drain_events():
+                            yield event
+                    raise
 
-            # Propagate errors from iteration, but first surface any pending events
-            try:
-                await iteration_task
-            except Exception:
-                # Make sure failure-related events (like ExecutorFailedEvent) are surfaced
+                for event in deferred_failure_events:
+                    yield event
+                deferred_failure_events.clear()
+
+                self._iteration += 1
+
+                # Drain any straggler events emitted at tail end
                 if await self._ctx.has_events():
                     for event in await self._ctx.drain_events():
                         yield event
-                raise
-            self._iteration += 1
 
-            # Drain any straggler events emitted at tail end
-            if await self._ctx.has_events():
-                for event in await self._ctx.drain_events():
-                    yield event
-
-            logger.info(f"Completed superstep {self._iteration}")
+                logger.info(f"Completed superstep {self._iteration}")
 
-            # Commit pending state changes at superstep boundary
-            self._state.commit()
+                # Commit pending state changes at superstep boundary
+                self._state.commit()
 
-            # Create checkpoint after each superstep iteration
-            await self.create_checkpoint_if_enabled()
+                # Create checkpoint after each superstep iteration. Keep
+                # ``committed`` false until this returns so residual pending
+                # staged during checkpoint prep is discarded on cancel/failure.
+                await self.create_checkpoint_if_enabled()
+                committed = True
 
-            yield WorkflowEvent.superstep_completed(iteration=self._iteration)
+                yield WorkflowEvent.superstep_completed(iteration=self._iteration)
 
-            # Check for convergence: no more messages to process
-            if not await self._ctx.has_messages():
-                break
+                # Check for convergence: no more messages to process
+                if not await self._ctx.has_messages():
+                    break
+            finally:
+                # Always discard on abort — including when iteration-task await above
+                # raised from executor cleanup — so staged writes cannot leak (#7859).
+                if not committed:
+                    self._state.discard()
+                self._blocking_reuse_until_cleanup = False
 
         logger.info(f"Workflow completed after {self._iteration} supersteps")
 
@@ -235,10 +273,18 @@ async def _prepare_checkpoint_state(self) -> None:
 
         This is used by checkpoint capture paths that need a complete, restorable
         state payload without necessarily writing to a checkpoint storage backend.
+
+        If staging executor/edge state fails or is cancelled before ``commit()``,
+        discard residual pending writes so a later run cannot commit a partial
+        checkpoint payload (#7859).
         """
-        await self._save_executor_states()
-        self._save_edge_runner_states()
-        self._state.commit()
+        try:
+            await self._save_executor_states()
+            self._save_edge_runner_states()
+            self._state.commit()
+        except BaseException:
+            self._state.discard()
+            raise
 
     async def create_checkpoint_if_enabled(self) -> None:
         """Create a checkpoint and save the checkpoint to the configured storage if one is configured.
@@ -249,9 +295,11 @@ async def create_checkpoint_if_enabled(self) -> None:
         if not self._ctx.has_checkpointing():
             return
 
+        prepared = False
         try:
             # Save executor states into committed state before creating the checkpoint.
             await self._prepare_checkpoint_state()
+            prepared = True
 
             checkpoint_id = await self._ctx.create_checkpoint(
                 self._workflow_name,
@@ -269,6 +317,10 @@ async def create_checkpoint_if_enabled(self) -> None:
             )
             self._previous_checkpoint_id = checkpoint_id
         except Exception as e:
+            # ``_prepare_checkpoint_state`` discards on its own failure; if we
+            # never reached that commit, clear any residual pending here too.
+            if not prepared:
+                self._state.discard()
             logger.warning(
                 "Failed to create checkpoint at iteration %d: %s. "
                 "Note that this does not fail the workflow run. "

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

@@ -409,6 +409,9 @@ def __init__(
         # ever iterating, the weakref dereferences to ``None`` once Python collects it,
         # so a subsequent ``run()`` is allowed.
         self._active_run: weakref.ref[ResponseStream[WorkflowEvent, WorkflowRunResult]] | None = None
+        # Strong hold for non-streaming exclusive entry points (e.g. cancel_pending_requests)
+        # that mutate the same runner/state without installing a ResponseStream weakref.
+        self._exclusive_run_hold: bool = False
 
     @property
     def status(self) -> WorkflowRunState:
@@ -1323,6 +1326,14 @@ async def cancel_pending_requests(
         if not all(isinstance(request_id, str) and request_id for request_id in selected_ids):
             raise ValueError("Pending workflow request IDs must be non-empty strings.")
 
+        # Share the same active/cleanup lock as ``run()`` so a cancellation continuation
+        # cannot start during ResponseStream cleanup (or while another run holds the lock)
+        # and commit pending State before the abandoned run discards it (#7859).
+        if self._is_run_active():
+            raise WorkflowException(
+                "Workflow is already running; concurrent runs are not allowed on the same instance."
+            )
+
         async def apply_cancellations() -> None:
             cancelled_events = await self._runner.context.cancel_request_info_events(selected_ids)
             for request_id, request_event in cancelled_events.items():
@@ -1337,12 +1348,16 @@ async def apply_cancellations() -> None:
                 )
                 await executor._cancel_pending_request(request_id, context)  # pyright: ignore[reportPrivateUsage]
 
-        if checkpoint_storage is not None:
-            self._runner.context.set_runtime_checkpoint_storage(checkpoint_storage)
-        runtime_tools = normalize_tools(tools) if tools is not None else None
-        self._runner.context.set_runtime_tools(runtime_tools)
+        # Acquire the hold first, but keep setup inside try/finally so a failure in
+        # normalize_tools / checkpoint restore cannot leave the workflow permanently locked.
+        self._exclusive_run_hold = True
         events: list[WorkflowEvent[Any]] = []
+        runtime_tools = None
         try:
+            if checkpoint_storage is not None:
+                self._runner.context.set_runtime_checkpoint_storage(checkpoint_storage)
+            runtime_tools = normalize_tools(tools) if tools is not None else None
+            self._runner.context.set_runtime_tools(runtime_tools)
             if checkpoint_id is not None:
                 await self._runner.restore_from_checkpoint(checkpoint_id, checkpoint_storage)
             async for event in self._run_workflow_with_tracing(
@@ -1357,6 +1372,7 @@ async def apply_cancellations() -> None:
                     continue
                 events.append(event)
         finally:
+            self._exclusive_run_hold = False
             if checkpoint_storage is not None:
                 self._runner.context.clear_runtime_checkpoint_storage()
             self._runner.context.clear_runtime_tools()
@@ -1411,5 +1427,11 @@ def _is_run_active(self) -> bool:
         Returns:
             True if a run is active, False otherwise.
         """
+        if self._exclusive_run_hold:
+            return True
+        # Runner cleanup can outlive a dropped ResponseStream (weakref cleared on GC).
+        # Keep the instance reserved until pending State is discarded (#7859).
+        if getattr(self._runner, "_blocking_reuse_until_cleanup", False):
+            return True
         existing_stream = self._active_run() if self._active_run is not None else None
         return existing_stream is not None

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

@@ -1590,3 +1590,104 @@ async def test_runner_drains_straggler_events_at_iteration_end():
     output_events = [e for e in events if e.type == "output"]
     # We should have output events from both executors
     assert len(output_events) >= 2
+
+@pytest.mark.asyncio
+async def test_failed_superstep_discards_pending_state_before_next_run() -> None:
+    """Pending State writes from a failed superstep must not leak into a later run (#7859)."""
+    from agent_framework import WorkflowBuilder
+
+    @dataclass
+    class Msg:
+        fail: bool
+
+    class FlakyExecutor(Executor):
+        @handler
+        async def run(self, message: Msg, ctx: WorkflowContext[Msg, str]) -> None:
+            if message.fail:
+                ctx.set_state("secret", "leaked-from-failed-run")
+                raise RuntimeError("simulated transient failure")
+            await ctx.yield_output("ok")
+
+    workflow = WorkflowBuilder(start_executor=FlakyExecutor(id="flaky")).build()
+
+    with pytest.raises(RuntimeError, match="simulated transient failure"):
+        async for _ in workflow.run(Msg(fail=True), stream=True):
+            pass
+
+    async for _ in workflow.run(Msg(fail=False), stream=True):
+        pass
+
+    committed = workflow._runner.state.export_state()  # pyright: ignore[reportPrivateUsage]
+    assert "secret" not in committed
+
+
+@pytest.mark.asyncio
+async def test_cancelled_superstep_discards_pending_state_before_next_run() -> None:
+    """Pending State writes from a cancelled superstep must not leak into a later run (#7859)."""
+    from agent_framework import WorkflowBuilder
+
+    @dataclass
+    class Msg:
+        cancel: bool
+
+    started = asyncio.Event()
+
+    class StagingThenBlockingExecutor(Executor):
+        @handler
+        async def run(self, message: Msg, ctx: WorkflowContext[Msg, str]) -> None:
+            if message.cancel:
+                ctx.set_state("secret", "leaked-from-cancelled-run")
+                started.set()
+                await asyncio.sleep(3600)
+            await ctx.yield_output("ok")
+
+    workflow = WorkflowBuilder(start_executor=StagingThenBlockingExecutor(id="blocker")).build()
+
+    async def run_and_cancel() -> None:
+        async for _ in workflow.run(Msg(cancel=True), stream=True):
+            pass
+
+    task = asyncio.create_task(run_and_cancel())
+    await started.wait()
+    task.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await task
+
+    async for _ in workflow.run(Msg(cancel=False), stream=True):
+        pass
+
+    committed = workflow._runner.state.export_state()  # pyright: ignore[reportPrivateUsage]
+    assert "secret" not in committed
+
+
+@pytest.mark.asyncio
+async def test_failed_superstep_discards_even_if_executor_cleanup_raises() -> None:
+    """Executor cleanup that raises after set_state must not skip State.discard (#7859)."""
+    from agent_framework import WorkflowBuilder
+
+    @dataclass
+    class Msg:
+        fail: bool
+
+    class CleanupRaisesExecutor(Executor):
+        @handler
+        async def run(self, message: Msg, ctx: WorkflowContext[Msg, str]) -> None:
+            if message.fail:
+                ctx.set_state("secret", "leaked-from-cleanup-raise")
+                try:
+                    raise RuntimeError("primary failure")
+                finally:
+                    raise RuntimeError("cleanup failure")  # noqa: B012
+            await ctx.yield_output("ok")
+
+    workflow = WorkflowBuilder(start_executor=CleanupRaisesExecutor(id="cleanup")).build()
+
+    with pytest.raises(RuntimeError):
+        async for _ in workflow.run(Msg(fail=True), stream=True):
+            pass
+
+    async for _ in workflow.run(Msg(fail=False), stream=True):
+        pass
+
+    committed = workflow._runner.state.export_state()  # pyright: ignore[reportPrivateUsage]
+    assert "secret" not in committed

python/packages/declarative/agent_framework_declarative/_workflows/_declarative_base.py

@@ -401,6 +401,14 @@ def set_state_data(self, data: DeclarativeStateData) -> None:
         """Set the full state data dict in state."""
         self._state.set(DECLARATIVE_STATE_KEY, data)
 
+    def commit(self) -> None:
+        """Commit pending runner State writes.
+
+        Used when an action publishes diagnostic state and then fails the
+        superstep. The runner discards uncommitted writes on failure (#7859).
+        """
+        self._state.commit()
+
     def get(self, path: str, default: Any = None) -> Any:
         """Get a value from the state using a dot-notated path.
 

python/packages/declarative/agent_framework_declarative/_workflows/_executors_http.py

@@ -206,6 +206,9 @@ async def handle_action(
 
         # Non-success path: still publish headers diagnostically, then raise.
         self._assign_response_headers(state, result)
+        # Runner discards pending State when the superstep fails (#7859 / #8306).
+        # Commit first so diagnostic headers remain readable after the error.
+        state.commit()
         raise DeclarativeActionError(f"HTTP request to '{url}' failed with status code {result.status_code}.")
 
     # ----- Field resolution ----------------------------------------------------