Skip to content

[.NET] Predictive state updates: stream a tool call's arguments to the UI as the model writes them #2245

Description

@danroth27

Note (edited 2026-08-02): Rewritten to lead with the scenario rather than a specific API shape, and to correct inaccuracies in the original filing.

Area: .NET SDK — sdks/dotnet (primarily AGUI.Server)

Scenario

A chat panel sits beside a document editor. The user asks for a short story; the model calls write_document.

  • Expected: the story appears in the editor a few words at a time as the model writes it, then an approve/reject prompt appears.
  • Actual (.NET): the editor stays empty for the whole generation, then the document appears at once.

AG-UI calls this predictive state updates — optimistically rendering a tool call's arguments while they're being generated, before the tool runs. It generalizes to any surface outside the chat transcript: forms filling in field by field, code appearing in an editor, a spreadsheet or itinerary being drafted, or simply rendering a pending tool call's arguments instead of a spinner. LangGraph, CrewAI, ADK, Strands and the reference Python server all support it, and it's one of the dojo's headline scenarios.

Current behavior

Microsoft.Extensions.AI coalesces a provider's argument deltas into a single completed FunctionCallContent before AGUI.Server sees it, so one TOOL_CALL_ARGS event carries the whole value. Measured with a local probe (shipped packages, no credentials — available on request): against a provider streaming a document over ~300 ms, the argument surfaced once, complete, after generation finished.

AGUIStreamOptions.MapCall(toolName, ...) is therefore also called once with the finished arguments. Anything built on it re-chunks a value it already holds in full — which is what the dojo's predictive_state_updates sample does, and why its output arrives all at once.

The raw-event escape hatch doesn't close this. AGUI.Server will emit any BaseEvent supplied via RawRepresentation, so emitting paced TOOL_CALL_ARGS is not the obstacle — the deltas are already gone by the time app code runs. The only workaround is provider-specific: reading FunctionArgumentsUpdate off RawRepresentation in a DelegatingChatClient.

The client side already works. A custom IAGUITransport observes the raw TOOL_CALL_ARGS stream, and predictive state is achievable that way today. That half is a papercut, not a gap (see below).

Proposal

A declarative call-side mapping, so app code does no chunking, no partial-JSON handling, and no hand-rolled state events:

var streamOptions = new AGUIStreamOptions();

// Existing, result-side: "when write_document returns, that's the new state."
streamOptions.MapResultAsStateSnapshot("write_document");

// Proposed, call-side: "while the model is writing write_document's `document` argument,
// that argument *is* the `document` state — show it as it arrives."
streamOptions.PredictState(stateKey: "document", toolName: "write_document", argument: "document");

app.MapAGUIServer("/predictive_state_updates", agent)
   .WithMetadata(streamOptions);

Two desired properties:

  1. Declarative — no manual chunking or partial-JSON handling in app code.
  2. The tool still runs normally — prediction is observational, so it shouldn't require intercepting or suppressing the tool call, or disabling function invocation, and shouldn't care whether the tool is frontend or backend.

Open question on prior art. LangGraph, CrewAI, ADK and Strands declare this mapping by emitting a CUSTOM event named PredictState ({state_key, tool, tool_argument}), consumed by the JS client's legacy CopilotKit compatibility layer — the path the dojo runs on today. The non-legacy AgentSubscriber.onToolCallArgsEvent (which already surfaces partialToolCallArgs) looks like the more natural fit for a new SDK, and a server-side declaration needn't travel over the wire at all.

Client-side papercuts

Both optional and separable from the main ask:

  • AGUIHttpTransport is internal, so it can't be wrapped. Observing the raw event stream requires re-implementing the HTTP transport (POST + ReadAGUIEventStreamAsync) purely to add a passthrough — ~20 lines of copied plumbing.
  • No partial-JSON helper. Accumulating deltas is easy; turning a truncated {"document":"Once upon a ti into a usable value is not. The JS client uses untruncate-json and exposes the result as partialToolCallArgs; a .NET equivalent would save every consumer hand-rolling quote/brace balancing.

Upstream dependency

The server-side fix needs argument deltas to reach AGUI.Server. MEAI's argument coalescing behavior prevents this. Coalescing is the right default for invoking a function, so this likely wants an opt-in rather than a behavior change. Consider tracking an optional non-coalescing behavior in dotnet/extensions.

Related: microsoft/agent-framework#4177 covers the StateBag half of state emission and also touches streaming tool-argument deltas.

@javiercn

Metadata

Metadata

Assignees

Labels

No labels
No labels

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions