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

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

@@ -81,6 +81,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
@@ -128,51 +131,83 @@ async def run_until_convergence(self) -> AsyncGenerator[WorkflowEvent, None]:
             # Run iteration concurrently with live event streaming: we poll
             # for new events while the iteration coroutine progresses.
             iteration_task = asyncio.create_task(self._run_iteration())
+            # Track commit so cancel/abort cleanup spans the whole superstep
+            # (polling → 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] = []
             try:
+                self._blocking_reuse_until_cleanup = True
                 while not iteration_task.done():
                     try:
                         # Wait briefly for any new event; timeout allows progress checks
                         event = await asyncio.wait_for(self._ctx.next_event(), timeout=0.05)
-                        yield event
+                        if event.type == "executor_failed":
+                            deferred_failure_events.append(event)
+                        else:
+                            yield event
                     except asyncio.TimeoutError:
                         # Periodically continue to let iteration advance
                         continue
-            except asyncio.CancelledError:
-                # Propagate cancellation to the iteration task to avoid orphaned work
-                iteration_task.cancel()
-                with contextlib.suppress(asyncio.CancelledError):
-                    await iteration_task
-                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
-                if await self._ctx.has_events():
-                    for event in await self._ctx.drain_events():
+                # 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
-                raise
-            self._iteration += 1
+                    deferred_failure_events.clear()
+                    if await self._ctx.has_events():
+                        for event in await self._ctx.drain_events():
+                            yield event
+                    raise
 
-            # Drain any straggler events emitted at tail end
-            if await self._ctx.has_events():
-                for event in await self._ctx.drain_events():
+                for event in deferred_failure_events:
                     yield event
+                deferred_failure_events.clear()
 
-            logger.info(f"Completed superstep {self._iteration}")
-
-            # Commit pending state changes at superstep boundary
-            self._state.commit()
+                self._iteration += 1
 
-            # Create checkpoint after each superstep iteration
-            await self.create_checkpoint_if_enabled()
-
-            yield WorkflowEvent.superstep_completed(iteration=self._iteration)
+                # Drain any straggler events emitted at tail end
+                if await self._ctx.has_events():
+                    for event in await self._ctx.drain_events():
+                        yield event
 
-            # Check for convergence: no more messages to process
-            if not await self._ctx.has_messages():
-                break
+                logger.info(f"Completed superstep {self._iteration}")
+
+                # Commit pending state changes at superstep boundary
+                self._state.commit()
+                committed = True
+
+                # Create checkpoint after each superstep iteration
+                await self.create_checkpoint_if_enabled()
+
+                yield WorkflowEvent.superstep_completed(iteration=self._iteration)
+
+                # Check for convergence: no more messages to process
+                if not await self._ctx.has_messages():
+                    break
+            except BaseException:
+                # Cancel during poll/drain, or an iteration task that ends cancelled,
+                # must still abandon staged writes before the commit boundary (#7859).
+                # Await cleanup with broad suppression so a raising executor ``finally``
+                # cannot skip discard below.
+                if not iteration_task.done():
+                    iteration_task.cancel()
+                    with contextlib.suppress(BaseException):
+                        await iteration_task
+                raise
+            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")
 

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

@@ -392,6 +392,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:
@@ -1307,6 +1310,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():
@@ -1321,6 +1332,7 @@ async def apply_cancellations() -> None:
                 )
                 await executor._cancel_pending_request(request_id, context)  # pyright: ignore[reportPrivateUsage]
 
+        self._exclusive_run_hold = True
         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
@@ -1341,6 +1353,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()
@@ -1395,5 +1408,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

@@ -1496,3 +1496,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

@@ -334,6 +334,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 ----------------------------------------------------