microsoft / microsoft/agent-framework
Python: [Feature]: AG-UI runs should survive client disconnect (detached execution + resumable event stream)
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 13.6k
- Forks
- 2.3k
- Avg merge
- 2d 45m
- Merged PRs (30d)
- 358
Description
### 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).
```mermaid
flowchart LR
subgraph today["Today: run lives inside the HTTP response"]
REQ["POST / (SSE)"] --> GEN["event_generator()"]
GEN --> RUN["run_agent_stream():
hydrate -> session -> agent stream
-> save snapshot -> RUN_FINISHED"]
X["client disconnect"] -. "closes generator" .-> GEN
GEN -. "GeneratorExit" .-> RUN
RUN -. "abandoned: no snapshot,
no RUN_FINISHED, interrupt lost" .-> DEAD[" "]
end
```
## 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 #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:
```python
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 ''`, 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=` 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`:
```python
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=`, 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
```markdown
```
### Language/SDK
Python
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Assessment
This issue has not been assessed yet.