[BREAKING] Python: Allow workflow checkpoint full replayability - #7347
[BREAKING] Python: Allow workflow checkpoint full replayability#7347TaoChenOSU wants to merge 3 commits into
Conversation
TaoChenOSU
added
python
GitHub Actions
Inactive
GitHub Actions
Inactive
GitHub Actions
Inactive
GitHub Actions
Inactive
GitHub Actions
Inactive
github-actions
Bot
changed the title
GitHub Actions
Inactive
GitHub Actions
Inactive
TaoChenOSU
changed the title
GitHub Actions
Inactive
agent-framework-automation
Bot
added
the
breaking change
GitHub Actions
Inactive
Python Test Coverage Report •
Python Unit Test Overview
|
|||||||||||||||||||||||||||||||||||||||||||||||||||||||
There was a problem hiding this comment.
Pull request overview
This PR updates the Python workflow engine’s checkpointing semantics so checkpoints can fully replay a run from its original input (including the start executor), simplifies runner responsibilities, and adds the ability to reuse a Workflow instance via a new reset() API.
Changes:
- Move the start executor execution into the normal superstep loop by seeding the initial input as an internal message and writing an iteration-0 “entry” checkpoint before any executors run.
- Simplify
RunnerImpl.run_until_convergence()to checkpoint only after each completed superstep (Workflow now owns entry/response-entry checkpoint creation). - Add a response-entry checkpoint for
responses=...continuations and introduceWorkflow.reset()backed by an in-memory pristine snapshot.
Reviewed changes
Copilot reviewed 10 out of 10 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| python/packages/core/agent_framework/_workflows/_workflow.py | Seeds initial input through internal edge + creates entry/response-entry checkpoints; captures pristine snapshot and adds Workflow.reset(). |
| python/packages/core/agent_framework/_workflows/_runner.py | Removes resume-flag logic; runner now checkpoints only after each superstep and maintains lineage via _previous_checkpoint_id. |
| python/packages/core/agent_framework/_workflows/_checkpoint.py | Documents that iteration_count is not globally unique across HIL lifecycles; ordering is by lineage/timestamp. |
| python/packages/core/agent_framework/_workflows/_agent_executor.py | Updates from_str empty-cache warning guard to align with new internal-source start message semantics; deprecates AgentExecutor.reset(). |
| python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py | Ensures any open assistant text message is closed before emitting the terminal RUN_FINISHED event. |
| python/packages/core/tests/workflow/test_workflow.py | Updates event-count expectations to account for the start executor now running as superstep 1. |
| python/packages/core/tests/workflow/test_runner.py | Updates/rewrites tests for new checkpoint lineage and removes tests tied to the old resume-flag/entry-checkpoint behavior. |
| python/packages/core/tests/workflow/test_checkpoint.py | Adds coverage for entry checkpoint input capture/replay, one-checkpoint-per-superstep lineage, and response-entry checkpoint replay. |
| python/packages/core/tests/workflow/test_workflow_reset.py | Adds tests for Workflow.reset() behavior and concurrency guard interaction. |
| python/packages/core/tests/core/test_function_invocation_logic.py | Minor formatting-only change (blank line). |
GitHub Actions
Inactive
GitHub Actions
Inactive
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
Automated Code Review
Reviewers: 5 | Confidence: 89%
✓ Correctness
This PR cleanly refactors checkpoint creation responsibility from the runner to the workflow, making all runs fully replayable. The entry checkpoint is correctly created after seding the message but before any executor runs, the response-entry checkpoint captures HIL responses in-flight, and the reset() mechanism properly uses deep-copy isolation and fresh lineage to avoid dangling checkpoint references. The internal self-edge routing for the start executor is correctly wired through the existing InternalEdgeGroup infrastructure, and the from_str guard in AgentExecutor is updated to match the new source ID. No correctness bugs found.
✓ Security Reliability
This PR makes a clean structural refactor that moves entry-checkpoint creation from the runner to the workflow, removes the fragile _resumed_from_checkpoint flag, and adds Workflow.reset() using deep-copied in-memory snapshots. The changes are well-protected: copy.deepcopy isolates the pristine snapshot from mutations, start_new_lineage=True prevents dangling parent references in persisted checkpoints, and the ag-ui fix ensures proper event ordering. The removal of the finally block is safe because the state it cleaned up (_resumed_from_checkpoint) no longer exists. No injection risks, resource leaks, missing validation, or unhandled failure modes were identified.
✓ Test Coverage
Test coverage for the core behavioral changes (entry checkpoints, response-entry checkpoints, per-superstep lineage, and Workflow.reset()) is thorough with well-structured integration and unit tests. Minor gaps exist: the new
start_new_lineage=Truebranch in_mark_resumedlacks a direct unit test (only tested indirectly via Workflow.reset()), and theAgentExecutor.reset()deprecation warning has no test verifying the warning is emitted.
✓ Failure Modes
The PR makes well-structured changes to workflow checkpointing with no new silent failure paths, lost errors, or stale-state hazards. The removed try/finally is safe because the flag it cleaned up was removed entirely. Entry checkpoint and response-entry checkpoint creation use the existing non-fatal checkpoint pattern. The reset() mechanism correctly uses start_new_lineage to avoid dangling parent references. No concrete failure-mode bugs found.
✗ Design Approach
I found one design-level correctness gap in the new reset/replay approach. The top-level workflow now correctly starts a fresh checkpoint lineage after
reset(), but that new-lineage behavior is not propagated into embedded sub-workflows restored throughWorkflowExecutor, so nested workflows with their own checkpoint storage can still emit checkpoints that point at an in-memory, non-persisted parent ID after a parent reset.
Flagged Issues
-
Workflow.reset()starts a new lineage only for the top-level runner (_workflow.py:1284-1286), but embeddedWorkflowExecutorrestores still callrestore_checkpoint(sub_workflow_checkpoint)withoutstart_new_lineage(_workflow_executor.py:469). Because sub-workflows are run normally viaself.workflow.run(...)and fresh runs now persist an entry checkpoint before execution (_workflow.py:672-674), a wrapped workflow with its own checkpoint storage will recreate the same dangling-lineage bug this PR fixed at the top level: its next persisted checkpoint will chain to the in-memory reset snapshot ID rather than to a stored checkpoint.
Automated review by TaoChenOSU's agents
| await self._runner.restore_checkpoint( | ||
| copy.deepcopy(self._initial_checkpoint), | ||
| start_new_lineage=True, |
There was a problem hiding this comment.
This only fixes the dangling-parent problem for the top-level workflow. During reset(), executor state restoration still routes embedded sub-workflows through WorkflowExecutor.on_checkpoint_restore(), which calls self.workflow._runner.restore_checkpoint(sub_workflow_checkpoint) without start_new_lineage (_workflow_executor.py:469). If the wrapped workflow also has checkpoint storage, its next fresh run will create an entry checkpoint chained to that in-memory snapshot ID, recreating the same broken lineage one level down. The new-lineage semantics need to propagate into nested workflow restores as well.
|
Flagged issue
Source: automated DevFlow PR review |
| source_span_ids=None, | ||
| # Seed the initial input through the start executor's internal self-edge. | ||
| start_id = self.start_executor_id | ||
| await self._runner.context.send_message( |
There was a problem hiding this comment.
What happens when the initial input doesn't match any of the start executor's handlers? Previously run(x) invoked the executor directly and _find_handler raised RuntimeError for an unhandleable type; now the seeded message goes through the edge runner's _can_handle check and is dropped as DROPPED_TYPE_MISMATCH, so the run converges to IDLE with zero outputs and no error. Could we validate the input against the start executor's handlers at seeding time so callers still get a loud failure?
I tasted this: a start executor handling only str given wf.run(12345) raises on main but returns an empty successful run on this branch.
| await self._runner.context.send_message( | ||
| WorkflowMessage( | ||
| data=message, | ||
| source_id=INTERNAL_SOURCE_ID(start_id), |
There was a problem hiding this comment.
Is the change of the initial message's observable source id from "Workflow" to internal:<start_id> intentional as a public break? Custom start executors that branch on ctx.get_source_executor_id() or source_executor_ids (the exact pattern AgentExecutor itself used, updated in this PR at _agent_executor.py:255) will now silently misclassify the initial input, and INTERNAL_SOURCE_ID isn't exported from agent_framework, so affected code can't cleanly match the new sentinel. Should we export the helper (or the prefix) and add this to the breaking-change notes with migration guidance?
| try: | ||
| # Capture the pristine initial checkpoint once, before any executor runs, | ||
| # so this instance can later be reset back to this state via ``reset()``. | ||
| await self._capture_initial_checkpoint_if_needed() |
There was a problem hiding this comment.
Should this capture be lazy or gated? It runs on every workflow's first run with no has_checkpointing() check, so every executor's on_checkpoint_save() now fires at run start (raising WorkflowCheckpointException on any failure) plus a full state export and deepcopy, even for workflows that never enabled checkpointing and never call reset(). A hook that raises, or snapshot state that isn't deepcopyable (client handles, locks), hard-fails runs that work today, and runner contexts whose build_checkpoint raises NotImplementedError (durabletask, azurefunctions) can never run at all. Could we capture on first reset() instead, or make the capture best-effort with reset() disabled when it fails?
I also tested this: an executor whose on_checkpoint_save raises runs fine on main without checkpointing, but every run() fails with WorkflowCheckpointException. This also contradicts the runner's own policy that checkpoint failure does not fail the run.
| Raises: | ||
| WorkflowException: If a run is currently active on this instance. | ||
| """ | ||
| if self._is_run_active(): |
There was a problem hiding this comment.
Could reset() reserve the instance synchronously the way run() does? We don't support concurrent operations on a workflow object, and run() enforces that with a synchronous guard before its first await, but reset() only checks _is_run_active() once and then yields inside executor restore hooks. A run() started in that window passes its own guard (reset isn't a run) and executes against partially restored state (I could reproduce silent state corruption with no exception). Since reset() targets instance reuse in long-lived services, could it hold the same exclusivity for its full duration so this misuse raises like the other orderings do?
| # snapshot and subsequent resets remain repeatable. Start a new lineage: the | ||
| # pristine snapshot is an in-memory checkpoint that was never persisted, so | ||
| # checkpoints created by the next run must not chain to its (storage-absent) id. | ||
| await self._runner.restore_checkpoint( |
There was a problem hiding this comment.
Could reset() also rewind the edge-runner state? FanInEdgeRunner retains unmatched inputs in _buffer, and the runner can consider the workflow converged while that buffer is still populated. Restoring this snapshot leaves the buffer intact, so a branch value from before reset can be combined with a branch value from the supposedly independent next run.
| # checkpoints created by the next run must not chain to its (storage-absent) id. | ||
| await self._runner.restore_checkpoint( | ||
| copy.deepcopy(self._initial_checkpoint), | ||
| start_new_lineage=True, |
There was a problem hiding this comment.
Should reset clear queued workflow events before applying the pristine checkpoint? InProcRunnerContext.apply_checkpoint() replaces messages and pending requests but leaves its event queue untouched, so an unconsumed event from a cancelled or abandoned streaming run is returned as part of the next independent run.
I reproduced reset() followed by run("new") returning ["old-2", "new-1", "new-2"], which can leak prior-run output across reuse boundaries.
Motivation & Context
Workflow checkpoints could not fully replay a run. The start executor was invoked outside the superstep/checkpoint loop, so the earliest ("entry") checkpoint captured the start executor's output messages and post-run state — never the raw input. Restoring the entry checkpoint therefore replayed from superstep 1 onward; the start executor and the original input were unrecoverable from any checkpoint. Human-in-the-loop (HIL) continuations had the same gap: responses were delivered and processed without ever being recorded, so the superstep that consumes them was not replayable.
This also made the runner's checkpointing convoluted: it decided whether to create an "entry" checkpoint based on how the workflow was invoked (a
_resumed_from_checkpointflag + iteration guards), mixing a Workflow-level concern into the runner whose only job should be running supersteps.This PR makes checkpoints fully replayable, gives the runner a single narrow checkpoint responsibility, and (building on the same pristine-snapshot capability) adds a way to reuse a Workflow instance.
Description & Review Guide
INTERNAL_SOURCE_ID(start_id), and the Workflow creates the entry checkpoint (iteration 0) capturing that input before any executor runs. The start executor now runs inside superstep 1 like any other executor, so fresh-run and checkpoint-resume paths are identical and a run is fully replayable from its input.run_until_convergencenow only checkpoints after each superstep. The pre-loop entry checkpoint and the_resumed_from_checkpointflag are removed; entry-checkpoint creation is now the Workflow's responsibility.responses=...now records a checkpoint capturing the responses in-flight before they are processed, so HIL continuations are fully replayable too.Workflow.reset(). A pristine initial checkpoint is captured on the first run, andreset()(backed byRunnerImpl.reset_to_checkpoint) rewinds an instance to it so a single Workflow can be reused across independent runs.AgentExecutor.reset().iteration_countshifts by +1, each run emits one extra superstep_started/superstep_completed pair, and themax_iterationsboundary shifts by one. This is a minor and low risk breaking change.iteration_countis no longer globally unique across a HIL lifecycle — the response-entry checkpoint shares the pending-request checkpoint's iteration. Checkpoint ordering is defined by theprevious_checkpoint_idlineage (andtimestamp), notiteration_count(documented onWorkflowCheckpoint).Workflow._execute_with_message_or_checkpointand_send_responses_internal; the runner simplification in run_until_convergence/_mark_resumed; and the iteration_count non-uniqueness semantics.Related Issue
Part of the multi-PR workflow engine refactor series (follows #6695, #6776, and #7097). No standalone tracking issue.
Contribution Checklist
breaking changelabel (or add "[BREAKING]" to the title prefix, before or after any language prefix) — a workflow keeps the label and title prefix in sync automatically.