microsoft / microsoft/agent-framework

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

Open
#7,230 0 comments 0 reactions 1 assignee View on GitHub

Nobody has claimed this yet.

ag-ui python
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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.