@assistant-ui/react-ag-ui 0.1.2 → 0.1.4

package/src/runtime/AgUiThreadController.test.ts

@@ -1,5 +1,5 @@
 import { describe, it, expect, vi } from "vitest";
-import type { AbstractAgent } from "@ag-ui/client";
+import { HttpAgent, type AbstractAgent } from "@ag-ui/client";
 import type {
   AppendMessage,
   ThreadAssistantMessage,
@@ -540,6 +540,25 @@ describe("AgUiThreadController activity deltas", () => {
 });
 
 describe("AgUiThreadController run state", () => {
+  const createStateAgent = (sentStates: unknown[]) =>
+    new HttpAgent({
+      url: "http://agent.test/agent",
+      initialState: { fromInitialState: true },
+      fetch: async (_url, init) => {
+        const body = JSON.parse(String(init?.body));
+        sentStates.push(body.state);
+        const events = [
+          { type: "RUN_STARTED", threadId: body.threadId, runId: body.runId },
+          { type: "STATE_SNAPSHOT", snapshot: { a: 1 } },
+          { type: "RUN_FINISHED", threadId: body.threadId, runId: body.runId },
+        ];
+        return new Response(
+          events.map((event) => `data: ${JSON.stringify(event)}\n\n`).join(""),
+          { headers: { "content-type": "text/event-stream" } },
+        );
+      },
+    });
+
   type Subscriber = {
     onStateSnapshotEvent?: (payload: {
       event: { type: "STATE_SNAPSHOT"; snapshot: { count: number } };
@@ -577,6 +596,119 @@ describe("AgUiThreadController run state", () => {
     expect(runAgent.mock.calls[2]?.[0].state).toBeNull();
     expect(agent.state).toBeNull();
   });
+
+  it("isolates mounted threads sharing one agent through runs and setState", async () => {
+    const sentStates: unknown[] = [];
+    const agent = createStateAgent(sentStates);
+    const { core: threadA } = createCore(undefined, agent);
+    const { core: threadB } = createCore(undefined, agent);
+
+    await threadA.append({ ...userMessage("A"), startRun: true });
+    expect(threadA.getState()).toEqual({ a: 1 });
+    expect.soft(threadB.getState()).toEqual({ fromInitialState: true });
+
+    threadA.setState({ x: 1 });
+    await flush();
+    expect(threadA.getState()).toEqual({ x: 1 });
+    expect.soft(threadB.getState()).toEqual({ fromInitialState: true });
+
+    await threadB.append({ ...userMessage("B"), startRun: true });
+    expect
+      .soft(sentStates)
+      .toEqual([{ fromInitialState: true }, { fromInitialState: true }]);
+    expect.soft(threadA.getState()).toEqual({ x: 1 });
+  });
+
+  it("starts a later thread from the shared agent's initial state after a run", async () => {
+    const sentStates: unknown[] = [];
+    const agent = createStateAgent(sentStates);
+    const { core: threadA } = createCore(undefined, agent);
+
+    await threadA.append({ ...userMessage("A"), startRun: true });
+    expect(agent.state).toEqual({ a: 1 });
+
+    const { core: threadB } = createCore(undefined, agent);
+    expect.soft(threadB.getState()).toEqual({ fromInitialState: true });
+    await threadB.append({ ...userMessage("B"), startRun: true });
+    expect(sentStates).toEqual([
+      { fromInitialState: true },
+      { fromInitialState: true },
+    ]);
+  });
+
+  it("copies the initial state before the agent state is mutated", () => {
+    const agent = new HttpAgent({
+      url: "http://agent.test/agent",
+      initialState: { nested: { count: 0 } },
+    });
+    const { core } = createCore(undefined, agent);
+
+    agent.state.nested.count = 1;
+    expect(core.getState()).toEqual({ nested: { count: 0 } });
+
+    const { core: laterThread } = createCore(undefined, agent);
+    expect(laterThread.getState()).toEqual({ nested: { count: 0 } });
+  });
+
+  it("records an undefined initial state without adopting later agent state", () => {
+    const agent = {} as AbstractAgent;
+    createCore(undefined, agent);
+    agent.state = { fromPreviousThread: true };
+
+    const { core } = createCore(undefined, agent);
+    expect(core.getState()).toBeUndefined();
+  });
+
+  it("leaves the state undefined for an agent without initialState", () => {
+    const agent = new HttpAgent({ url: "http://agent.test/agent" });
+    const { core } = createCore(undefined, agent);
+
+    expect(agent.state).toEqual({});
+    expect(core.getState()).toBeUndefined();
+
+    const swapped = new HttpAgent({ url: "http://agent.test/other" });
+    core.updateOptions({
+      agent: swapped,
+      logger: makeLogger(),
+      showThinking: true,
+    });
+    expect(core.getState()).toBeUndefined();
+  });
+
+  it("seeds a new agent's initial state on swap when the thread has no snapshot", () => {
+    const { core } = createCore();
+    const agent = new HttpAgent({
+      url: "http://agent.test/agent",
+      initialState: { a: 1 },
+    });
+    const listener = vi.fn();
+    core.getSnapshot();
+    core.subscribe(listener);
+
+    core.updateOptions({ agent, logger: makeLogger(), showThinking: true });
+
+    expect(agent.state).toEqual({ a: 1 });
+    expect(core.getState()).toEqual({ a: 1 });
+    expect(core.getSnapshot().state).toEqual({ a: 1 });
+    expect(listener).toHaveBeenCalledOnce();
+  });
+
+  it("preserves the thread snapshot on swap and records the new agent's initial state", async () => {
+    const { core } = createCore();
+    core.setState({ fromThread: true });
+    const sentStates: unknown[] = [];
+    const agent = createStateAgent(sentStates);
+
+    core.updateOptions({ agent, logger: makeLogger(), showThinking: true });
+
+    expect(agent.state).toEqual({ fromInitialState: true });
+    expect(core.getState()).toEqual({ fromThread: true });
+    await core.append({ ...userMessage("after swap"), startRun: true });
+    expect(sentStates).toEqual([{ fromThread: true }]);
+
+    const { core: laterThread } = createCore(undefined, agent);
+    expect(laterThread.getState()).toEqual({ fromInitialState: true });
+  });
 });
 
 describe("AgUiThreadController setState", () => {

package/src/runtime/AgUiThreadController.ts

@@ -82,6 +82,26 @@ import {
   withToolApprovalDecision,
 } from "./adapter/tool-approval";
 
+// The SDK does not retain initialState.
+const agentInitialStates = new WeakMap<
+  AbstractAgent,
+  ReadonlyJSONValue | undefined
+>();
+
+const getAgentInitialState = (agent: AbstractAgent) => {
+  if (!agentInitialStates.has(agent)) {
+    const state = agent.state;
+    agentInitialStates.set(
+      agent,
+      state === undefined ||
+        (isPlainObject(state) && Object.keys(state).length === 0)
+        ? undefined
+        : (structuredClone(state) as ReadonlyJSONValue),
+    );
+  }
+  return agentInitialStates.get(agent);
+};
+
 // AbstractAgent.runAgent declares two parameters. HttpAgent ignores a third and
 // is cancelled through agent.abortRun(); the run options stay for subclasses
 // that inherit the base no-op abortRun and have no other cancellation hook.
@@ -280,22 +300,35 @@ const replaceTextParts = (
   return result;
 };
 
+/**
+ * Settles nested sub-agent messages with `status`. Without `runSubagentIds`,
+ * only running messages settle. With it, only that run's sub-agent messages
+ * settle, and an incomplete `status` also settles one still waiting on its own
+ * tool call: the run that could answer it has ended.
+ */
 const settleNestedContent = (
   content: ThreadAssistantMessage["content"],
   status: MessageStatus,
+  runSubagentIds?: ReadonlySet<string>,
 ) =>
-  mapToolCallPartsDeep(content, (part) =>
-    part.messages
-      ? {
-          ...part,
-          messages: part.messages.map((message) =>
-            message.role === "assistant" && message.status?.type === "running"
-              ? { ...message, status }
-              : message,
-          ),
-        }
-      : part,
-  ).content;
+  mapToolCallPartsDeep(content, (part) => {
+    if (!part.messages) return part;
+    let changed = false;
+    const messages = part.messages.map((message) => {
+      if (message.role !== "assistant") return message;
+      const settles = !runSubagentIds
+        ? message.status?.type === "running"
+        : runSubagentIds.has(message.id) &&
+          (message.status?.type === "running" ||
+            (status.type === "incomplete" &&
+              message.status?.type === "requires-action" &&
+              message.status.reason === "tool-calls"));
+      if (!settles) return message;
+      changed = true;
+      return { ...message, status };
+    });
+    return changed ? { ...part, messages } : part;
+  }).content;
 
 type AgUiRuntimeCallbackName = "onError" | "onCancel";
 
@@ -372,6 +405,8 @@ export class AgUiThreadController {
   // The agent that started the active run. updateOptions can swap this.agent
   // mid-run, and cancelling has to reach the agent holding the live request.
   private activeRunAgent: AbstractAgent | null = null;
+  // Counts runs started, so a finished run can tell a newer one began.
+  private runCount = 0;
   private history: ThreadHistoryAdapter | undefined;
   private readonly assistantHistoryParents = new Map<string, string | null>();
   private readonly subagentParentToolCallIds = new Map<
@@ -380,12 +415,22 @@ export class AgUiThreadController {
   >();
   private readonly snapshotHistoryIds = new Set<string>();
   private readonly persistedHistoryIds = new Set<string>();
+  // Ids whose append has not succeeded yet. A write queued behind that append
+  // runs as an append too, since an update may not move the stored head.
+  private readonly unappendedHistoryIds = new Set<string>();
   private readonly historyWrites = createHistoryWriteChain();
   private _historyImportPromise: Promise<void> | undefined;
   private _loadPromise: Promise<void> | undefined;
 
   constructor(options: CoreOptions) {
     this.agent = options.agent;
+    const initialState = getAgentInitialState(this.agent);
+    if (initialState !== undefined) {
+      this.dispatch({
+        type: "stateSnapshot",
+        snapshot: initialState,
+      });
+    }
     this.logger = options.logger;
     this.showThinking = options.showThinking;
     this.joinStrategy = options.joinStrategy;
@@ -421,7 +466,18 @@ export class AgUiThreadController {
   }
 
   updateOptions(options: CoreOptions) {
+    const agentChanged = this.agent !== options.agent;
     this.agent = options.agent;
+    if (agentChanged) {
+      const initialState = getAgentInitialState(this.agent);
+      if (
+        this.state.stateSnapshot === undefined &&
+        initialState !== undefined
+      ) {
+        this.dispatch({ type: "stateSnapshot", snapshot: initialState });
+        this.publishUpdate();
+      }
+    }
     this.logger = options.logger;
     this.showThinking = options.showThinking;
     this.joinStrategy = options.joinStrategy;
@@ -487,6 +543,7 @@ export class AgUiThreadController {
       this.subagentParentToolCallIds.clear();
       this.snapshotHistoryIds.clear();
       this.persistedHistoryIds.clear();
+      this.unappendedHistoryIds.clear();
       for (const { message } of repo.messages) {
         this.persistedHistoryIds.add(message.id);
       }
@@ -1683,6 +1740,8 @@ export class AgUiThreadController {
     const assistantParentId = parent ? parentId : this.session.headId;
     let assistantMessageId: string | undefined;
     let snapshotAssistantId: string | undefined;
+    const importedRunAssistants = new Set<string>();
+    const runSubagentIds = new Set<string>();
     let replayInterruptMessageId: string | undefined;
     // A snapshot the preserve gate declines still evicts the in-flight
     // assistant; recreating under the cached id on the next content-bearing
@@ -2225,9 +2284,38 @@ export class AgUiThreadController {
         adoptServerMessageId(serverId, true, frameOpenIds);
       },
     });
+    const settleRunSubagents = (status: MessageStatus) => {
+      if (runSubagentIds.size === 0) return;
+      for (const message of this.session.getMessages()) {
+        if (message.role !== "assistant") continue;
+        const content = settleNestedContent(
+          message.content,
+          status,
+          runSubagentIds,
+        );
+        if (content === message.content) continue;
+        this.updateAssistantMessage(message.id, { content });
+        this.persistSettledAssistant(message.id);
+      }
+    };
+    // Writes this run's other assistants in branch order: before the answer
+    // settles, the ones above it; after, the ones below. An adapter that
+    // moves its head on each first write then ends on the branch's tail.
+    const persistRunAssistants = (aboveAnswer: boolean) => {
+      for (const message of this.session.getMessages()) {
+        if (message.id === assistantMessageId) {
+          if (aboveAnswer) return;
+          continue;
+        }
+        if (message.role === "assistant" && !runStartMessageIds.has(message.id))
+          this.persistAssistantHistory(message.id);
+      }
+    };
     let runFinished = false;
     const dispatch = (event: AgUiEvent) => {
       if (this.abortController !== abortController) return;
+      if (event.type === "SUBAGENT_STARTED")
+        runSubagentIds.add(event.subagentRunId);
       if (event.type === "RUN_FINISHED") {
         runFinished = true;
         if (event.outcome?.type === "interrupt")
@@ -2250,6 +2338,7 @@ export class AgUiThreadController {
         event.type === "RUN_FINISHED" ||
         event.type === "RUN_ERROR" ||
         event.type === "RUN_CANCELLED";
+      if (settlingRun) persistRunAssistants(true);
       const beforeSnapshot =
         event.type === "MESSAGES_SNAPSHOT" ? this.session.getMessages() : [];
       try {
@@ -2263,6 +2352,31 @@ export class AgUiThreadController {
           candidateContexts.size ? reconcileReplaySnapshot : undefined,
         );
         if (event.type === "MESSAGES_SNAPSHOT") {
+          // Assistants a snapshot adds belong to the run's history even when
+          // the run never streams into them. Every import resets the pending
+          // marks, so each snapshot marks them again. The import also routes
+          // every message to update; a run assistant this client never wrote
+          // gets an append first, which is what moves the stored head.
+          const previousIds = new Set(beforeSnapshot.map(({ id }) => id));
+          for (const message of this.session.getMessages()) {
+            if (
+              message.role !== "assistant" ||
+              runStartMessageIds.has(message.id)
+            )
+              continue;
+            if (!previousIds.has(message.id))
+              importedRunAssistants.add(message.id);
+            if (
+              !this.persistedHistoryIds.has(message.id) &&
+              this.historyWrites.pending(message.id) === undefined
+            )
+              this.snapshotHistoryIds.delete(message.id);
+          }
+          for (const messageId of importedRunAssistants) {
+            const item = this.session.tryGetMessage(messageId);
+            if (item)
+              this.markPendingAssistantHistory(messageId, item.parentId);
+          }
           for (const message of beforeSnapshot) {
             const context = replayContexts.get(message.id);
             const current = this.session.tryGetMessage(message.id)?.message;
@@ -2294,6 +2408,16 @@ export class AgUiThreadController {
               tail?.role === "assistant" ? tail.id : undefined;
           }
         }
+        if (event.type === "RUN_ERROR") {
+          settleRunSubagents({
+            type: "incomplete",
+            reason: "error",
+            ...(event.message !== undefined ? { error: event.message } : {}),
+          });
+        } else if (event.type === "RUN_CANCELLED") {
+          settleRunSubagents({ type: "incomplete", reason: "cancelled" });
+        }
+        if (settlingRun) persistRunAssistants(false);
       } finally {
         settlingRun = false;
       }
@@ -2302,6 +2426,7 @@ export class AgUiThreadController {
     this.abortController = abortController;
     const runAgentInstance = this.agent;
     this.activeRunAgent = runAgentInstance;
+    const runNumber = ++this.runCount;
 
     let cancelRun = () => dispatch({ type: "RUN_CANCELLED" });
     abortSignal.addEventListener(
@@ -2385,6 +2510,7 @@ export class AgUiThreadController {
         current?.role === "assistant" &&
         current.status?.type === "running"
       ) {
+        persistRunAssistants(true);
         // Finalization is transport cleanup, not proof of a terminal event.
         const status: MessageStatus = {
           type: "incomplete",
@@ -2397,6 +2523,25 @@ export class AgUiThreadController {
           content: settleNestedContent(current.content, status),
         });
       }
+      if (this.abortController === abortController && !runFinished) {
+        // The run ended without RUN_FINISHED, so no sub-agent call it left
+        // open can still be answered.
+        const settled = assistantMessageId
+          ? this.session.tryGetMessage(assistantMessageId)?.message
+          : undefined;
+        settleRunSubagents(
+          settled?.status?.type === "incomplete"
+            ? settled.status
+            : {
+                type: "incomplete",
+                reason: "error",
+                error: "The connection ended before the answer finished.",
+              },
+        );
+      }
+      // Stop may have finished this run already; a newer run owns its own
+      // writes.
+      if (this.runCount === runNumber) persistRunAssistants(false);
       this.finishRun(abortController);
     }
 
@@ -2735,11 +2880,14 @@ export class AgUiThreadController {
       oldId,
       collidesWithExisting ? undefined : newId,
     );
+    if (this.unappendedHistoryIds.delete(oldId) && !collidesWithExisting)
+      this.unappendedHistoryIds.add(newId);
     if (pendingWrite && !collidesWithExisting) {
       void pendingWrite.then(
         () => {
           this.persistedHistoryIds.delete(oldId);
           this.persistedHistoryIds.add(newId);
+          this.unappendedHistoryIds.delete(newId);
         },
         () => {},
       );
@@ -2775,7 +2923,7 @@ export class AgUiThreadController {
                 update.content as ThreadAssistantMessage["content"],
               )
           : assistant.content;
-      const calls = replay ? Array.from(iterateToolCallParts(content)) : [];
+      const calls = Array.from(iterateToolCallParts(content));
       // A retained historical call needs existing or current stream authority.
       const pending =
         pendingToolCallIds ??
@@ -2789,7 +2937,11 @@ export class AgUiThreadController {
           : undefined);
       if (replay && latestStatus?.type !== "running" && pending)
         replay.admittedToolCallIds = pending;
-      if (replay) {
+      if (
+        replay ||
+        (latestStatus?.type === "requires-action" &&
+          latestStatus.reason === "interrupt")
+      ) {
         // Project gates against the content the user can actually answer.
         const interrupts =
           latestStatus?.type === "requires-action" &&
@@ -2805,10 +2957,10 @@ export class AgUiThreadController {
           new Set(calls.map((part) => part.toolCallId)),
         );
         const originalCalls = new Map(
-          Array.from(iterateToolCallParts(replay.originalContent), (part) => [
-            part.toolCallId,
-            part,
-          ]),
+          Array.from(
+            iterateToolCallParts(replay?.originalContent ?? assistant.content),
+            (part) => [part.toolCallId, part],
+          ),
         );
         content = mapToolCallPartsDeep(content, (part) => {
           const approval = approvals.get(part.toolCallId);
@@ -3301,6 +3453,9 @@ export class AgUiThreadController {
             }
           }
           activeAssistantId = nextActiveAssistantId;
+          if (activeAssistantId && joinStrategy !== "concat-content") {
+            aggregator.rebaseMessage(activeAssistantId);
+          }
           if (hasNewPendingToolCallAnswer && head) {
             this.updateAssistantMessage(head.id, {
               status: { type: "running" },
@@ -3592,10 +3747,28 @@ export class AgUiThreadController {
             previous?.status?.type === "running" &&
             snapshot.status?.type === "requires-action" &&
             snapshot.status.reason === "tool-calls";
+          const previousToolCallIds = new Set(
+            Array.from(
+              iterateToolCallParts(previous?.content ?? []),
+              (call) => call.toolCallId,
+            ),
+          );
+          const preservesSettledStatus =
+            convertedMessage.id !== activeAssistantId &&
+            this.isTerminalStatus(previous?.status) &&
+            snapshot.status?.type === "requires-action" &&
+            snapshot.status.reason === "tool-calls" &&
+            Array.from(iterateToolCallParts(snapshot.content)).every(
+              (call) =>
+                isResolvedToolCall(call) ||
+                previousToolCallIds.has(call.toolCallId),
+            );
           converted[index] = previous
             ? {
                 ...snapshot,
-                ...(preservesRunningStatus ? { status: previous.status } : {}),
+                ...(preservesRunningStatus || preservesSettledStatus
+                  ? { status: previous.status }
+                  : {}),
                 content: this.mergeSnapshotAssistantContent(
                   previous.content,
                   snapshot.content,
@@ -3796,7 +3969,9 @@ export class AgUiThreadController {
 
     if (shouldUpdate) {
       const write = this.historyWrites.chain(messageId, () =>
-        update.call(history, { parentId, message }),
+        this.unappendedHistoryIds.has(messageId)
+          ? this.appendOnce(history, messageId, { parentId, message })
+          : update.call(history, { parentId, message }),
       );
       this.assistantHistoryParents.delete(messageId);
       void write.then(
@@ -3834,6 +4009,15 @@ export class AgUiThreadController {
     );
   }
 
+  private async appendOnce(
+    history: ThreadHistoryAdapter,
+    messageId: string,
+    item: { parentId: string | null; message: ThreadMessage },
+  ): Promise<void> {
+    await history.append(item);
+    this.unappendedHistoryIds.delete(messageId);
+  }
+
   private appendHistoryItem(
     parentId: string | null,
     message: ThreadMessage,
@@ -3842,9 +4026,10 @@ export class AgUiThreadController {
     const pending = this.historyWrites.pending(message.id);
     if (pending) return pending;
 
-    const append = this.history.append.bind(this.history);
+    const history = this.history;
+    this.unappendedHistoryIds.add(message.id);
     const write = this.historyWrites.chain(message.id, () =>
-      append({ parentId, message }),
+      this.appendOnce(history, message.id, { parentId, message }),
     );
     void write.then(
       () => {

package/src/runtime/adapter/run-aggregator.ts

@@ -267,6 +267,17 @@ export class RunAggregator {
     return this.toolCalls.has(toolCallId);
   }
 
+  rebaseMessage(messageId: string): void {
+    const snapshotText = this.pendingSnapshotText.get(messageId);
+    this.resetMessageParts();
+    this.lastTextMessageId = messageId;
+    this.reportedServerMessageId = messageId;
+    if (snapshotText !== undefined) {
+      this.pendingSnapshotText.set(messageId, snapshotText);
+    }
+    this.startTextMessage(ROOT_SCOPE, messageId);
+  }
+
   handle(event: AgUiEvent): void {
     switch (event.type) {
       case "RUN_STARTED": {