Python: add checkpointing support to AgentFrameworkWorkflow.run() in agent-framework-ag-ui - #6646
Open
vaibhav-patel wants to merge 11 commits into
Open
Conversation
…ag-ui The ag-ui AgentFrameworkWorkflow.run() previously accepted only a RunAgentInput payload and exposed no way to use the core workflow's checkpointing/state-persistence, unlike the core agent-framework workflow implementations. This left ag-ui workflows without resumable execution. Add optional checkpoint_storage and checkpoint_id keyword arguments to run(), threaded through run_workflow_stream() into the core Workflow.run(). This delegates to the existing core capability instead of reinventing it and keeps the public surface consistent with Workflow.run(): - checkpoint_storage enables checkpoint creation at each superstep boundary. - checkpoint_id resumes a run from a persisted checkpoint; incoming messages are forwarded only as request-info responses (never as a new start-executor message) to honor the core's message/checkpoint_id mutual exclusivity, and responses + checkpoint_id performs a restore-then-send in one call. Both can also be supplied via the input_data keys __ag_ui_checkpoint_storage and __ag_ui_checkpoint_id so the FastAPI endpoint (which calls run(input_data) positionally) can opt in without changing its call site; explicit keyword arguments take precedence. Checkpoint resume bypasses the AG-UI thread snapshot hydration early-returns so it always reaches the core restore path. Backward compatible: run(input_data) keeps working unchanged, and the non-checkpoint path still calls run_workflow_stream(input_data, workflow) with its original two-argument convention. Adds focused tests covering checkpoint creation, resume-from-checkpoint, input-data-keyed params, and the unchanged default path. Fixes microsoft#6632.
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds workflow checkpointing support to the agent-framework-ag-ui adapter so AgentFrameworkWorkflow.run() can create checkpoints and resume from a persisted checkpoint_id, mirroring the core agent_framework.Workflow.run() API.
Changes:
- Extended
AgentFrameworkWorkflow.run()to acceptcheckpoint_storage/checkpoint_id(and optionally read them frominput_data), ensuring checkpoint resumes bypass snapshot-only hydration paths. - Updated
run_workflow_stream()to forward checkpointing parameters to the core workflow and to avoid pre-run early-returns that would otherwise skip restores. - Added focused tests covering checkpoint creation, resume, input-data-keyed params, and the unchanged default path.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
python/packages/ag-ui/agent_framework_ag_ui/_workflow.py |
Adds checkpoint kwargs to AgentFrameworkWorkflow.run() and plumbs them through to the stream runner, while preserving existing snapshot behavior. |
python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py |
Extends run_workflow_stream() with checkpoint parameters and ensures checkpoint resumes reach workflow.run(checkpoint_id=...). |
python/packages/ag-ui/tests/ag_ui/test_workflow_agent.py |
Adds checkpointing tests for the AG-UI workflow wrapper. |
moonbox3
reviewed
Jul 9, 2026
Contributor
Python Test Coverage Report •
Python Unit Test Overview
|
||||||||||||||||||||||||||||||||||||||||
Contributor
|
@vaibhav-patel please also look at the failing code quality CI/CD checks thanks |
…responses; fix CI lint/typing A checkpoint-only resume no longer clobbers the stored AG-UI thread snapshot: the snapshot builder is seeded with the prior stored history so the saved snapshot keeps the earlier replayable transcript plus the newly produced output. Resume responses are now coerced against the post-restore pending requests on a checkpoint restore, so a JSON function_approval_response resumes through AG-UI after a cold restore instead of failing with a response-type mismatch. Also update the test-double workflow run() overrides to match the new keyword-only parent signature and re-sort the workflow test imports so ruff and the typing checkers pass.
Contributor
Author
|
@moonbox3 CI is green now — fixed the Thanks for the thorough review — the two resume bugs you reproduced are fixed too (details on the inline threads). |
moonbox3
reviewed
Jul 13, 2026
Reading pending request_info events for resume-response coercion previously restored the checkpoint into the live workflow, which invoked every executor's on_checkpoint_restore hook. workflow.run(checkpoint_id=...) then restored again, running those hooks a second time. Custom restore hooks are not required to be idempotent, so this could duplicate restoration work or break workflows that expect exactly one restore per resume. Load the persisted WorkflowCheckpoint directly from storage (runtime override or the workflow's build-time context storage) and read its pending_request_info_events instead. This exposes the same post-restore pending set for the resume contract and response coercion without mutating workflow state or running any restore hook, leaving workflow.run(checkpoint_id=...) as the single restore per resume. Add a regression test asserting on_checkpoint_restore runs exactly once on a checkpointed ag-ui resume.
TaoChenOSU
reviewed
Jul 13, 2026
Contributor
|
Any follow up here, @vaibhav-patel? Thank you. |
…checkpointing-takeover # Conflicts: # python/packages/ag-ui/agent_framework_ag_ui/_workflow.py
… surfaces Checkpoint storage is now configured on AgentFrameworkWorkflow (or the FastAPI endpoint) instead of being smuggled through input_data keys, and a run resumes by supplying its checkpoint id in the AG-UI forwarded props. With storage always in hand, resume-response coercion reads the pending request set straight from the persisted checkpoint via the public CheckpointStorage.load(), replacing the private runner-context fallback, and the core run call forwards checkpoint arguments directly, relying on core validation for conflicting parameters. Requesting a resume without configured storage now fails with a clear error.
Contributor
|
Helping get this across the finish line. |
moonbox3
temporarily deployed
to
github-app-auth
August 3, 2026 06:54 — with 
GitHub Actions
Inactive
GitHub Actions
Inactive
moonbox3
approved these changes
Aug 3, 2026
moonbox3
enabled auto-merge
August 3, 2026 06:57
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Fixes #6632. Adds workflow checkpointing to the ag-ui adapter's
AgentFrameworkWorkflow.run(), bringing it to parity with the coreagent_framework.Workflow.run()checkpointing API.run()now accepts two optional keyword arguments:checkpoint_storage: CheckpointStorage | None— enables checkpoint creation at eachsuperstep (forwarded to the core workflow).
checkpoint_id: str | None— resumes a run from a persisted checkpoint.This delegates to the existing core capability rather than reinventing it, so the public
surface mirrors
Workflow.run(checkpoint_storage=..., checkpoint_id=...).Details
messageis never passed alongsidecheckpoint_id; incoming messages are forwarded only as request-inforesponses(so
responses+checkpoint_idis a restore-then-send).input_datakeys__ag_ui_checkpoint_storage/__ag_ui_checkpoint_id, letting the FastAPI endpoint (which callsrun(input_data)positionally) opt in without a signature change; explicit keyword args take precedence.
bypasses snapshot-hydration early-returns to reach the core restore path.
Backward compatibility
The existing
run(RunAgentInput)/run(input_data)signature is preserved — the newparams are optional and default to
None, and the non-checkpoint path still callsrun_workflow_stream(input_data, workflow)with its original two-argument convention.Tests
New focused tests cover checkpoint creation, resume-from-checkpoint, input-data-keyed
params, and the unchanged default path. Full ag-ui suite green (824 passed); ruff
lint/format clean; no new pyright errors.