Python: Keep a dropped run's finalizer from discarding state writes

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

@@ -125,9 +125,10 @@ async def run_until_convergence(self) -> AsyncGenerator[WorkflowEvent, None]:
             yield WorkflowEvent.superstep_started(iteration=self._iteration + 1)
 
             committed = False
+            superstep_epoch = self._state.begin_epoch()
             try:
                 # Wake on either a live event or iteration completion, including silent supersteps.
-                iteration_task = asyncio.create_task(self._run_iteration())
+                iteration_task = self._state.bound_context().run(asyncio.create_task, self._run_iteration())
                 event_task: asyncio.Task[WorkflowEvent] | None = None
                 try:
                     while not iteration_task.done():
@@ -167,7 +168,8 @@ async def run_until_convergence(self) -> AsyncGenerator[WorkflowEvent, None]:
                 self._state.commit()
                 committed = True
             finally:
-                if not committed:
+                # A stale run's finalizer must not discard pending writes owned by a newer superstep.
+                if not committed and self._state.epoch is superstep_epoch:
                     self._state.discard()
 
             # Create checkpoint after each superstep iteration
@@ -243,6 +245,7 @@ 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.
         """
+        self._state.begin_epoch()
         await self._save_executor_states()
         self._save_edge_runner_states()
         self._state.commit()

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

@@ -1,9 +1,17 @@
 # Copyright (c) Microsoft. All rights reserved.
 
 import copy
+import logging
 from collections.abc import Callable
+from contextvars import Context, ContextVar, copy_context
 from typing import Any
 
+logger = logging.getLogger(__name__)
+
+# The State and epoch that writes from the current task belong to. Set for each superstep's iteration
+# task so executor tasks inherit it; writes made outside any superstep are always accepted.
+_bound_epoch: ContextVar[tuple["State", object] | None] = ContextVar("_bound_epoch", default=None)
+
 
 class State:
     """Manages shared state across executors within a workflow.
@@ -18,6 +26,8 @@ class State:
         - `get()` checks pending buffer first, then committed state
         - `commit()` moves all pending changes to committed state (called by Runner at superstep boundary)
         - `discard()` clears pending changes without committing
