Python: Give the AG-UI Thread Snapshot lifecycle a single owner module - #7479
Open
moonbox3 wants to merge 2 commits into
Open
Python: Give the AG-UI Thread Snapshot lifecycle a single owner module#7479moonbox3 wants to merge 2 commits into
moonbox3 wants to merge 2 commits into
Conversation
Both the agent and workflow runners independently implemented the thread snapshot lifecycle: hydration replay, the load-once stored read, resume message seeding, the stored/request/deferred-default state overlay, and the save whose storage failures must never surface on an already-streamed run. The two copies had already drifted in small ways (one hydrate helper re-checked a store the caller had verified; the two cancelled-resume-id helpers differed on missing-id handling). Introduce ThreadSnapshotSession in _snapshot_session.py as the one owner of that lifecycle, opened once per run and inert when no store or scope is configured so callers stop branching on configuration. Rewire both runners onto it, consolidate _cancelled_resume_interrupt_ids in _run_common (defensive variant) and _event_messages_to_snapshot_dicts in the new module, and delete the superseded per-runner copies. The session interface is covered by dedicated tests; existing suites pin runner behavior. Public exports are unchanged.
moonbox3
temporarily deployed
to
github-app-auth
August 3, 2026 01:41 — with 
GitHub Actions
Inactive
GitHub Actions
Inactive
moonbox3
temporarily deployed
to
github-app-auth
August 3, 2026 01:41 — with
GitHub Actions
Inactive
GitHub Actions
Inactive
moonbox3
temporarily deployed
to
github-app-auth
August 3, 2026 01:41 — with
GitHub Actions
Inactive
GitHub Actions
Inactive
agent-framework-automation
Bot
added
the
python
moonbox3
temporarily deployed
to
github-app-auth
August 3, 2026 01:42 — with
GitHub Actions
Inactive
GitHub Actions
Inactive
Contributor
There was a problem hiding this comment.
Pull request overview
This PR centralizes the AG-UI Thread Snapshot lifecycle for Python runs into a single owner module (ThreadSnapshotSession), removing duplicated snapshot hydration/state overlay/resume seeding/save logic from both the agent runner and workflow runner. This aims to keep snapshot behavior consistent across runners while preserving the existing runtime semantics (including “swallow-and-log” snapshot save failures after streaming has completed).
Changes:
- Introduces
ThreadSnapshotSessionin a new private module to own snapshot open/load, hydration replay, effective-state overlay, resume message seeding, saving, and interrupt clearing. - Rewires both the agent runner and workflow runner to use the session, deleting the previously duplicated per-runner helpers and consolidating cancelled-resume interrupt-id extraction.
- Adds a new test suite covering the session interface (including save-failure swallow-and-log semantics and interrupt clearing behavior).
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated no comments.
Show a summary per file
| File | Description |
|---|---|
| python/packages/ag-ui/tests/ag_ui/test_snapshot_session.py | Adds interface-level tests for ThreadSnapshotSession covering open/hydrate/effective_state/resume seeding/save/clear semantics. |
| python/packages/ag-ui/agent_framework_ag_ui/_workflow.py | Switches workflow snapshot lifecycle operations to ThreadSnapshotSession and removes duplicated helpers. |
| python/packages/ag-ui/agent_framework_ag_ui/_snapshot_session.py | New per-run session module that becomes the single owner of snapshot lifecycle interactions. |
| python/packages/ag-ui/agent_framework_ag_ui/_run_common.py | Moves and consolidates _cancelled_resume_interrupt_ids for shared use across runners. |
| python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py | Switches agent runner snapshot lifecycle operations to ThreadSnapshotSession and removes duplicated helpers. |
Contributor
Python Test Coverage Report •
Python Unit Test Overview
|
|||||||||||||||||||||||||||||||||||||||||||||
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
The hydration test accessed run_id, snapshot, and messages on values typed as BaseEvent, which fails the tests/samples type checkers. Narrow each event with isinstance assertions before reading its fields.
moonbox3
temporarily deployed
to
github-app-auth
August 3, 2026 02:24 — with
GitHub Actions
Inactive
GitHub Actions
Inactive
moonbox3
enabled auto-merge
August 3, 2026 02:53
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.
Motivation & Context
The AG-UI Thread Snapshot is the package's central continuity record, but its lifecycle had no single owner: the agent runner and the workflow runner each carried their own implementation of hydration replay, the load-once stored read, resume message seeding, the stored/request/deferred-default state overlay, and the save whose storage failure must never surface on an already-streamed run. The copies had already drifted in small ways: one hydrate helper re-checked a store its caller had verified, and the two cancelled-resume-id helpers differed on missing-id handling (the workflow copy would raise
KeyErroron a cancelled interrupt entry with noid; the agent copy skips it).Description & Review Guide
_snapshot_session.pywithThreadSnapshotSession, opened once per run. When no store or scope is configured the session is inert (reads return nothing, writes are no-ops), so runners no longer branch on configuration at every call site. Interface:open,enabled,stored,hydrate_events,effective_state,resume_seeded_messages,save,clear_interrupts._hydrate_thread_snapshot,_save_thread_snapshot(agent),_hydrate_workflow_thread_snapshot(workflow), plus the duplicated_event_messages_to_snapshot_dicts(now in_snapshot_session.py) and_cancelled_resume_interrupt_ids(now in_run_common.py, consolidated on the defensive variant).InMemoryAGUIThreadSnapshotStoreplus a failing-store stub for the swallow-and-log path. No mocks of internals._WorkflowSnapshotBuilderintentionally stays in_workflow.py; moving it is out of scope here.Impact: none intended at runtime. 977 tests pass (957 baseline plus 20 new); ruff, formatting, and pyright are clean; the 3 remaining mypy errors are identical on main. Public exports are unchanged. One observable difference: the save-failure log line is now emitted from
_snapshot_sessioninstead of the runner modules; the message text is unchanged.Review focus: the equivalence of the rewired runner paths to the deleted inline logic, especially the agent runner's approval-gated resume seeding and the workflow save guard conditions.
Related Issue
Internal refactor, no linked 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.