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

[antsok]

Description

1. Current behavior

_save_thread_snapshot is invoked only at run-completion sites in _agent_run.py (the confirmation-turn path and the main RUN_FINISHED path). A harness agent's Execute todo auto-loop is one AG-UI run that can span many model roundtrips and minutes of tool execution. For its whole duration, the last persisted snapshot is the previous run's end state.

When the client disconnects mid-run (mobile screen-off, page reload, network blip), the SSE generator is aborted (http.disconnect), so the completion code — including the snapshot save — never executes. Meanwhile the underlying agent work can continue running detached.

2. Observed live (Azure Container Apps deployment, harness + AG-UI)

Case A — hard data loss + discarded work (2026-07-14, ~21:35–22:02 UTC). A user refreshed the page ~6 minutes into a long Execute run:

• the stream died; no Completed streaming was ever logged for the run — the snapshot save was skipped;
• the reloaded client hydrated the previous run's snapshot: the conversation visibly rolled back to the end of planning; even the user's plan-approval message was gone;
• the detached run kept executing for many more minutes (event counter reached 3459+ after the disconnect; file_access_write/file_memory_write tool calls kept succeeding up to ~20 minutes later) — tokens billed, results appended to a session no snapshot will ever reference. Worst of both worlds: neither cancelled nor persisted.

Case B — stale-window confusion (2026-07-15, 14:27–14:32 UTC). Screen-off killed the SSE during a 5196-event run; the restored tab reloaded 18 seconds before the run finished and hydrated the previous run's snapshot (including its already-superseded pending-approval interrupt). The run then completed and saved correctly; the user saw "all progress gone" and a stale interrupt card until a later manual refresh. Recoverable — but indistinguishable, from the user's point of view, from case A.

3. Requested

Any of the following (in descending order of preference; they compose):

1. Incremental snapshot saves at safe boundaries. Persist the thread snapshot after each model roundtrip / executed tool-result batch inside the run loop (optionally throttled, e.g. at most once per N seconds). The snapshot format needs no change — it is the same messages/state/interrupt/session_state payload, just written more often. Hydration then returns the latest roundtrip instead of the previous run's end.
2. Deterministic disconnect semantics. On client disconnect, either (a) detach and run to completion including the final snapshot save (the save must not live only inside the event-stream generator), or (b) cancel the run promptly. Either is defensible; continue-then-discard is not. A config knob (on_client_disconnect: "complete_and_save" | "cancel") would cover both camps.
3. A snapshot-save hook/config on the endpoint (e.g. snapshot_save_mode: "run_end" | "roundtrip", or a callback) so hosts can opt into incremental persistence without the framework changing its default.

4. Why it matters

The harness positions AG-UI agents for long-running autonomous loops; at that duration, client-connection lifetime and run lifetime diverge routinely (mobile browsers suspend tabs aggressively). Today the durability contract is "nothing since the last completed run", which for a 5–30 minute auto-loop run means the entire turn. With #7084, private Session Continuation State shares the same fate. Server-side the work usually succeeds — the loss is purely a persistence-timing artifact.

5. Workarounds evaluated (application-side)

• Post-reload polling of the snapshot store: repairs case B only, after the run eventually completes.
• Periodic app-initiated saves are not possible: the endpoint owns the snapshot lifecycle and the authoritative in-run message list; a host writing snapshots concurrently would race the endpoint's own run-end save.
• Keep-alive tricks cannot help: the disconnect originates client-side (tab suspension), not from idle timeouts.

Code Sample

Language/SDK

Both