[BREAKING] Python: Allow workflow checkpoint full replayability

[Tao Chen (TaoChenOSU)]

Motivation & Context

Supersedes #7347 with reset() stripped out to a separate PR.

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.

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.
  • Deprecated the unused AgentExecutor.reset() and RunnerContext.reset_for_new_run().
• 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.

[github-actions]

Python Test Coverage Report •

File	Stmts	Miss	Cover	Missing
packages/ag-ui/agent_framework_ag_ui
_workflow_run.py	685	73	89%	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.py	208	17	91%	166, 190, 231, 255, 275–276, 356–358, 360, 370–371, 494, 521–522, 594, 600
_checkpoint.py	158	1	99%	309
_runner.py	191	5	97%	387–389, 398–399
_runner_context.py	174	16	90%	67, 81–82, 84–85, 87, 453, 470, 483, 491, 494–496, 541, 554, 558
_workflow.py	362	25	93%	61, 63, 68, 92, 97, 158, 194, 413–415, 417–418, 442, 476, 604, 643, 934, 955, 1003, 1015, 1021, 1026, 1052–1054
_workflow_builder.py	207	13	93%	223, 500–504, 506, 609, 624, 653, 669, 677, 712
TOTAL	45616	4482	90%

Python Unit Test Overview

Tests	Skipped	Failures	Errors	Time
9421	34 💤	0 ❌	0 🔥	2m 35s ⏱️

[Copilot]

Pull request overview

Enables fully replayable workflow checkpointing by ensuring the initial input and human-in-the-loop (responses) deliveries are captured as checkpoints before the corresponding consuming supersteps run, and by simplifying the runner’s checkpoint responsibility to “after each superstep”.

Changes:

• Seed initial input via the start executor’s internal self-edge and create an iteration-0 entry checkpoint before any executor runs.
• Simplify RunnerImpl.run_until_convergence() to checkpoint only after each completed superstep; resume bookkeeping now relies on _previous_checkpoint_id.
• Add response-entry checkpointing (checkpoint responses in-flight before they’re processed) and strengthen restore behavior by dropping stale queued events in apply_checkpoint().

Reviewed changes

Copilot reviewed 13 out of 13 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
python/packages/core/tests/workflow/test_workflow.py	Updates fan-in/out event-count expectations and adds regression tests for input-type handling and WorkflowMessage unwrapping.
python/packages/core/tests/workflow/test_runner.py	Updates resumed-run assertions to use _previous_checkpoint_id and replaces resumed-flag tests with per-superstep checkpoint lineage checks.
python/packages/core/tests/workflow/test_request_info_event_rehydrate.py	Adds tests ensuring apply_checkpoint() drops stale events and preserves the streaming flag.
python/packages/core/tests/workflow/test_checkpoint.py	Adds end-to-end tests for entry-checkpoint replay, per-superstep checkpoint lineage, and response-entry replay.
python/packages/core/tests/core/test_function_invocation_logic.py	Minor formatting-only change.
python/packages/core/agent_framework/_workflows/_workflow.py	Implements input seeding + entry checkpoint creation and response-entry checkpoint creation.
python/packages/core/agent_framework/_workflows/_runner.py	Removes _resumed_from_checkpoint and narrows runner checkpointing to post-superstep only; resume sets _previous_checkpoint_id.
python/packages/core/agent_framework/_workflows/_runner_context.py	Deprecates reset_for_new_run() and ensures apply_checkpoint() clears stale queued events without resetting streaming.
python/packages/core/agent_framework/_workflows/_checkpoint.py	Documents that iteration_count is not globally unique across HIL lifecycles; lineage is via previous_checkpoint_id.
python/packages/core/agent_framework/_workflows/_agent_executor.py	Aligns “empty cache” warning logic with internal-source seeding (INTERNAL_SOURCE_ID).
python/packages/core/agent_framework/init.pyi	Exposes INTERNAL_SOURCE_ID in the public typing surface.
python/packages/core/agent_framework/init.py	Exposes INTERNAL_SOURCE_ID in runtime exports.
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.

[github-actions]

Flagged issue

Fresh runs can inherit the previous run's checkpoint parent (_workflow.py:675-678, _runner.py:253-267), which contradicts the new "entry checkpoint begins a new lineage" invariant asserted in test_checkpoint.py:476-481.

---

Source: automated DevFlow PR review

[github-actions]

Automated Code Review

Reviewers: 5 | Confidence: 90%

✓ Correctness

The PR is well-implemented. The architectural shift of entry-checkpoint responsibility from runner to workflow is sound. Internal self-edge routing (via InternalEdgeGroup) is an existing mechanism correctly reused for seding initial input. Checkpoint lineage, event queue cleanup on restore, and the response-entry checkpoint mechanics are all correct. The ag-ui fix for draining open messages before RUN_FINISHED is a valid ordering correction. No correctness issues found.

✓ Security Reliability

