Repository navigation
Python: preserve AG-UI input on failed runs and prevent duplicate snapshot messages - #9798
Ran Shemtov (ranst91) wants to merge 4 commits into
Conversation
There was a problem hiding this comment.
🟡 Changes recommended
Checkpoint persistence can lose recoverable input after session-ID rebinding and prematurely clear unresolved interrupt prompts, leaving hydration incomplete after a failure or disconnect.
2 open issues
What changed in this PR
This PR strengthens persistence in the Python AG-UI adapter so failed runs retain submitted input and differing client history does not duplicate saved messages.
Changes:
- Filter already-stored message IDs from reconciled incoming history.
- Checkpoint input before provider execution for configured stateless snapshot sessions.
- Preserve attachment filenames during multimodal snapshot normalization and add regression coverage.
| File | Description |
|---|---|
python/packages/ag-ui/tests/ag_ui/test_run.py |
Expect submitted input to survive response invalidation. |
python/packages/ag-ui/tests/ag_ui/test_run_common.py |
Test deduplication when client history differs from saved history. |
python/packages/ag-ui/tests/ag_ui/test_message_adapters.py |
Test filename preservation across normalization. |
python/packages/ag-ui/tests/ag_ui/test_endpoint.py |
Cover failed-run persistence, cancellation, media, and approval-history replay. |
python/packages/ag-ui/agent_framework_ag_ui/_run_common.py |
Exclude stored IDs from incoming history suffixes. |
python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py |
Retain filenames in normalized binary parts. |
python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py |
Add an input checkpoint before provider execution. |
Review effort: Balanced
| await save_thread_snapshot( | ||
| persisted_messages=persisted_input, | ||
| state=latest_state_snapshot, | ||
| interrupt=snapshot_interrupts(), |
| await save_thread_snapshot( | ||
| persisted_messages=persisted_input, | ||
| state=latest_state_snapshot, | ||
| interrupt=snapshot_interrupts(), | ||
| ) |
|
Ran Shemtov (@ranst91), thanks for the detailed investigation and validation. This is a substantial community contribution, so the statement that the AG-UI team initiated it is not enough by itself to establish the required agreed direction. Please link the public repository issue that maintainers agreed to implement, or have a maintainer explicitly confirm here that the internal request is sufficient for this PR. Review and approval remain blocked until that direction and the current CI failures are resolved. |

Motivation & Context
This change addresses three persistence gaps in the Python AG-UI adapter:
Our original duplication case was triggered by a CopilotKit wrapper that removed historical tool messages. We fixed that wrapper separately. The reconciliation change here provides additional protection when incoming history differs from saved history.
Description & Review Guide
Existing duplicated snapshots are not rewritten. Unfinished assistant output is not committed by the input checkpoint.
Please review the checkpoint boundary, approval continuation behavior, and preservation of message identity and attachment metadata.
Validation
At commit
bf82841e7:These checks used a custom JSON snapshot store. They do not establish coverage for every storage implementation or browser/Intelligence workflow.
Related Issue
This contribution was initiated by the AG-UI team. An agreed Microsoft issue has not yet been linked.
AI Assistance
Implementation, tests, investigation, and this description were AI-assisted. Human contributor review and responsibility acknowledgement remain pending.