Repository navigation
Python: fix(core): keep a dropped run's late finalizer from discarding the next run's state writes - #9477
Conversation
GitHub Actions
Active
GitHub Actions
Active
agent-framework-automation
Bot
added
the
python
There was a problem hiding this comment.
🟡 Changes recommended
Remaining ownership races can commit abandoned cleanup writes or discard successor checkpoint snapshots.
2 open issues
What changed in this PR
Addresses #9473 by protecting successor workflow state from a dropped stream’s late finalizer.
Changes:
- Guards runner cleanup with a per-drive epoch.
- Clears abandoned pending writes before successor setup.
- Adds a regression test for delayed stream finalization.
| File | Description |
|---|---|
| python/packages/core/tests/workflow/test_workflow.py | Tests successor writes, outputs, status, and abandoned-write removal. |
| python/packages/core/agent_framework/_workflows/_workflow.py | Sweeps pending state at run entry. |
| python/packages/core/agent_framework/_workflows/_runner.py | Guards pending-state discard with an epoch check. |
Review effort: Balanced
| self._run_epoch += 1 | ||
| my_epoch = self._run_epoch |
There was a problem hiding this comment.
Fixed in c4d09dd, and the mechanism moved earlier than you suggested. The epoch is now claimed by Runner.begin_run() from Workflow.run()'s synchronous section, before any await point, so by the time _save_executor_states() stages the first snapshot the stale finalizer's epoch check already fails. run_until_convergence reads the epoch instead of advancing it. Regression pin: test_workflow_run_advances_runner_epoch_synchronously asserts the epoch has advanced before run() returns (red if the claim moves back into the generator).
| # 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() |
There was a problem hiding this comment.
Partially addressed in c4d09dd, with the residual stated honestly. _run_core now yields once at its start so the loop can drive the already-scheduled finalizer, then sweeps the pending map again before this run stages anything; a regression test (..._cancellation_cleanup_writes_do_not_leak) covers the dominant interleaving where the predecessor's finally-block write lands before the successor's sweep. The residual window you describe is real: a finalizer slower than that single yield can still leak a cleanup write into the successor's first commit. Closing it for good needs per-run isolation of pending writes (keying the pending map by run epoch, or draining the predecessor through a tracked handle), which is a bigger State-layer change than this PR's scope; I would rather land that separately if you agree it is worth doing.
Copilot review on microsoft#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.
GitHub Actions
Active
…g 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 (microsoft#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.
Copilot review on microsoft#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.
c4d09dd to
e9ed684
Compare
GitHub Actions
Active
|
Rebased onto current main (0200599, which includes #8819) — the two changes interleave cleanly. Worth stating the relationship precisely, since #8819 works the same failure-adjacent area: it discards pending state when a superstep fails (the Re-verified after the rebase: |
Why
Fixes #9473. Dropping a workflow stream mid-superstep and starting a new run on the same
Workflowbefore the dropped stream's async-generator finalizer executes silently loses the new run's state writes: the stale finalizer'sState.discard()clears the shared pending map, which by then holds the successor's writes. The successor still finishesIDLEwith its outputs; the write is just gone.What changed
Two cooperating pieces, matching the identity-check defense
_run_corealready uses for_active_run:Runnerstamps each drive ofrun_until_convergencewith an epoch, and the superstepfinallyonly discards pending writes while it is still the latest drive. A stale finalizer (the successor has already bumped the epoch) leaves the pending map alone.Workflow.run()sweeps a dropped predecessor's leftover pending writes synchronously, next to the existing stale runtime-checkpoint-storage clearing. This has to happen inrun(), not at the runner's entry:_run_workflow_with_tracingcommits this run's kwargs before the runner loop starts, and without the sweep that commit would fold the predecessor's leftover writes into the committed state.Verification
test_workflow_dropped_stream_finalizer_preserves_successor_state_writesintest_workflow.py, following the reported repro's exact timing (iterate run A's stream toexecutor_invoked, drop +gc.collect, start run B before the finalizer runs, let the finalizer run, then finish B). It asserts B's write survives, B's outputs and IDLE state are intact, and A's never-finished write is swept rather than committed. Fails on unpatched main with exactly the reported symptom (b_keymissing); passes here.State.commitconfirms the ordering: the successor's kwargs commit no longer carries the abandoneda_key, and the stale finalizer performs no discard.test_workflow.py,test_runner.py,test_state.py,test_workflow_states.pyand the checkpoint storage/conformance files; ruff check/format clean; pyright on the two touched modules shows no new errors (the pre-existing baseline is unchanged).Verified locally on macOS, Python 3.14.2, CPU only.