AG-UI Streaming State Management#

AG-UI streaming state management in @assistant-ui/react-ag-ui covers how the runtime processes the event stream from an AG-UI agent and keeps message state consistent from the start of a run through its settlement. The two central components are:

• AgUiThreadRuntimeCore — owns the thread's message array, orchestrates runs, and routes special events (cross-run tool results, snapshots, state deltas) before they reach the aggregator.
• RunAggregator — accumulates individual AG-UI events into full ChatModelRunResult snapshots that are applied to the in-progress assistant message on every event.

The entry point for any new run is startRun(), which creates a fresh RunAggregator, wires up the event dispatcher, creates an AbortController, and calls agent.runAgent(). Each event emitted by the agent is dispatched through handleEvent().

Run Lifecycle and Event Flow#

startRun()
 ├─ insert optimistic assistant placeholder (isOptimistic: true, status: "running")
 ├─ create RunAggregator { emit: applyUpdate, onServerMessageId: reassignAssistantId }
 ├─ agent.runAgent(input, subscriber) — streams events
 │ └─ every event → handleEvent(aggregator, event)
 │ ├─ STATE_SNAPSHOT / STATE_DELTA → update stateSnapshot directly
 │ ├─ MESSAGES_SNAPSHOT → importMessagesSnapshot()
 │ ├─ TOOL_CALL_RESULT (cross-run) → applyCrossRunToolResult()
 │ ├─ ACTIVITY_SNAPSHOT (cross-run, mcp-apps) → applyCrossRunActivitySnapshot()
 │ └─ everything else → aggregator.handle(event)
 │ └─ aggregator.emit() → applyUpdate() → updateAssistantMessage()
 └─ finishRun() → clears AbortController, sets isRunningFlag = false

handleEvent() is the traffic director 1. Four event types are intercepted at the core level and never reach the aggregator: STATE_SNAPSHOT, STATE_DELTA, MESSAGES_SNAPSHOT, and cross-run instances of TOOL_CALL_RESULT and ACTIVITY_SNAPSHOT. Every other event passes through to the aggregator's handle() method.

The RunAggregator handles run-lifecycle, text, reasoning, and tool-call events 2. Its emit() method reconstructs the full content array on every event using a partOrder array that preserves insertion order across text, reasoning, and tool-call parts 3. Each emitted snapshot is applied to the assistant message via updateAssistantMessage().

The onServerMessageId callback fires when TEXT_MESSAGE_START delivers a canonical server ID; reassignAssistantId() swaps the optimistic placeholder for the real ID, stripping isOptimistic from the message metadata so the core repository doesn't retain a phantom branch.

MESSAGES_SNAPSHOT: Mid-Run History Reset#

MESSAGES_SNAPSHOT is the mechanism by which the backend can authoritatively replace the thread's message array mid-run. When received, importMessagesSnapshot() calls fromAgUiMessages() to convert the raw AG-UI messages (folding role: "tool" messages into the preceding assistant message's tool-call parts, with normalization added in v0.0.19 4), then passes the result to applyExternalMessages(), which performs a hard replacement of the repository's message array.

This is a full reset: the existing in-memory message state is discarded and rebuilt from the snapshot. This makes MESSAGES_SNAPSHOT appropriate for backends (like LangGraph) that send a final authoritative updates event alongside incremental stream events, but it also means any frontend-only state (e.g. tool results added via addToolResult before the snapshot arrives) will be lost unless the backend includes them in the snapshot.

fromAgUiMessages() is also called on the page-load path (via ThreadHistoryAdapter.load() → applyExternalMessageRepository()) — the same conversion code handles both live snapshots and persisted history 5.

Tool Call Lifecycle and Cross-Run Routing#

During a run, tool calls follow: TOOL_CALL_START → TOOL_CALL_ARGS (incremental deltas) → TOOL_CALL_END → TOOL_CALL_RESULT. The RunAggregator tracks each tool call by ID and emits a new snapshot on every event, building up argsText incrementally and recording the result when it arrives 6.