+        - `begin_epoch()` starts a new owner of pending writes; writes, deletes and clears from tasks
+          bound to an earlier epoch (for example a dropped run's cancelled executor) are ignored
 
     Reserved Keys:
         Keys starting with underscore (_) are reserved for internal framework use.
@@ -28,6 +38,30 @@ def __init__(self) -> None:
         """Initialize the state."""
         self._committed: dict[str, Any] = {}
         self._pending: dict[str, Any] = {}
+        self._epoch = object()
+
+    @property
+    def epoch(self) -> object:
+        """The current write epoch."""
+        return self._epoch
+
+    def begin_epoch(self) -> object:
+        """Start a new write epoch and return it."""
+        self._epoch = object()
+        return self._epoch
+
+    def bound_context(self) -> Context:
+        """Return a copy of the current context whose writes, and its tasks' writes, belong to the current epoch."""
+        context = copy_context()
+        context.run(_bound_epoch.set, (self, self._epoch))
+        return context
+
+    def _is_stale_write(self) -> bool:
+        bound = _bound_epoch.get()
+        if bound is None or bound[0] is not self or bound[1] is self._epoch:
+            return False
+        logger.debug("Ignoring a state write from a superstep that no longer owns this state.")
+        return True
 
     def set(self, key: str, value: Any) -> None:
         """Set a value in the pending state buffer.
@@ -43,6 +77,8 @@ def set(self, key: str, value: Any) -> None:
             the .NET behavior and the superstep execution model where all executors
             in a superstep see the same committed state at the start.
         """
+        if self._is_stale_write():
+            return
         self._pending[key] = copy.deepcopy(value)
 
     def get(self, key: str, default: Any = None) -> Any:
@@ -88,6 +124,8 @@ def delete(self, key: str) -> None:
         to indicate deletion at commit time. If it only exists in pending,
         it is removed from pending.
         """
+        if self._is_stale_write():
+            return
         if key not in self._pending and key not in self._committed:
             raise KeyError(f"Key '{key}' not found in state.")
 
@@ -100,6 +138,8 @@ def delete(self, key: str) -> None:
 
     def clear(self) -> None:
         """Clear both committed and pending state."""
+        if self._is_stale_write():
+            return
         self._committed.clear()
         self._pending.clear()
 

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

@@ -589,6 +589,11 @@ async def _run_workflow_with_tracing(
                 if not is_continuation:
                     self._runner.reset_iteration_count()
 
+                # Pending writes here come from a superstep or checkpoint save that never committed; do not
+                # commit them, and ignore later writes from that superstep's executors if it is still finalizing.
+                self._runner.state.begin_epoch()
+                self._runner.state.discard()
+
                 # Store run kwargs in State so executors can access them.
                 # Per-run kwargs semantics:
                 # - On a fresh message run, prior kwargs go away (set to {}

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

@@ -39,6 +39,7 @@
     handler,
     response_handler,
 )
+from agent_framework._workflows._const import EXECUTOR_STATE_KEY
 
 
 @dataclass
@@ -1483,6 +1484,203 @@ async def test_workflow_partial_stream_does_not_clobber_successor_active_run() -
     await asyncio.sleep(0)
 
 
+async def test_workflow_partial_stream_does_not_discard_successor_state_writes() -> None:
+    """A stale runner finalizer must not discard a successor run's pending state writes.
+
+    Same GC-finalizer race as
+    ``test_workflow_partial_stream_does_not_clobber_successor_active_run`` but inside a
+    superstep: stream A is dropped while its executor has staged a write, stream B stages
+    its own write, then A's deferred ``finally`` runs. It must only discard A's superstep,
+    so B's write is committed when B completes.
+    """
+    a_staged = asyncio.Event()
+    b_staged = asyncio.Event()
+    release_b = asyncio.Event()
+
+    class _Writer(Executor):
+        @handler
+        async def handle(self, message: str, ctx: WorkflowContext[str, str]) -> None:
+            if message == "A":
+                ctx.set_state("a_key", "from-run-A")
+                a_staged.set()
+                await asyncio.Event().wait()
+            else:
+                ctx.set_state("b_key", "from-run-B")
+                b_staged.set()
+                await release_b.wait()
+                await ctx.yield_output("b-done")
+
+    workflow = WorkflowBuilder(start_executor=_Writer(id="writer")).build()
+
+    # Drive stream A until its executor has staged a write and is blocked, then drop it.
+    stream_a = workflow.run("A", stream=True)
+    aiter_a = stream_a.__aiter__()
+    while (await aiter_a.__anext__()).type != "executor_invoked":
+        pass
+    await a_staged.wait()
+    del stream_a
+    del aiter_a
+    gc.collect()
+
+    async def consume_b() -> WorkflowRunResult:
+        return await workflow.run("B", stream=True).get_final_response()
+
+    task_b = asyncio.create_task(consume_b())
+    await asyncio.wait_for(b_staged.wait(), 5)
+
+    # Let stream A's scheduled aclose run its finally while B's write is pending.
+    for _ in range(20):
+        await asyncio.sleep(0.01)
+    release_b.set()
+    result_b = await asyncio.wait_for(task_b, 5)
+
+    assert result_b.get_outputs() == ["b-done"]
+    assert result_b.get_final_state() == WorkflowRunState.IDLE
+    assert workflow._runner.state.get("b_key") == "from-run-B"  # pyright: ignore[reportPrivateUsage]
+    assert workflow._runner.state.get("a_key") is None  # pyright: ignore[reportPrivateUsage]
+
+
+async def test_workflow_partial_stream_does_not_discard_successor_entry_checkpoint_snapshots() -> None:
+    """A stale runner finalizer must not discard snapshots staged by a successor's entry checkpoint.
+
+    Stream A is dropped mid-superstep. Stream B's entry checkpoint stages the writer's snapshot,
+    then suspends in the next executor's ``on_checkpoint_save`` while A's deferred ``finally``
+    runs. B's entry checkpoint must still hold B's writer snapshot.
+    """
+    a_staged = asyncio.Event()
+    hook_suspended = asyncio.Event()
+    release_hook = asyncio.Event()
+    gate_armed = False
+
+    class _Writer(Executor):
+        def __init__(self, id: str) -> None:
+            super().__init__(id=id)
+            self.saves = 0
+
+        @handler
+        async def handle(self, message: str, ctx: WorkflowContext[str, str]) -> None:
+            if message == "A":
+                a_staged.set()
+                await asyncio.Event().wait()
+            else:
+                await ctx.yield_output("b-done")
+
+        async def on_checkpoint_save(self) -> dict[str, Any]:
+            self.saves += 1
+            return {"saves": self.saves}
+
+    class _SuspendingSave(Executor):
+        @handler
+        async def handle(self, message: str, ctx: WorkflowContext) -> None:
+            pass
+
+        async def on_checkpoint_save(self) -> dict[str, Any]:
+            if gate_armed:
+                hook_suspended.set()
+                await release_hook.wait()
+            return {}
+
+    storage = InMemoryCheckpointStorage()
+    writer = _Writer(id="writer")
+    workflow = (
+        WorkflowBuilder(start_executor=writer, checkpoint_storage=storage)
+        .add_edge(writer, _SuspendingSave(id="gate"))
+        .build()
+    )
+
+    stream_a = workflow.run("A", stream=True)
+    aiter_a = stream_a.__aiter__()
+    while (await aiter_a.__anext__()).type != "executor_invoked":
+        pass
+    await a_staged.wait()
+    a_checkpoint_ids = set(await storage.list_checkpoint_ids(workflow_name=workflow.name))
+    del stream_a
+    del aiter_a
+    gc.collect()
+
+    gate_armed = True
+    task_b = asyncio.create_task(workflow.run("B", stream=True).get_final_response())
+    await asyncio.wait_for(hook_suspended.wait(), 5)
+
+    # Let stream A's scheduled aclose run its finally while B's writer snapshot is pending.
+    for _ in range(20):
+        await asyncio.sleep(0.01)
+    release_hook.set()
+    result_b = await asyncio.wait_for(task_b, 5)
+    assert result_b.get_outputs() == ["b-done"]
+
+    checkpoints = await storage.list_checkpoints(workflow_name=workflow.name)
+    [b_entry] = [c for c in checkpoints if c.checkpoint_id not in a_checkpoint_ids and c.iteration_count == 0]
+    assert b_entry.state[EXECUTOR_STATE_KEY]["writer"] == {"saves": 2}
+
+
+async def test_workflow_partial_stream_cancellation_write_not_committed_by_successor() -> None:
+    """A dropped run's cancellation cleanup must not write into a successor run's state.
+
+    Stream A is dropped while its executor is blocked. Stream B passes its run start and stages its
+    own write, then A is finalized: its executor's ``finally`` writes state during cancellation.
+    B's superstep commit must keep B's write and must not persist A's cleanup write.
+    """
+    a_blocked = asyncio.Event()
+    b_staged = asyncio.Event()
+    release_b = asyncio.Event()
+
+    class _Router(Executor):
+        @handler
+        async def handle(self, message: str, ctx: WorkflowContext[str]) -> None:
+            await ctx.send_message(message, target_id=f"worker_{message.lower()}")
+
+    class _WorkerA(Executor):
+        @handler
+        async def handle(self, message: str, ctx: WorkflowContext) -> None:
+            try:
+                a_blocked.set()
+                await asyncio.Event().wait()
+            finally:
+                ctx.set_state("a_cleanup", "from-run-A")
+
+    class _WorkerB(Executor):
+        @handler
+        async def handle(self, message: str, ctx: WorkflowContext[str, str]) -> None:
+            ctx.set_state("b_key", "from-run-B")
+            b_staged.set()
+            await release_b.wait()
+            await ctx.yield_output("b-done")
+
+    # Separate workers so B's superstep is not serialized behind A's blocked handler.
+    router = _Router(id="router")
+    workflow = (
+        WorkflowBuilder(start_executor=router)
+        .add_edge(router, _WorkerA(id="worker_a"))
+        .add_edge(router, _WorkerB(id="worker_b"))
+        .build()
+    )
+
+    stream_a = workflow.run("A", stream=True)
+    aiter_a = stream_a.__aiter__()
+    while (event := await aiter_a.__anext__()).type != "executor_invoked" or event.executor_id != "worker_a":
+        pass
+    await a_blocked.wait()
+    # Keep A's generator so its finalization can be ordered after B's run start.
+    generator_a = cast(Any, stream_a)._iterator
+    del stream_a
+    del aiter_a
+    gc.collect()
+
+    task_b = asyncio.create_task(workflow.run("B", stream=True).get_final_response())
+    await asyncio.wait_for(b_staged.wait(), 5)
+
+    # Finalize A the way the event loop does for a dropped stream.
+    await generator_a.aclose()
+    release_b.set()
+    result_b = await asyncio.wait_for(task_b, 5)
+
+    assert result_b.get_outputs() == ["b-done"]
+    committed = workflow._runner.state.export_state()  # pyright: ignore[reportPrivateUsage]
+    assert committed["b_key"] == "from-run-B"
+    assert "a_cleanup" not in committed
+
+
 async def test_workflow_stale_runtime_checkpoint_storage_not_inherited() -> None:
     """A new run must not inherit a prior run's leftover runtime checkpoint storage.