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

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

@@ -0,0 +1,226 @@
+import { describe, expect, it } from "vitest";
+import type { AppendMessage } from "@assistant-ui/core";
+import { HttpAgent } from "@ag-ui/client";
+import { AgUiThreadController } from "./AgUiThreadController";
+import { makeLogger } from "./logger";
+
+const tick = () => new Promise((r) => setTimeout(r, 0));
+const settle = async () => {
+  for (let i = 0; i < 8; i++) await tick();
+};
+const userMessage = {
+  role: "user",
+  content: [{ type: "text", text: "hello" }],
+  attachments: [],
+  metadata: { custom: {} },
+  createdAt: new Date(0),
+  parentId: null,
+  sourceId: null,
+  runConfig: {},
+  startRun: true,
+} as unknown as AppendMessage;
+
+type Row = { id: string; status: string | undefined };
+
+// Follows the ThreadHistoryAdapter contract: both writes upsert by message
+// id, and only append moves the stored head.
+function createHistory(writes: string[]) {
+  const items = new Map<string, any>();
+  let headId: string | null = null;
+  const store = (item: any, moveHead: boolean) => {
+    if (moveHead) headId = item.message.id;
+    items.set(item.message.id, structuredClone(item));
+  };
+  return {
+    rows: () =>
+      new Map<string, Row>(
+        [...items.values()].map((item) => [
+          item.message.id,
+          { id: item.message.id, status: item.message.status?.type },
+        ]),
+      ),
+    load: async () => ({ headId, messages: [...items.values()] }),
+    append: async (item: any) => {
+      writes.push(`append ${item.message.id}`);
+      store(item, true);
+    },
+    update: async (item: any) => {
+      writes.push(`update ${item.message.id}`);
+      store(item, false);
+    },
+  };
+}
+
+async function reload(history: ReturnType<typeof createHistory>) {
+  const core = new AgUiThreadController({
+    agent: new HttpAgent({ url: "https://agent.test" }),
+    logger: makeLogger(),
+    showThinking: false,
+    history: history as any,
+  });
+  await core.__internal_load();
+  return core.getMessages().map((m: any) => `${m.id}:${m.status?.type ?? "-"}`);
+}
+
+async function run(events: (requestMessages: any[]) => object[][]) {
+  const writes: string[] = [];
+  const history = createHistory(writes);
+  let write!: (e: object) => void;
+  let close!: () => void;
+  let requestMessages: any[] = [];
+  const agent = new HttpAgent({
+    url: "https://agent.test",
+    fetch: async (_url: unknown, init: RequestInit) => {
+      const encoder = new TextEncoder();
+      let controller!: ReadableStreamDefaultController<Uint8Array>;
+      const body = new ReadableStream<Uint8Array>({
+        start: (c) => void (controller = c),
+      });
+      requestMessages = JSON.parse(init.body as string).messages ?? [];
+      write = (e) =>
+        controller.enqueue(encoder.encode(`data: ${JSON.stringify(e)}\n\n`));
+      close = () => controller.close();
+      return new Response(body, {
+        headers: { "Content-Type": "text/event-stream" },
+      });
+    },
+  });
+  const core = new AgUiThreadController({
+    agent,
+    logger: makeLogger(),
+    showThinking: false,
+    history: history as any,
+  });
+  const sending = core.append(userMessage).catch(() => {});
+  for (let i = 0; i < 50 && !write; i++) await tick();
+  for (const batch of events(requestMessages)) {
+    for (const e of batch) write(e);
+    await settle();
+  }
+  close();
+  await sending;
+  await settle();
+  const shown = core
+    .getMessages()
+    .map((m: any) => `${m.id}:${m.status?.type ?? "-"}`);
+  return {
+    rows: history.rows(),
+    writes,
+    shown,
+    reloaded: await reload(history),
+  };
+}
+
+const start = [
+  { type: "RUN_STARTED", threadId: "t", runId: "r" },
+  { type: "TEXT_MESSAGE_START", messageId: "answer", role: "assistant" },
+  { type: "TEXT_MESSAGE_CONTENT", messageId: "answer", delta: "a1" },
+];
+const end = { type: "TEXT_MESSAGE_END", messageId: "answer" };
+const moveSnapshot = (req: any[]) => ({
+  type: "MESSAGES_SNAPSHOT",
+  messages: [
+    ...req,
+    { id: "answer", role: "assistant", content: "a1" },
+    { id: "other", role: "assistant", content: "o" },
+  ],
+});
+
+describe("ag-ui: a snapshot's new head and the history adapter", () => {
+  it("stores the new head after one head-move snapshot and RUN_FINISHED", async () => {
+    const { rows, writes, shown, reloaded } = await run((req) => [
+      start,
+      [moveSnapshot(req)],
+      [end, { type: "RUN_FINISHED", threadId: "t", runId: "r" }],
+    ]);
+    expect(shown).toContain("other:complete");
+    expect(rows.get("other")?.status, writes.join(", ")).toBe("complete");
+    expect(reloaded.slice(1), writes.join(", ")).toEqual(shown.slice(1));
+  });
+
+  it("stores the new head after one head-move snapshot and RUN_ERROR", async () => {
+    const { rows, writes, shown, reloaded } = await run((req) => [
+      start,
+      [moveSnapshot(req)],
+      [{ type: "RUN_ERROR", message: "boom" }],
+    ]);
+    // Which message takes the error depends on whether the run moved onto
+    // "other"; either way history holds what the thread shows.
+    const other = shown.find((entry) => entry.startsWith("other:"));
+    expect(["other:incomplete", "other:complete"]).toContain(other);
+    expect(`other:${rows.get("other")?.status}`, writes.join(", ")).toBe(other);
+    expect(reloaded.slice(1), writes.join(", ")).toEqual(shown.slice(1));
+  });
+
+  it("stores a snapshot-only answer that a second snapshot repeats", async () => {
+    const snapshot = (req: any[]) => ({
+      type: "MESSAGES_SNAPSHOT",
+      messages: [...req, { id: "only", role: "assistant", content: "o" }],
+    });
+    const { rows, writes, shown, reloaded } = await run((req) => [
+      [{ type: "RUN_STARTED", threadId: "t", runId: "r" }, snapshot(req)],
+      [snapshot(req)],
+      [{ type: "RUN_FINISHED", threadId: "t", runId: "r" }],
+    ]);
+    expect(rows.get("only")?.status, writes.join(", ")).toBe("complete");
+    expect(reloaded.slice(1), writes.join(", ")).toEqual(shown.slice(1));
+  });
+
+  it("stores an answer that arrives only in a snapshot", async () => {
+    const { rows, writes, shown, reloaded } = await run((req) => [
+      [
+        { type: "RUN_STARTED", threadId: "t", runId: "r" },
+        {
+          type: "MESSAGES_SNAPSHOT",
+          messages: [...req, { id: "only", role: "assistant", content: "o" }],
+        },
+      ],
+      [{ type: "RUN_FINISHED", threadId: "t", runId: "r" }],
+    ]);
+    expect(shown).toContain("only:complete");
+    expect(rows.get("only")?.status, writes.join(", ")).toBe("complete");
+    expect(writes.filter((w) => w.endsWith(" only"))).toHaveLength(1);
+    expect(reloaded.slice(1), writes.join(", ")).toEqual(shown.slice(1));
+  });
+
+  it("keeps the streamed answer as the head after a snapshot adds an earlier step", async () => {
+    const { writes, shown, reloaded } = await run((req) => [
+      start,
+      [
+        {
+          type: "MESSAGES_SNAPSHOT",
+          messages: [
+            ...req,
+            { id: "step", role: "assistant", content: "s" },
+            { id: "answer", role: "assistant", content: "a1" },
+          ],
+        },
+      ],
+      [end, { type: "RUN_FINISHED", threadId: "t", runId: "r" }],
+    ]);
+    expect(shown.slice(1)).toEqual(["step:complete", "answer:complete"]);
+    expect(reloaded.slice(1), writes.join(", ")).toEqual(shown.slice(1));
+  });
+
+  it("control: a second snapshot after the head move stores the new head", async () => {
+    const { rows, writes, shown, reloaded } = await run((req) => [
+      start,
+      [moveSnapshot(req)],
+      [moveSnapshot(req)],
+      [end, { type: "RUN_FINISHED", threadId: "t", runId: "r" }],
+    ]);
+    expect(rows.get("other")?.status, JSON.stringify({ writes, shown })).toBe(
+      "complete",
+    );
+    expect(reloaded.slice(1), writes.join(", ")).toEqual(shown.slice(1));
+  });
+
+  it("control: the streamed answer without a snapshot is stored", async () => {
+    const { rows, shown, reloaded } = await run(() => [
+      start,
+      [end, { type: "RUN_FINISHED", threadId: "t", runId: "r" }],
+    ]);
+    expect(rows.get("answer")?.status).toBe("complete");
+    expect(reloaded).toEqual(shown);
+  });
+});

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.snapshot-nested-settle.test.ts