This PR is a well-designed refactoring that improves workflow checkpoint replayability. The removal of the _resumed_from_checkpoint flag in favor of a simpler _previous_checkpoint_id pointer eliminates a class of state-leakage bugs. The new entry checkpoint and response-entry checkpoint creation paths are correctly placed, input validation is added at the trust boundary, and stale event cleanup in apply_checkpoint prevents data leakage between runs. No security or reliability issues found.

✓ Test Coverage

Test coverage for the core behavioral changes is strong — entry checkpoints, response-entry checkpoints, replay semantics, lineage chains, and apply_checkpoint behaviors all have dedicated tests with meaningful assertions. The main gap is the newly deprecated reset_for_new_run method which now emits a DeprecationWarning but has no test verifying the warning is actually raised. The removal of the try/finally block in the runner is safe because the state it protected (_resumed_from_checkpoint) no longer exists, and existing tests (test_runner_accepts_new_run_after_previous_failure, test_runner_run_until_convergence_not_completed) still cover the failure-recovery paths.

✓ Failure Modes

The PR is well-designed with no significant new failure modes. The removal of _resumed_from_checkpoint and its try/finally is safe since the flag is entirely eliminated. Entry and response-entry checkpoint failures follow the existing best-effort semantics (log-and-continue). The apply_checkpoint event-queue clearing is correctly ordered before the runner loop. The streaming flag preservation through checkpoint restore is explicitly tested. No silent data loss, state leakage, or cancellation race conditions were identified.

✗ Design Approach

I found one design regression in the checkpoint lineage changes: fresh message runs on a reused workflow instance no longer reset the parent-checkpoint pointer, so the new iteration-0 entry checkpoint can be silently chained onto the previous run instead of starting a new lineage.

Flagged Issues

• Fresh runs can inherit the previous run's checkpoint parent (_workflow.py:675-678, _runner.py:253-267), which contradicts the new "entry checkpoint begins a new lineage" invariant asserted in test_checkpoint.py:476-481.

---

Automated review by TaoChenOSU's agents

[github-actions]

Automated Code Review

Reviewers: 5 | Confidence: 89%

✓ Correctness

This PR correctly refactors workflow checkpointing for full replayability. The initial input is now seeded through the pre-existing internal self-edge mechanism (INTERNAL_SOURCE_ID), the entry checkpoint properly captures it at iteration 0 before any executor runs, and the response-entry checkpoint captures HIL responses before processing. The runner simplification (removing _resumed_from_checkpoint flag and the pre-loop entry checkpoint) is clean since entry-checkpoint responsibility is now fully owned by the Workflow. The apply_checkpoint change correctly clears stale events while preserving the streaming flag. The ag-ui fix ensures proper event ordering before terminal events. All changes are consistent with the test assertions.

✓ Security Reliability

The PR cleanly refactors checkpoint creation responsibility from the Runner to the Workflow, removes the now-unnecessary _resumed_from_checkpoint flag, and adds proper entry/response-entry checkpoints for full replayability. The apply_checkpoint correctly drops stale events before restoring, and the can_handle validation at the input boundary prevents silent message drops. No injection risks, resource leaks, or unhandled failure modes were found. The create_checkpoint_if_enabled method already gracefully handles checkpoint creation failures (logs a warning and continues), so the new callsites at lines 678 and 1039 inherit that fault tolerance. The asyncio.gather for sending responses at line 1030 is preceded by thorough type validation, making partial-mutation scenarios practically unreachable.

✓ Test Coverage

The PR provides strong test coverage for the new checkpoint replayability behavior. Three new comprehensive integration tests in test_checkpoint.py cover entry checkpoints, per-superstep lineage with replay, and response-entry checkpoints. The runner-level test was appropriately replaced with a simpler one that matches the simplified implementation. Two new unit tests verify apply_checkpoint drops stale events and preserves the streaming flag. The only notable gap is the missing test for the newly added DeprecationWarning on reset_for_new_run(), which is minor since the deprecation is a transitional concern.

✓ Failure Modes

This PR cleanly restructures checkpoint responsibility between the Workflow (entry checkpoints) and Runner (per-superstep checkpoints). The removal of _resumed_from_checkpoint and its try/finally cleanup is safe because the flag's role (suppressing duplicate entry checkpoints in the runner) is entirely subsumed by the new architecture. The apply_checkpoint change correctly drops stale events while preserving the streaming flag. The create_checkpoint_if_enabled best-effort semantics (swallow failures, log warning) provide graceful degradation — a failed entry checkpoint means the first superstep checkpoint still captures state after the start executor runs. No concrete silent-failure, lost-error, or stale-state paths were found that are introduced by this diff.

✓ Design Approach

I found one design-level correctness issue in the new entry-checkpoint flow. Fresh message runs on a reused workflow instance reset only the iteration counter, not the stored checkpoint parent, so the entry checkpoint created for the next independent run can silently chain onto the previous run’s lineage instead of starting a new one.

Suggestions

• Consider adding a test that verifies reset_for_new_run() emits a DeprecationWarning (e.g., with pytest.warns(DeprecationWarning, match='reset_for_new_run'):). The deprecation was added at _runner_context.py:483 but no test validates the warning is emitted, which means a future refactor could silently remove it.

---

Automated review by TaoChenOSU's agents