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

package/src/runtime/AgUiThreadController.snapshot-head-move.test.ts

@@ -58,9 +58,12 @@ const answerRecord = {
   ],
 };
 const otherRecord = { id: "other", role: "assistant", content: "Next step." };
-const taskAndSubagentEvents = [
+const taskCallEvents = [
   { type: "TOOL_CALL_START", toolCallId: "task", toolCallName: "task" },
   { type: "TOOL_CALL_END", toolCallId: "task" },
+];
+const taskAndSubagentEvents = [
+  ...taskCallEvents,
   {
     type: "SUBAGENT_STARTED",
     subagentRunId: "sub",
@@ -106,7 +109,53 @@ function statuses(core: AgUiThreadController) {
   return { answer: answer.status?.type, subagent: nested?.status?.type };
 }
 
+function taskPartFor(core: AgUiThreadController, messageId: string) {
+  const message = core.getMessages().find(({ id }) => id === messageId);
+  if (message?.role !== "assistant") return undefined;
+  return message.content.find(
+    (part) => part.type === "tool-call" && part.toolCallId === "task",
+  );
+}
+
 describe("AgUiThreadController snapshot head moves", () => {
+  it("reopens a settled snapshot-only assistant when a later snapshot adds a tool call", async () => {
+    const answer = {
+      id: "m1",
+      role: "assistant",
+      content: "Let me ask.",
+    };
+    const core = createController((wire) => [
+      start,
+      { type: "MESSAGES_SNAPSHOT", messages: [...wire.messages, answer] },
+      {
+        type: "MESSAGES_SNAPSHOT",
+        messages: [
+          ...wire.messages,
+          {
+            ...answer,
+            toolCalls: [
+              {
+                id: "ask",
+                type: "function",
+                function: { name: "ask_user", arguments: "{}" },
+              },
+            ],
+          },
+        ],
+      },
+      finished,
+    ]);
+
+    await core.append(userMessage);
+
+    expect(
+      core.getMessages().find(({ id }) => id === "m1")?.status,
+    ).toMatchObject({
+      type: "requires-action",
+      reason: "tool-calls",
+    });
+  });
+
   it("settles the moved-off answer when the snapshot has no tool call", async () => {
     const core = createController((wire) => [
       start,
@@ -143,6 +192,136 @@ describe("AgUiThreadController snapshot head moves", () => {
     });
   });
 
+  it("keeps a moved-off answer settled and owns the result after repeated snapshots", async () => {
+    const core = createController((wire) => [
+      start,
+      ...answerText,
+      ...taskCallEvents,
+      snapshot(wire, otherRecord),
+      snapshot(wire, otherRecord),
+      taskResult,
+      finished,
+    ]);
+    await core.append(userMessage);
+
+    expect(core.isRunning()).toBe(false);
+    expect(statuses(core).answer).toBe("complete");
+    expect(
+      core.getMessages().find(({ id }) => id === "other")?.status?.type,
+    ).toBe("complete");
+    expect(taskPartFor(core, "answer")).toMatchObject({ result: "done" });
+    expect(taskPartFor(core, "other")).toBeUndefined();
+    expect(
+      core.getMessages().find(({ id }) => id === "answer")?.content,
+    ).toEqual([
+      { type: "text", text: "Let me check." },
+      {
+        type: "tool-call",
+        toolCallId: "task",
+        toolName: "task",
+        args: {},
+        argsText: "{}",
+        result: "done",
+        unstable_toolMessageId: "tool-task",
+      },
+    ]);
+    expect(
+      core.getMessages().find(({ id }) => id === "other")?.content,
+    ).toEqual([{ type: "text", text: "Next step." }]);
+  });
+
+  it("preserves full content and settled statuses with a sub-agent after repeated snapshots", async () => {
+    const core = createController((wire) => [
+      start,
+      ...answerText,
+      ...taskAndSubagentEvents,
+      snapshot(wire, otherRecord),
+      snapshot(wire, otherRecord),
+      subFinished,
+      taskResult,
+      finished,
+    ]);
+    await core.append(userMessage);
+
+    expect(core.isRunning()).toBe(false);
+    expect(statuses(core)).toEqual({
+      answer: "complete",
+      subagent: "complete",
+    });
+    expect(
+      core.getMessages().find(({ id }) => id === "other")?.status?.type,
+    ).toBe("complete");
+    expect(
+      core.getMessages().find(({ id }) => id === "answer")?.content,
+    ).toEqual([
+      { type: "text", text: "Let me check." },
+      {
+        type: "tool-call",
+        toolCallId: "task",
+        toolName: "task",
+        args: {},
+        argsText: "{}",
+        result: "done",
+        unstable_toolMessageId: "tool-task",
+        messages: [
+          expect.objectContaining({
+            id: "sub",
+            role: "assistant",
+            status: { type: "complete", reason: "unknown" },
+            content: [{ type: "text", text: "working" }],
+          }),
+        ],
+      },
+    ]);
+    expect(
+      core.getMessages().find(({ id }) => id === "other")?.content,
+    ).toEqual([{ type: "text", text: "Next step." }]);
+  });
+
+  it.each([false, true])(
+    "preserves current text when the moved result arrives during streaming: %s",
+    async (duringStreaming) => {
+      const core = createController((wire) => [
+        start,
+        ...answerText,
+        ...taskCallEvents,
+        snapshot(wire, otherRecord),
+        snapshot(wire, otherRecord),
+        ...(!duringStreaming ? [taskResult] : []),
+        { type: "TEXT_MESSAGE_START", messageId: "other", role: "assistant" },
+        { type: "TEXT_MESSAGE_CONTENT", messageId: "other", delta: " Before." },
+        ...(duringStreaming ? [taskResult] : []),
+        { type: "TEXT_MESSAGE_CONTENT", messageId: "other", delta: " After." },
+        { type: "TEXT_MESSAGE_END", messageId: "other" },
+        finished,
+      ]);
+      await core.append(userMessage);
+
+      expect(core.isRunning()).toBe(false);
+      expect(statuses(core).answer).toBe("complete");
+      expect(
+        core.getMessages().find(({ id }) => id === "answer")?.content,
+      ).toEqual([
+        { type: "text", text: "Let me check." },
+        {
+          type: "tool-call",
+          toolCallId: "task",
+          toolName: "task",
+          args: {},
+          argsText: "{}",
+          result: "done",
+          unstable_toolMessageId: "tool-task",
+        },
+      ]);
+      expect(core.getMessages().find(({ id }) => id === "other")).toMatchObject(
+        {
+          status: { type: "complete" },
+          content: [{ type: "text", text: "Next step. Before. After." }],
+        },
+      );
+    },
+  );
+
   it("does not show a complete answer with a running sub-agent after a result snapshot", async () => {
     const core = createController((wire) => [
       start,

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.
@@ -386,6 +406,13 @@ export class AgUiThreadController {
 
   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 +448,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;
@@ -2775,7 +2813,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 +2827,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 +2847,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 +3343,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 +3637,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,

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": {