Python: Make AG-UI runs disconnect durable - #8672
Eduard van Valkenburg (eavanvalkenburg) wants to merge 4 commits into
Conversation
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
Critical lifecycle and concurrency defects, plus snapshot-consistency gaps, remain unresolved.
Get a fresh assessment by requesting another Copilot review.
Review effort: Balanced
Findings: 2
Open (2)
What changed in this PR
Adds opt-in disconnect-durable AG-UI runs and incremental stateless-agent snapshots.
Changes:
- Adds bounded detached execution and scoped mutation guards.
- Persists snapshots at model, tool-result, and approval boundaries.
- Expands durability documentation and tests.
Required fixes:
- Critical —
_endpoint.py:350: An immediate disconnect before the response generator starts can prevent drainer creation, blocking the producer and strandingactive_runs. Make draining independent of generator finalization and test this case. - Critical —
_endpoint.py:87: Hydration detection overlooks resume payloads inforwarded_props, allowing concurrent mutation. Reuse the shared resume extractor or check all supported forms. - Moderate —
_agent_run.py:3388: Snapshotting onfinish_reasonoccurs before model-roundtrip finalization and can persist stale continuation state. Save only after finalization. - Moderate —
_agent_run.py:3438:mcp_server_tool_resultis omitted from incremental tool-result snapshots. Include it and add coverage without a finish reason. - Moderate —
_agent_run.py:3589: Approval snapshotting lacks the service-session exclusion, risking non-atomic continuation state. Keep service-session snapshots terminal-only. - Moderate —
_agent_run.py:3589: Approval snapshots are limited toResponseStream; plainAsyncIterablestreams can lose the interrupt snapshot. Save for all waiting approvals.
| File | Description |
|---|---|
python/packages/ag-ui/tests/ag_ui/test_run.py |
Snapshot ordering and approval tests |
python/packages/ag-ui/tests/ag_ui/test_endpoint.py |
Detached execution, conflict, and disconnect tests |
python/packages/ag-ui/README.md |
Durability and detached-run documentation |
python/packages/ag-ui/AGENTS.md |
Package implementation guidance |
python/packages/ag-ui/agent_framework_ag_ui/_snapshot_session.py |
Snapshot save lifecycle documentation |
python/packages/ag-ui/agent_framework_ag_ui/_endpoint.py |
Detached execution and concurrency protection |
python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py |
Incremental and terminal snapshot persistence |
💡 Add a code-review agent skill for context-aware, tailored reviews. Learn more in the docs.
There was a problem hiding this comment.
MAF Automated Review — Iteration 1
Result: Findings reported
Scope: full PR (1 commit(s)): a70fed415d72
Model: gpt-5.6-sol-fast
Overview
The PR adds opt-in endpoint-owned producers with bounded per-request queues and retained drainers, and advances stateless snapshots at model, tool-result, and approval boundaries while preserving terminal-only service-session and workflow snapshots. Its tests cover disconnect backpressure, scope isolation, approval ordering, error handling, and unchanged default cancellation. Two residual issues remain: detached task admission and lifetime are unbounded across client-controlled thread IDs, and model-boundary snapshots are committed before core stream finalization and context-provider side effects complete.
Reviewed the supplied pull-request change set across correctness, security/reliability, architecture, and failure behavior.
2 verified findings remained after source verification (1 high, 1 medium) across 2 files. Details are attached to the affected lines below.
Affected areas: python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py, python/packages/ag-ui/agent_framework_ag_ui/_endpoint.py
Jose Alvarez (jpalvarezl)
left a comment
There was a problem hiding this comment.
The hydration capacity bypass in ea839d799 works: snapshot reads remain available while detached capacity is occupied. I rechecked the remaining lifecycle/finalization cases on that commit and can still reproduce these three issues:
-
P1: A connected fast producer can lose its entire SSE response. In
python/packages/ag-ui/agent_framework_ag_ui/_endpoint.py:360-369, the producer can fill the 16-item queue before the response reader task is scheduled.QueueFullthen marks that still-connected reader abandoned and clears the queue. Using the existingStubAgentwith 32 immediate text updates and a normal ASGI POST with no disconnect, attached execution returns 37 events includingRUN_FINISHED; detached execution returns HTTP 200 with an empty body. This reproduces with bothkeepalive_seconds=Noneand the default 15. Can startup/backpressure be coordinated without treating queue fullness as a disconnect, with a connected fast-producer regression test? -
P2: The next update after
finish_reasonis not necessarily a finalized-turn boundary. Follow-up to the existing finalization discussion, rather than a separate duplicate topic:_agent_run.py:3390-3394saves when the stream advances to any next update. A trailing usage-only update can arrive before the inner result/finalization hook. With a realAgentandFunctionInvocationLayer, a tool-call/finish-reason update followed by usage and thenResponseInvalidatedExceptionleaves callc1in the persisted snapshot although it was rejected and the tool never executed. This also reproduces when the exception is raised specifically from the inner result hook. Could the save wait for an actual successful finalization boundary, with trailing-metadata/finalizer-failure coverage? -
P2: An exactly-full unstarted stream retains completed-run capacity. Follow-up to the existing reader-startup discussion: at
_endpoint.py:379-380, a runner that produces exactly 16 events and completes blocks onqueue.put(None)if its response body is never advanced. The startup watchdog sets abandonment but does not unblock this sentinel write. Withmax_detached_runs=1, the completed run still occupies capacity after the watchdog and unrelated mutations receive 503 until expiry (3600 seconds by default). Hydration now bypasses the limit, but mutations remain affected. Could completion signaling release capacity without depending on an unstarted consumer? Please cover exactly queue-capacity events as well as the existing over-capacity case.
All three reproductions are bounded offline cases; no live provider calls were needed.

Motivation & Context
AG-UI FastAPI runs currently execute inside the SSE response generator, so a client disconnect can close the only consumer before the run reaches snapshot, checkpoint, and approval-state finalization. Long-running stateless agent turns also persist Thread Snapshots only at terminal paths, leaving the previous completed run as the latest recoverable state throughout tool loops.
This change makes disconnect completion an explicit endpoint option and advances stateless agent snapshots at completed tool-result and approval boundaries without introducing a resumable event-log service.
Description & Review Guide
detached_runs=Trueendpoint execution with connected-reader backpressure, bounded producer admission, abandoned-run expiry, and unchanged default cancellation behavior.detached_runs=Trueis set.Related Issue
Fixes #7215
Partially addresses #7230 (disconnect-safe execution only; resumable event replay remains out of scope).
No open pull request currently addresses either 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.