preserveToolResults() 7: Because the aggregator rebuilds content entirely from stream events, any results injected via addToolResult() (front-end or human-in-the-loop tool execution) would be clobbered by the next snapshot. preserveToolResults re-merges previously resolved tool-call parts back into each incoming snapshot before it is applied, preventing the aggregator from overwriting front-end results.

Cross-run routing: A TOOL_CALL_RESULT event may arrive for a tool call that belongs to a previous run (e.g. a backend that asynchronously delivers results). handleEvent checks whether the aggregator owns the toolCallId; if not, it routes the result directly to the owning historical message via applyCrossRunToolResult(), which applies the result and calls maybeCompleteAfterToolResults() — but deliberately does not call maybeResumeAfterToolResults(), because the delivering run is already in flight and a resume would reset the head past it.

MCP app ACTIVITY_SNAPSHOT: For live runs, ACTIVITY_SNAPSHOT events go to the aggregator 8, which records the MCP app's resourceUri, serverId, and CallToolResult on the tool-call entry, and prevents a subsequent TOOL_CALL_RESULT from overwriting the richer snapshot result. For restored/cross-run tool calls (those not in the current aggregator), ACTIVITY_SNAPSHOT is handled by applyCrossRunActivitySnapshot(), added in PR #5110. fromAgUiMessages() also reconstructs mcp.app from _meta["ui/resourceUri"] on the conversion side, added in PR #5225.

State Consistency Mechanisms#

Deferred tool-call continuation 9: If a TOOL_CALL_RESULT arrives while a run is still draining (after RUN_FINISHED but before the stream closes), the follow-up resume run is stored in pendingResumeMessageId and dispatched only after startRun returns 10. This guarantees no two runs ever overlap.

Interrupt guards 11: assertNoPendingInterrupts() throws before append, reload, or resume if the last assistant message is in requires-action / interrupt status. The thread cannot be advanced past an unresolved interrupt.

State snapshot 12: STATE_SNAPSHOT events replace stateSnapshot wholesale; STATE_DELTA events apply a JSON Patch (via fast-json-patch) incrementally, with validation and immutable application (mutateDocument: false). The current state is forwarded to the agent on every buildRunInput() call.

autoCancelPendingToolCalls 13: When enabled (default), starting a new append or reload while there are unresolved tool calls will auto-cancel them with { error: "Tool call cancelled by user" } before starting the run, keeping the message in a consistent terminal state.

History deduplication: recordedHistoryIds (a Set<string>) prevents double-appending to the ThreadHistoryAdapter when the same message is processed multiple times (e.g. after an ID reassignment) 5.

Known Issues: argsText Divergence#

When a backend sends a final snapshot whose argsText is not a byte-prefix of the already-streamed partial_json — e.g. LangGraph's updates event re-stringifying Anthropic's formatted JSON into compact JSON — ToolInvocationTracker takes the A.2 path: it freezes the args stream at the incomplete prefix rather than appending the divergent text (see EDGE_CASES.md). This is the root cause described in issue #5098 and the broader #5130 family 14.

Fixed (PR #5107): When a backend result and a divergent-args snapshot arrive together, setResponse() now enqueues the result chunk into the outer stream before closing argsText. ToolExecutionStream skips the stale-prefix parse when it sees the result is already present (toolCallIdsWithBackendResult set). Chunk delivery order is deterministic: a direct controller.enqueue beats a chunk traveling through the args sub-stream pipeline 14.

Open (#5130): When no backend result exists at close time (front-end tools, HITL/interrupt flows), the stale prefix is parsed by ToolExecutionStream, throws Function parameter parsing failed, and the parse-error ToolResponse is auto-submitted. On LangGraph interrupt flows this can abandon a pending interrupt and brick the thread with Anthropic 400 errors. Two directions remain open:

1. Divergent-complete snapshot adoption — let the tracker hand the complete snapshot argsText to the execution path; requires a replace primitive on TextStreamController (currently append-only).
2. Suppress internal parse-error auto-submit — mark or suppress ToolResponse objects generated from stale-prefix parse failures so hosts can distinguish them from genuine tool failures.