Skip to content

[BREAKING] Python: Allow workflow checkpoint full replayability - #7347

Closed
Tao Chen (TaoChenOSU) wants to merge 5 commits into
mainfrom
fix/allow-workflow-checkpoint-full-replayability
Closed

[BREAKING] Python: Allow workflow checkpoint full replayability#7347
Tao Chen (TaoChenOSU) wants to merge 5 commits into
mainfrom
fix/allow-workflow-checkpoint-full-replayability

Conversation

@TaoChenOSU

@TaoChenOSU Tao Chen (TaoChenOSU) commented Jul 27, 2026

Copy link
Copy Markdown
Contributor

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_checkpoint flag + 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

  • What are the major changes?
    • Record the first message. The initial input is now seeded through the start executor's existing internal self-edge 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.
    • Runner simplification. run_until_convergence now only checkpoints after each superstep. The pre-loop entry checkpoint and the _resumed_from_checkpoint flag are removed; entry-checkpoint creation is now the Workflow's responsibility.
    • Response-entry checkpoint. Delivering 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, and reset() (backed by RunnerImpl.reset_to_checkpoint) rewinds an instance to it so a single Workflow can be reused across independent runs.
    • Deprecated the unused AgentExecutor.reset().
  • What is the impact of these changes?
    • Every run performs one additional superstep (the start executor is now a superstep): checkpoint iteration_count shifts by +1, each run emits one extra superstep_started/superstep_completed pair, and the max_iterations boundary shifts by one. This is a minor and low risk breaking change.
    • iteration_count is no longer globally unique across a HIL lifecycle — the response-entry checkpoint shares the pending-request checkpoint's iteration. Checkpoint ordering is defined by the previous_checkpoint_id lineage (and timestamp), not iteration_count (documented on WorkflowCheckpoint).
    • Not breaking for persisted checkpoints: the change reuses existing internal self-edges, so the graph signature hash is unchanged and existing checkpoints still validate and restore.
  • What do you want reviewers to focus on?
    • The input seeding + entry-checkpoint placement in Workflow._execute_with_message_or_checkpoint and _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

  • The code builds clean without any errors or warnings
  • All unit tests pass, and I have added new tests where possible
  • The PR follows the Contribution Guidelines
  • This PR is linked to an issue and there is no other open PR for this issue (see Related Issue above).
  • This is not a breaking change. If it is a breaking change, add the breaking change label (or add "[BREAKING]" to the title prefix, before or after any language prefix) — a workflow keeps the label and title prefix in sync automatically.

@TaoChenOSU Tao Chen (TaoChenOSU) self-assigned this Jul 27, 2026
Copilot AI review requested due to automatic review settings July 27, 2026 17:38
@TaoChenOSU Tao Chen (TaoChenOSU) added python Usage: [Issues, PRs], Target: Python workflows Usage: [Issues, PRs], Target: Workflows labels Jul 27, 2026
@github-actions github-actions Bot changed the title Allow workflow checkpoint full replayability Python: Allow workflow checkpoint full replayability Jul 27, 2026
@TaoChenOSU Tao Chen (TaoChenOSU) changed the title Python: Allow workflow checkpoint full replayability [BREAKING] Python: Allow workflow checkpoint full replayability Jul 27, 2026
@agent-framework-automation agent-framework-automation Bot added the breaking change Usage: [PRs], Target: all PRs that introduce changes that are not backward compatible label Jul 27, 2026
@github-actions

github-actions Bot commented Jul 27, 2026

Copy link
Copy Markdown
Contributor

Python Test Coverage

Python Test Coverage Report •
FileStmtsMissCoverMissing
packages/ag-ui/agent_framework_ag_ui
   _workflow_run.py6857389%78, 92, 94, 96, 99, 226–229, 278, 289, 294, 319, 355–358, 386, 391, 419, 429, 440, 445, 448, 461, 471, 474, 479, 482, 497–499, 504, 506–507, 511, 513, 530, 536–537, 547–548, 552–553, 577–578, 611, 619, 688, 708, 724, 739, 840–854, 886–887, 908–909, 984, 1046
packages/core/agent_framework/_workflows
   _agent_executor.py2101891%167, 191, 232, 256, 276–277, 357–359, 361, 376, 381–382, 505, 532–533, 605, 611
   _checkpoint.py158199%309
   _runner.py191497%390–392, 402
   _runner_context.py1741690%67, 81–82, 84–85, 87, 453, 470, 483, 491, 494–496, 541, 554, 558
   _workflow.py3932693%62, 64, 69, 93, 98, 159, 195, 429–431, 433–434, 458, 492, 620, 659, 956, 977, 1025, 1037, 1043, 1048, 1074–1076, 1319
TOTAL45635447490% 

Python Unit Test Overview

Tests Skipped Failures Errors Time
9426 34 💤 0 ❌ 0 🔥 2m 25s ⏱️

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 introduce Workflow.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).

Comment thread python/packages/core/agent_framework/_workflows/_workflow.py
Comment thread python/packages/core/agent_framework/_workflows/_workflow.py Outdated

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Automated Code Review

Reviewers: 5 | Confidence: 83% | Result: All clear

Reviewed: Correctness, Security Reliability, Test Coverage, Failure Modes, Design Approach


Automated review by TaoChenOSU's agents

@TaoChenOSU
Tao Chen (TaoChenOSU) marked this pull request as ready for review July 27, 2026 21:34

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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=True branch in _mark_resumed lacks a direct unit test (only tested indirectly via Workflow.reset()), and the AgentExecutor.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 through WorkflowExecutor, 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 embedded WorkflowExecutor restores still call restore_checkpoint(sub_workflow_checkpoint) without start_new_lineage (_workflow_executor.py:469). Because sub-workflows are run normally via self.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

Comment on lines +1284 to +1286
await self._runner.restore_checkpoint(
copy.deepcopy(self._initial_checkpoint),
start_new_lineage=True,

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@github-actions

Copy link
Copy Markdown
Contributor

Flagged issue

Workflow.reset() starts a new lineage only for the top-level runner (_workflow.py:1284-1286), but embedded WorkflowExecutor restores still call restore_checkpoint(sub_workflow_checkpoint) without start_new_lineage (_workflow_executor.py:469). Because sub-workflows are run normally via self.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.


Source: automated DevFlow PR review

Comment thread python/packages/core/agent_framework/_workflows/_workflow.py Outdated
await self._runner.context.send_message(
WorkflowMessage(
data=message,
source_id=INTERNAL_SOURCE_ID(start_id),

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread python/packages/core/agent_framework/_workflows/_workflow.py
Comment thread python/packages/core/agent_framework/_workflows/_workflow.py
# 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(

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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,

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@eavanvalkenburg

Copy link
Copy Markdown
Member

if we are introducign breaking changes here, we should at least have remediation plans and samples that show how to deal with this change!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

breaking change Usage: [PRs], Target: all PRs that introduce changes that are not backward compatible python Usage: [Issues, PRs], Target: Python workflows Usage: [Issues, PRs], Target: Workflows

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants