Skip to content

Python: [Feature]: AG-UI runs should survive client disconnect (detached execution + resumable event stream) #7230

Description

@antsok

Description

1. Summary

add_agent_framework_fastapi_endpoint executes the entire run inside the SSE response generator: event_generator() (_endpoint.py) does async for event in protocol_runner.run(input_data), and everything run_agent_stream owns — stored-snapshot hydration, the fresh AgentSession, the agent stream, approval-state save, _save_thread_snapshot, and RUN_FINISHED — happens while that one HTTP response is being streamed.

When the client goes away mid-run (a mobile browser backgrounding the tab suspends its fetch within seconds; also network blips, proxy resets, laptop lid closes), Starlette stops pulling the generator and closes it. The consequence is worse than a dropped connection:

  • the agent run is abandoned mid-flight — tool calls in progress are torn down, the analysis silently dies;
  • _save_thread_snapshot never fires — it runs after the stream drains, just before RUN_FINISHED (_agent_run.py), so the package's own thread-snapshot resume story breaks: the store keeps the previous run's snapshot, and a client that reconnects hydrates stale history with no trace of the interrupted turn;
  • a pending approval interrupt emitted in that run is lost the same way (persisted only via the same end-of-run save);
  • the client, if it comes back, has no protocol-level way to discover any of this — the AG-UI spec defines no reconnection mechanism and documents MESSAGES_SNAPSHOT/STATE_SNAPSHOT as the recovery primitives ("synchronizing after connection interruptions"), but there is nothing to synchronize to while a run is (or was) in flight.

Mobile backgrounding makes this an every-session event for any long-running agent, not an edge case. Request: make run execution disconnect-safe in the endpoint (phase A), and optionally expose a resumable event stream so clients can re-attach to an in-flight run (phase B).

flowchart LR
    subgraph today["Today: run lives inside the HTTP response"]
        REQ["POST / (SSE)"] --> GEN["event_generator()"]
        GEN --> RUN["run_agent_stream():<br/>hydrate -> session -> agent stream<br/>-> save snapshot -> RUN_FINISHED"]
        X["client disconnect"] -. "closes generator" .-> GEN
        GEN -. "GeneratorExit" .-> RUN
        RUN -. "abandoned: no snapshot,<br/>no RUN_FINISHED, interrupt lost" .-> DEAD[" "]
    end
Loading

