Python: fix(core): keep a dropped run's late finalizer from discarding the next run's state writes

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

@@ -81,6 +81,13 @@ def __init__(
         self._max_iterations = max_iterations
         self._state = state
 
+        # Incremented once per drive of ``run_until_convergence``. A dropped
+        # stream's late finalizer (GeneratorExit arrives after a successor run
+        # already drives this shared state) must not discard the successor's
+        # pending writes, so the superstep cleanup below is epoch-guarded
+        # (#9473).
+        self._run_epoch = 0
+
         # Checkpointing related attributes
         self._previous_checkpoint_id: CheckpointID | None = None
 
@@ -105,6 +112,20 @@ def reset_iteration_count(self) -> None:
         """
         self._iteration = 0
 
+    def begin_run(self) -> int:
+        """Stamp a new run's ownership of this shared state, synchronously.
+
+        Advances the run epoch and sweeps a dropped predecessor's leftover
+        pending writes. Must run before the new run stages anything (the entry
+        checkpoint's executor snapshots land in the same pending map), so
+        Workflow.run() calls it from its synchronous section, ahead of any
+        await point (#9473). The stale finalizer's epoch check then fails even
+        during the successor's checkpoint staging.
+        """
+        self._run_epoch += 1
+        self._state.discard()
+        return self._run_epoch
+
     async def run_until_convergence(self) -> AsyncGenerator[WorkflowEvent, None]:
         """Run the workflow until no more messages are sent.
 
@@ -114,6 +135,12 @@ async def run_until_convergence(self) -> AsyncGenerator[WorkflowEvent, None]:
         to entering this loop, since only the workflow knows how the run was started
         (fresh input vs. checkpoint resume vs. responses).
         """
+        # The epoch is advanced synchronously by Workflow.run() (see
+        # begin_run), before this generator's first pull: the entry checkpoint
+        # staging that precedes this loop already writes state, so claiming
+        # ownership here would leave that window uncovered (#9473 review).
+        my_epoch = self._run_epoch
+
         # Emit any events already produced prior to entering loop
         if await self._ctx.has_events():
             logger.info("Yielding pre-loop events")
@@ -167,7 +194,10 @@ async def run_until_convergence(self) -> AsyncGenerator[WorkflowEvent, None]:
                 self._state.commit()
                 committed = True
             finally:
-                if not committed:
+                if not committed and self._run_epoch == my_epoch:
+                    # Stale finalizers (a successor run has already bumped the
+                    # epoch) leave the pending writes alone: they now belong to
+                    # the successor, which swept ours at its own entry (#9473).
                     self._state.discard()
 
             # Create checkpoint after each superstep iteration

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

@@ -3,6 +3,7 @@
 # ruff:file-ignore[unnecessary-assign-before-yield]
 from __future__ import annotations
 
+import asyncio
 import functools
 import hashlib
 import json
@@ -876,6 +877,14 @@ def run(
         # its async-generator finalizer ran. Clear it so this run starts clean and does
         # not silently inherit the prior run's runtime checkpoint storage.
         self._runner.context.clear_runtime_checkpoint_storage()
+        # A dropped run can also leave uncommitted writes pending on the shared
+        # state (its finalizer, epoch-guarded in the runner, skips the discard
+        # once this run takes over). Claim ownership and sweep them here,
+        # synchronously and before this run stages anything: the entry
+        # checkpoint's executor snapshots land in the same pending map, so the
+        # epoch must already have advanced by the time they are written
+        # (#9473 review).
+        self._runner.begin_run()
         runtime_tools = normalize_tools(tools) if tools is not None else None
         self._runner.context.set_runtime_tools(runtime_tools)
 
@@ -928,6 +937,18 @@ async def _run_core(
         # here it already points at our own ``ResponseStream``.
         my_active_run = self._active_run
 
+        # A dropped predecessor's finalizer may already be scheduled but not yet
+        # finished: if its handler writes state from a cancellation cleanup path
+        # (a `finally` or cancel handler), those writes land in the shared pending
+        # map after run()'s synchronous sweep. Yield once so the loop can drive the
+        # scheduled finalizer, then sweep again before this run stages its own
+        # writes. A finalizer that has not even been scheduled yet cannot have
+        # written anything; one scheduled but slower than this yield leaves a
+        # narrow residual window, which per-run state isolation would close
+        # (#9473 review).
+        await asyncio.sleep(0)
+        self._runner.state.discard()
+
         # Enable runtime checkpointing if storage provided.
         if checkpoint_storage is not None:
             self._runner.context.set_runtime_checkpoint_storage(checkpoint_storage)

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

@@ -2205,3 +2205,163 @@ async def test_resolve_pause_checkpoint_id_uses_builder_storage() -> None:
     pause_id = workflow.get_last_checkpoint_id()
     resolved = await workflow.resolve_pause_checkpoint_id({"approval-1"})
     assert resolved == pause_id
+
+
+@dataclass
+class _StagingMessage:
+    run: str
+
+
+async def test_workflow_dropped_stream_finalizer_preserves_successor_state_writes() -> None:
+    """A dropped run's late finalizer must not discard the next run's state writes.
+
+    Repro for #9473:
+
+    1. Run A stages ``a_key`` and blocks mid-superstep; its stream is dropped,
+       so ``aclose`` is scheduled on the loop but has not run yet.
+    2. Run B starts on the same workflow before that finalizer executes and
+       stages ``b_key``.
+    3. The loop runs A's stale finalizer. Without the epoch guard its
+       ``State.discard()`` clears B's pending ``b_key`` as well, and B still
+       finishes IDLE with the write silently gone.
+    """
+    a_staged = asyncio.Event()
+    b_staged = asyncio.Event()
+    release_b = asyncio.Event()
+
+    class StagingExecutor(Executor):
+        @handler
+        async def handle(self, message: _StagingMessage, ctx: WorkflowContext[_StagingMessage, str]) -> None:
+            if message.run == "A":
+                ctx.set_state("a_key", "from-run-A")
+                a_staged.set()
+                await asyncio.Event().wait()  # never released; run A is abandoned here
+            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=StagingExecutor(id="staging_exec")).build()
+
+    # Step 1: drive stream A until its executor has been invoked, then staged
+    # a_key and blocked inside the superstep.
+    stream_a = workflow.run(_StagingMessage(run="A"), stream=True)
+    aiter_a = stream_a.__aiter__()
+    while True:
+        event = await aiter_a.__anext__()
+        if event.type == "executor_invoked":
+            break
+    await a_staged.wait()
+
+    # Step 2: drop stream A; GC schedules the async-gen close without running it.
+    del stream_a
+    del aiter_a
+    del event
+    gc.collect()
+
+    # Step 3: start run B before A's finalizer executes. B sweeps A's leftover
+    # write at entry, then stages its own.
+    async def consume_b() -> tuple[list[str], WorkflowRunState | None]:
+        outputs: list[str] = []
+        final_state: WorkflowRunState | None = None
+        async for ev in workflow.run(_StagingMessage(run="B"), stream=True):
+            if ev.type == "output":
+                outputs.append(ev.data)
+            elif ev.type == "status":
+                final_state = ev.state
+        return outputs, final_state
+
+    task_b = asyncio.create_task(consume_b())
+    await asyncio.wait_for(b_staged.wait(), 5)
+
+    # Step 4: let the loop drive A's stale finalizer through GeneratorExit.
+    for _ in range(20):
+        await asyncio.sleep(0.01)
+
+    # Step 5: release B; it converges and commits its superstep.
+    release_b.set()
+    outputs, final_state = await asyncio.wait_for(task_b, 5)
+
+    assert outputs == ["b-done"]
+    assert final_state == WorkflowRunState.IDLE
+    state = workflow._runner.state.export_state()  # pyright: ignore[reportPrivateUsage]
+    assert state.get("b_key") == "from-run-B"
+    # A's unfinished write was swept at B's entry, not committed into the state.
+    assert "a_key" not in state
+
+
+async def test_workflow_dropped_run_cancellation_cleanup_writes_do_not_leak() -> None:
+    """A dropped run whose handler writes from a cancellation cleanup path must
+    not leak those writes into the successor's committed state (#9477 review).
+
+    The successor's synchronous sweep in run() happens after the predecessor's
+    finalizer has completed here, which is the dominant interleaving: the late
+    cleanup write lands in the shared pending map and is swept before the
+    successor stages anything.
+    """
+    a_staged = asyncio.Event()
+    a_cleaned = asyncio.Event()
+
+    class CleanupWriter(Executor):
+        @handler
+        async def handle(self, message: _StagingMessage, ctx: WorkflowContext[_StagingMessage, str]) -> None:
+            try:
+                if message.run == "A":
+                    ctx.set_state("a_key", "from-run-A")
+                    a_staged.set()
+                    await asyncio.Event().wait()
+            finally:
+                if message.run == "A":
+                    # Cancellation cleanup writing state, the shape the review
+                    # flagged: this lands in pending after the handler is torn down.
+                    ctx.set_state("a_cleanup", "from-run-A-cleanup")
+                    a_cleaned.set()
+
+    workflow = WorkflowBuilder(start_executor=CleanupWriter(id="cleanup_writer")).build()
+
+    stream_a = workflow.run(_StagingMessage(run="A"), stream=True)
+    aiter_a = stream_a.__aiter__()
+    while True:
+        event = await aiter_a.__anext__()
+        if event.type == "executor_invoked":
+            break
+    await a_staged.wait()
+    del stream_a
+    del aiter_a
+    del event
+    gc.collect()
+
+    # Drive the loop until A's finalizer has cancelled the handler and its
+    # finally block has written the cleanup key.
+    for _ in range(50):
+        await asyncio.sleep(0.01)
+        if a_cleaned.is_set():
+            break
+    assert a_cleaned.is_set(), "predecessor's cancellation cleanup never ran"
+
+    result_b = await workflow.run(_StagingMessage(run="B"))
+    assert result_b.get_final_state() == WorkflowRunState.IDLE
+    state = workflow._runner.state.export_state()  # pyright: ignore[reportPrivateUsage]
+    assert "a_key" not in state
+    assert "a_cleanup" not in state
+
+
+async def test_workflow_run_advances_runner_epoch_synchronously() -> None:
+    """run() must advance the runner epoch before returning, synchronously.
+
+    Regression pin for the #9477 review finding: the entry checkpoint stages
+    executor snapshots before run_until_convergence's generator body first
+    executes, so an epoch claimed there would leave the staging window
+    unguarded; a dropped predecessor's finalizer running mid-staging could then
+    match the epoch and discard the successor's staged snapshots. Claiming it
+    in run()'s synchronous section closes the window.
+    """
+    workflow = WorkflowBuilder(start_executor=IncrementExecutor(id="epoch_exec", limit=2, increment=1)).build()
+    epoch_before = workflow._runner._run_epoch  # pyright: ignore[reportPrivateUsage]
+    stream = workflow.run(NumberMessage(data=0), stream=True)
+    # No await between run() and the assertion: the epoch must already have
+    # advanced before this run stages or commits anything.
+    assert workflow._runner._run_epoch == epoch_before + 1  # pyright: ignore[reportPrivateUsage]
+    async for _ in stream:
+        pass