[BREAKING] Python: Allow workflow checkpoint full replayability - #7347
[BREAKING] Python: Allow workflow checkpoint full replayability#7347Tao Chen (TaoChenOSU) wants to merge 5 commits into
Conversation
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). |
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 |
| 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?
There was a problem hiding this comment.
I don't think the source id of the initial message is ever a public contract. If we look at the comment we made on the source_executor_ids property of the WorkflowContext, it states that "This is a list to support fan_in scenarios where multiple sources send aggregated messages to the same executor.". Customers are not expected to be using it in the start executor.
That being said, this is nonetheless still a breaking change on our side, but I think this is acceptable if we include it in our change log.
And I have exported the INTERNAL_SOURCE_ID.
| # 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.
There was a problem hiding this comment.
Good call! This is a gap that already exists on checkpoint restoration today. Issue created to track separately: #7371
| # 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.
|
if we are introducign breaking changes here, we should at least have remediation plans and samples that show how to deal with this change! |
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.