2. Mechanics

  1. agent_endpoint returns EventSourceResponse(event_generator(), ping=keepalive_seconds, ...) (_endpoint.py). The generator is the only thing keeping protocol_runner.run(input_data) alive.
  2. On client disconnect, the ASGI server cancels the response; sse-starlette/Starlette close the generator; GeneratorExit propagates into run_agent_stream at whatever yield it is suspended on.
  3. Everything downstream of that point is skipped — including the two persistence calls that make AG-UI resume work at all:
    • await _save_thread_snapshot(...) (both call sites in _agent_run.py run only after the stream drains);
    • _save_tool_approval_state(...) next to them.
  4. keepalive_seconds (PR Python: Add AG-UI FastAPI SSE keepalive support #6980) cannot help: keepalives keep an idle connection alive through proxies, but a suspended mobile tab is the disconnect — no amount of pinging reaches a frozen renderer.

Impact matrix:

State After a mid-run disconnect today
agent run (LLM/tool loop) cancelled mid-flight, work lost
thread snapshot store stale (previous run) — reconnecting client hydrates old history
pending approval interrupt lost (never persisted)
tool-approval state store not saved for the run
client UI on return frozen "running" state forever; no RUN_FINISHED, no error

3. Reproduction sketch

Any agent slow enough to still be streaming when the client dies:

app = FastAPI()
add_agent_framework_fastapi_endpoint(
    app, agent=slow_agent, path="/",
    snapshot_store=InMemoryAGUIThreadSnapshotStore(),
    snapshot_scope_resolver=lambda r: "scope",
)
  1. curl -N -X POST http://host/ -d '<RunAgentInput with thread_id T>', kill curl after the first few events.
  2. Observe server-side: the run stops (no further tool/LLM activity), no snapshot is written for T.
  3. POST / again for T with empty messages → hydration replays the previous snapshot; the interrupted turn never existed.

On a phone the same happens by backgrounding the browser for ~30s during a run.

4. Application-level workaround (works, but everyone must reinvent it)

It is only possible because the endpoint consumes exactly one seam, agent.run(input_data):

  • a mixin composed onto our outermost AgentFrameworkAgent subclass pumps super().run(...) in a background asyncio.create_task (created in the request task, so ContextVars/scope carry over) that appends into a per-(scope, thread_id) replayable event buffer; the endpoint's generator merely reads the buffer, so GeneratorExit detaches the reader and the run — including both upstream persistence calls — always completes;
  • a run-status field and a GET /threads/{id}/live?after=<seq> SSE route replay the buffer (re-encoded with the upstream EventEncoder for byte parity, same keepalive contract) and follow the live run, prefixed by a synthetic MESSAGES_SNAPSHOT built from the run's input messages so a client that lost its transcript rehydrates fully;
  • per-thread run serialization via the registry.

Costs of doing this outside the framework: we re-implement SSE framing/keepalive, duplicate scope handling the endpoint already resolves, must track upstream's private lifecycle to stay correct, our buffer is in-process only (no store protocol to plug Redis into), and every AG-UI application hits the exact same wall.

5. Proposed out-of-the-box design

Phase A — disconnect-safe runs (small, high value). Opt-in endpoint parameter in the spirit of keepalive_seconds:

add_agent_framework_fastapi_endpoint(app, agent, path="/", detached_runs=True)

event_generator drives protocol_runner.run(input_data) in a background task feeding a queue and streams from the queue; on client disconnect the task keeps consuming to completion. Snapshot + approval-state persistence then always run, and the existing hydration path already gives reconnecting clients a correct (finished-run) view. No new protocol surface; fixes the data-loss half outright.

Phase B — resumable event stream (the full reconnect story). On top of A:

  • buffer the run's encoded events behind a small store protocol (mirroring AGUIThreadSnapshotStore: in-memory default, host-pluggable Redis/etc. for multi-replica), keyed by resolved scope + thread_id, bounded + TTL'd;
  • surface run status (active, run_id, seq) — naturally alongside the stored snapshot record;
  • mount a resume route with the endpoint (e.g. GET {path}runs/{thread_id}?after=<seq>, or honor SSE Last-Event-ID on it): replay buffered events from after, then follow live until the run ends. When replaying from 0, prefix a MESSAGES_SNAPSHOT built from the run's input messages — the stored snapshot predates the in-flight run, so this is what restores the client's transcript;
  • serialize runs per thread while one is detached-active.

A buffered replay also recovers something snapshot-based resume never can: reasoning
traces
. REASONING_* events exist only in the live stream — thread snapshots have no
representation for them — so today any resumed thread shows reasoning headers with no
content. The event buffer carries the run's reasoning verbatim, making phase B the only
path by which thinking survives a reconnect.

This would make MAF the reference implementation of the recovery flow the AG-UI spec implies but leaves to implementers (snapshots as the documented "synchronizing after connection interruptions" primitive, transport left open).

6. Closing thoughts

  • The endpoint already owns every ingredient: the event encoder, keepalive framing, the snapshot store and scope resolver, the approval state store, and the run lifecycle itself. App-side implementations must shadow all of them.
  • Phase A is arguably a robustness fix for the package's own persistence contract (_save_thread_snapshot silently skipped is surprising regardless of reconnection plans).
  • A store protocol for the buffer solves multi-replica cleanly at the framework level; app-side buffers are process-local.

Code Sample

Language/SDK

Python

Metadata

Metadata

Assignees

Labels

ag-uiUsage: [Issues, PRs], Target: AG-UI protocol integrationpythonUsage: [Issues, PRs], Target: Python

Projects

Status
No status

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions