From 53a4fc2241f657726b669099440009a4a80af499 Mon Sep 17 00:00:00 2001 From: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Date: Sun, 18 Oct 2026 06:13:04 +0800 Subject: [PATCH 1/2] Python: fix(core): keep a dropped run's late finalizer from discarding the next run's state writes A dropped workflow stream's async-generator finalizer runs after the caller has already started a successor run on the same instance. The superstep cleanup then discarded the shared State's pending writes, which by then belonged to the successor: the new run still finished IDLE with its outputs, but its set_state writes were silently gone (#9473). The runner now stamps each drive with an epoch and only discards pending writes while it is still the latest drive, and run() sweeps a dropped predecessor's leftover pending writes up front, next to the existing stale runtime-checkpoint-storage clearing, so they are neither kept nor folded into the successor's first kwargs commit. --- .../agent_framework/_workflows/_runner.py | 15 +++- .../agent_framework/_workflows/_workflow.py | 6 ++ .../core/tests/workflow/test_workflow.py | 84 +++++++++++++++++++ 3 files changed, 104 insertions(+), 1 deletion(-) diff --git a/python/packages/core/agent_framework/_workflows/_runner.py b/python/packages/core/agent_framework/_workflows/_runner.py index f41183cb1c2..10e8d0412c1 100644 --- a/python/packages/core/agent_framework/_workflows/_runner.py +++ b/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 @@ -114,6 +121,9 @@ 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). """ + self._run_epoch += 1 + 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 +177,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 diff --git a/python/packages/core/agent_framework/_workflows/_workflow.py b/python/packages/core/agent_framework/_workflows/_workflow.py index c60617dabd5..c591262f43b 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/python/packages/core/agent_framework/_workflows/_workflow.py @@ -876,6 +876,12 @@ 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). Sweep them here, before this run stages and + # commits its own run kwargs, so the leftover writes are neither kept + # nor folded into this run's first commit (#9473). + self._runner.state.discard() runtime_tools = normalize_tools(tools) if tools is not None else None self._runner.context.set_runtime_tools(runtime_tools) diff --git a/python/packages/core/tests/workflow/test_workflow.py b/python/packages/core/tests/workflow/test_workflow.py index 76a285e6d93..8f3104d4176 100644 --- a/python/packages/core/tests/workflow/test_workflow.py +++ b/python/packages/core/tests/workflow/test_workflow.py @@ -2205,3 +2205,87 @@ 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 From e9ed6844466fdc7277cc8e8ecc4ca7ade60c260f Mon Sep 17 00:00:00 2001 From: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Date: Sun, 18 Oct 2026 17:29:01 +0800 Subject: [PATCH 2/2] Python: fix(core): claim the run epoch before the successor stages state Copilot review on #9477 found the epoch advanced too late: run_until_convergence is an async generator whose entry code runs on the first pull, but the successor's entry checkpoint stages executor snapshots before that, so a dropped run's finalizer could still match the epoch and discard the staged snapshots mid-flight. The epoch now advances in Workflow.run()'s synchronous section via Runner.begin_run(), before any await point, and the runner's loop reads it. Also sweep once more after a loop yield at _run_core's start: a dropped predecessor whose handler writes state from a cancellation cleanup path can otherwise land those writes after the synchronous sweep. A residual window remains for a finalizer slower than that yield; per-run state isolation would close it for good. --- .../agent_framework/_workflows/_runner.py | 19 ++++- .../agent_framework/_workflows/_workflow.py | 23 +++++- .../core/tests/workflow/test_workflow.py | 76 +++++++++++++++++++ 3 files changed, 113 insertions(+), 5 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_runner.py b/python/packages/core/agent_framework/_workflows/_runner.py index 10e8d0412c1..8a0e08a70c2 100644 --- a/python/packages/core/agent_framework/_workflows/_runner.py +++ b/python/packages/core/agent_framework/_workflows/_runner.py @@ -112,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. @@ -121,7 +135,10 @@ 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). """ - self._run_epoch += 1 + # 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 diff --git a/python/packages/core/agent_framework/_workflows/_workflow.py b/python/packages/core/agent_framework/_workflows/_workflow.py index c591262f43b..111f8c7187f 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/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 @@ -878,10 +879,12 @@ def run( 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). Sweep them here, before this run stages and - # commits its own run kwargs, so the leftover writes are neither kept - # nor folded into this run's first commit (#9473). - self._runner.state.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) @@ -934,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) diff --git a/python/packages/core/tests/workflow/test_workflow.py b/python/packages/core/tests/workflow/test_workflow.py index 8f3104d4176..28ada97a58a 100644 --- a/python/packages/core/tests/workflow/test_workflow.py +++ b/python/packages/core/tests/workflow/test_workflow.py @@ -2289,3 +2289,79 @@ async def consume_b() -> tuple[list[str], WorkflowRunState | None]: 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