Skip to content

Python: Make AG-UI runs disconnect durable - #8672

Open
Eduard van Valkenburg (eavanvalkenburg) wants to merge 4 commits into
microsoft:mainfrom
eavanvalkenburg:ag-ui-disconnect-durability
Open

Eduard van Valkenburg (eavanvalkenburg) wants to merge 4 commits into
microsoft:mainfrom
eavanvalkenburg:ag-ui-disconnect-durability

Conversation

@eavanvalkenburg

@eavanvalkenburg Eduard van Valkenburg (eavanvalkenburg) commented Sep 23, 2026 •

Copy link
Copy Markdown
Member

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

  • What are the major changes?
    • Add opt-in detached_runs=True endpoint execution with connected-reader backpressure, bounded producer admission, abandoned-run expiry, and unchanged default cancellation behavior.
    • Allow empty snapshot hydration during an active detached run while rejecting concurrent mutations for the same scoped thread with HTTP 409.
    • Persist configured stateless agent Thread Snapshots after function/MCP tool-result batches and approval safe points; keep model-only, service-session, and workflow snapshots terminal-only so unfinalized output and continuation state remain atomic.
    • Add deterministic disconnect, startup, exact-capacity, fast-producer, scope-isolation, invalidation, workflow, approval, and snapshot-ordering coverage, plus endpoint documentation.
  • What is the impact of these changes?
    • Applications can opt into finite AG-UI runs that complete persistence after the HTTP reader disconnects.
    • Existing endpoints retain their current disconnect behavior unless detached_runs=True is set.
    • Reconnecting clients recover the latest committed Thread Snapshot; this does not add in-flight SSE event replay.
  • What do you want reviewers to focus on?
    • Producer/reader startup, cancellation, bounded-queue backpressure, and expiry cleanup.
    • Scoped-thread mutation exclusion versus hydration reads.
    • Snapshot ordering at completed tool/approval boundaries and provider-owned continuation state.

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

  • 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.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 High severity

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 stranding active_runs. Make draining independent of generator finalization and test this case.
  • Critical — _endpoint.py:87: Hydration detection overlooks resume payloads in forwarded_props, allowing concurrent mutation. Reuse the shared resume extractor or check all supported forms.
  • Moderate — _agent_run.py:3388: Snapshotting on finish_reason occurs before model-roundtrip finalization and can persist stale continuation state. Save only after finalization.
  • Moderate — _agent_run.py:3438: mcp_server_tool_result is 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 to ResponseStream; plain AsyncIterable streams 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.

Comment thread python/packages/ag-ui/agent_framework_ag_ui/_endpoint.py
Comment thread python/packages/ag-ui/agent_framework_ag_ui/_endpoint.py

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Comment thread python/packages/ag-ui/agent_framework_ag_ui/_endpoint.py
Comment thread python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py Outdated

@jpalvarezl Jose Alvarez (jpalvarezl) left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:

  1. 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. QueueFull then marks that still-connected reader abandoned and clears the queue. Using the existing StubAgent with 32 immediate text updates and a normal ASGI POST with no disconnect, attached execution returns 37 events including RUN_FINISHED; detached execution returns HTTP 200 with an empty body. This reproduces with both keepalive_seconds=None and the default 15. Can startup/backpressure be coordinated without treating queue fullness as a disconnect, with a connected fast-producer regression test?

  2. P2: The next update after finish_reason is not necessarily a finalized-turn boundary. Follow-up to the existing finalization discussion, rather than a separate duplicate topic: _agent_run.py:3390-3394 saves when the stream advances to any next update. A trailing usage-only update can arrive before the inner result/finalization hook. With a real Agent and FunctionInvocationLayer, a tool-call/finish-reason update followed by usage and then ResponseInvalidatedException leaves call c1 in 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?

  3. 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 on queue.put(None) if its response body is never advanced. The startup watchdog sets abandonment but does not unblock this sentinel write. With max_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.

This branch was successfully deployed

1 active deployment
github-app-auth — 901d4a35 Deployed Sep 25, 2026 by eavanvalkenburg via add_label #23778
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

documentation Usage: [Issues, PRs], Target: documentation in the code base and learn docs python Usage: [Issues, PRs], Target: Python

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Python: [Feature]: persist AG-UI thread snapshots incrementally (and handle client disconnect), not only at run completion

3 participants