@@ -0,0 +1,210 @@
+import { describe, expect, it } from "vitest";
+import type { AppendMessage } from "@assistant-ui/core";
+import { HttpAgent } from "@ag-ui/client";
+import { AgUiThreadController } from "./AgUiThreadController";
+import { makeLogger } from "./logger";
+
+const tick = () => new Promise((r) => setTimeout(r, 0));
+const settle = async () => {
+  for (let i = 0; i < 8; i++) await tick();
+};
+const userMessage = {
+  role: "user",
+  content: [{ type: "text", text: "hello" }],
+  attachments: [],
+  metadata: { custom: {} },
+  createdAt: new Date(0),
+  parentId: null,
+  sourceId: null,
+  runConfig: {},
+  startRun: true,
+} as unknown as AppendMessage;
+
+type End = "close" | "error" | "stop" | "finish";
+
+async function run(end: End, withSnapshot = true, history?: unknown) {
+  let write!: (e: object) => void;
+  let close!: () => void;
+  let requestMessages: any[] = [];
+  const agent = new HttpAgent({
+    url: "https://agent.test",
+    fetch: async (_url: unknown, init: RequestInit) => {
+      const encoder = new TextEncoder();
+      let controller!: ReadableStreamDefaultController<Uint8Array>;
+      let closed = false;
+      const body = new ReadableStream<Uint8Array>({
+        start: (c) => void (controller = c),
+      });
+      requestMessages = JSON.parse(init.body as string).messages ?? [];
+      write = (e) => {
+        if (!closed)
+          controller.enqueue(encoder.encode(`data: ${JSON.stringify(e)}\n\n`));
+      };
+      close = () => {
+        if (closed) return;
+        closed = true;
+        controller.close();
+      };
+      init.signal?.addEventListener("abort", () => {
+        if (closed) return;
+        closed = true;
+        try {
+          controller.error(new DOMException("aborted", "AbortError"));
+        } catch {}
+      });
+      return new Response(body, {
+        headers: { "Content-Type": "text/event-stream" },
+      });
+    },
+  });
+  const core = new AgUiThreadController({
+    agent,
+    logger: makeLogger(),
+    showThinking: false,
+    ...(history ? { history: history as any } : {}),
+  });
+  const sending = core.append(userMessage).catch(() => {});
+  for (let i = 0; i < 50 && !write; i++) await tick();
+  const events: object[] = [
+    { type: "RUN_STARTED", threadId: "t", runId: "r" },
+    { type: "TEXT_MESSAGE_START", messageId: "answer", role: "assistant" },
+    { type: "TEXT_MESSAGE_CONTENT", messageId: "answer", delta: "a1" },
+    { type: "TEXT_MESSAGE_END", messageId: "answer" },
+    { type: "TOOL_CALL_START", toolCallId: "task", toolCallName: "task" },
+    { type: "TOOL_CALL_END", toolCallId: "task" },
+    {
+      type: "SUBAGENT_STARTED",
+      subagentRunId: "sub",
+      name: "worker",
+      parentToolCallId: "task",
+    },
+    {
+      type: "TOOL_CALL_START",
+      toolCallId: "ncall",
+      toolCallName: "search",
+      subagentRunId: "sub",
+    },
+    { type: "TOOL_CALL_END", toolCallId: "ncall", subagentRunId: "sub" },
+    { type: "SUBAGENT_FINISHED", subagentRunId: "sub" },
+  ];
+  for (const e of events) write(e);
+  await settle();
+  if (withSnapshot) {
+    write({
+      type: "MESSAGES_SNAPSHOT",
+      messages: [
+        ...requestMessages,
+        {
+          id: "answer",
+          role: "assistant",
+          content: "a1",
+          toolCalls: [
+            {
+              id: "task",
+              type: "function",
+              function: { name: "task", arguments: "{}" },
+            },
+          ],
+        },
+        {
+          id: "s-call",
+          role: "assistant",
+          content: "",
+          subagentRunId: "sub",
+          toolCalls: [
+            {
+              id: "ncall",
+              type: "function",
+              function: { name: "search", arguments: "{}" },
+            },
+          ],
+        },
+      ],
+    });
+    await settle();
+  }
+  const before = nested(core);
+  if (end === "close") close();
+  else if (end === "error")
+    write({ type: "RUN_ERROR", message: "server failed" });
+  else if (end === "finish")
+    write({ type: "RUN_FINISHED", threadId: "t", runId: "r" });
+  else await core.cancel();
+  await settle();
+  close();
+  await sending;
+  await settle();
+  return { before, after: nested(core) };
+}
+
+function nested(core: AgUiThreadController) {
+  const answer = core.getMessages().find((m: any) => m.id === "answer") as any;
+  const call = answer?.content.find(
+    (p: any) => p.type === "tool-call" && p.toolCallId === "task",
+  );
+  const sub = call?.messages?.find((m: any) => m.id === "sub");
+  return {
+    answer: answer?.status?.type,
+    sub: sub?.status?.type,
+    subReason: sub?.status?.reason,
+    ncallOpen: sub?.content.some(
+      (p: any) => p.toolCallId === "ncall" && p.result === undefined,
+    ),
+  };
+}
+
+describe("ag-ui: a nested message waiting on its own call after the run ends", () => {
+  for (const end of ["close", "error", "stop"] as const) {
+    it(`settles the nested message when the run ends by ${end}`, async () => {
+      const { before, after } = await run(end);
+      // the snapshot imported the sub-agent's message waiting on ncall
+      expect(before.sub, JSON.stringify(before)).toBe("requires-action");
+      expect(after.answer, JSON.stringify(after)).toBe("incomplete");
+      expect(after.sub, JSON.stringify(after)).not.toBe("requires-action");
+    });
+  }
+
+  it("control: after RUN_FINISHED the nested call stays open and is not settled as an error", async () => {
+    const { after } = await run("finish");
+    expect(after.answer, JSON.stringify(after)).toBe("requires-action");
+    expect(after.ncallOpen, JSON.stringify(after)).toBe(true);
+    expect(after.subReason, JSON.stringify(after)).not.toBe("error");
+  });
+
+  it("control: without the snapshot the nested message is the stream's (complete) and settles", async () => {
+    const { before, after } = await run("close", false);
+    expect(before.sub, JSON.stringify(before)).toBe("complete");
+    expect(after.sub, JSON.stringify(after)).not.toBe("requires-action");
+  });
+
+  it("appends the answer again when its first append fails under a queued update", async () => {
+    const writes: string[] = [];
+    let headId: string | null = null;
+    const rows = new Set<string>();
+    let failed = false;
+    const history = {
+      load: async () => ({ headId: null, messages: [] }),
+      append: async (item: any) => {
+        const id = item.message.id;
+        writes.push(`append ${id}`);
+        if (id === "answer" && !failed) {
+          failed = true;
+          await tick();
+          throw new Error("write failed");
+        }
+        rows.add(id);
+        headId = id;
+      },
+      // Follows the ThreadHistoryAdapter contract: update does not move the head.
+      update: async (item: any) => {
+        writes.push(`update ${item.message.id}`);
+        rows.add(item.message.id);
+      },
+    };
+    const { after } = await run("error", true, history);
+    expect(after.sub, JSON.stringify(after)).not.toBe("requires-action");
+    expect(writes, writes.join(", ")).not.toContain("update answer");
+    expect(rows.has("answer"), writes.join(", ")).toBe(true);
+    expect(headId, writes.join(", ")).toBe("answer");
+  });
+});