@assistant-ui/react-ag-ui 0.1.5 → 0.1.7

package/src/runtime/AgUiThreadController.disconnect.test.ts

@@ -234,7 +234,7 @@ describe("AgUiThreadController HTTP stream closure", () => {
     });
   });
 
-  it("settles a nested assistant when a snapshot advances past its step", async () => {
+  it("settles a cut-off answer and its nested assistant when a snapshot advanced past them", async () => {
     const core = createController([
       start,
       ...text,
@@ -294,14 +294,68 @@ describe("AgUiThreadController HTTP stream closure", () => {
     await core.append(userMessage);
 
     const answer = core.getMessages().find(({ id }) => id === "answer");
-    expect(answer?.status?.type).toBe("complete");
+    expect(answer?.status).toMatchObject({
+      type: "incomplete",
+      reason: "error",
+    });
     expect(answer?.content[1]).toMatchObject({
       type: "tool-call",
       toolCallId: "task",
-      messages: [{ id: "sub", status: { type: "complete" } }],
+      messages: [
+        { id: "sub", status: { type: "incomplete", reason: "error" } },
+      ],
     });
+    expect(
+      core.getMessages().find(({ id }) => id === "next")?.status?.type,
+    ).toBe("complete");
   });
 
+  it.each([
+    {
+      name: "mid-text after the move",
+      events: [
+        { type: "TEXT_MESSAGE_CONTENT", messageId: "answer", delta: " more" },
+      ],
+      status: { type: "incomplete", reason: "error" },
+    },
+    {
+      name: "right after the move",
+      events: [],
+      status: { type: "incomplete", reason: "error" },
+    },
+    {
+      name: "after the frame closed",
+      before: [{ type: "TEXT_MESSAGE_END", messageId: "answer" }],
+      events: [],
+      status: { type: "complete" },
+    },
+  ])(
+    "settles a moved-off answer at EOF $name",
+    async ({ before = [], events, status }) => {
+      const core = createController([
+        start,
+        ...text,
+        ...before,
+        {
+          type: "MESSAGES_SNAPSHOT",
+          messages: [
+            { id: "question", role: "user", content: "hello" },
+            { id: "answer", role: "assistant", content: "partial answer" },
+            { id: "next", role: "assistant", content: "next step" },
+          ],
+        },
+        ...events,
+      ]);
+
+      await core.append(userMessage);
+
+      expect(core.isRunning()).toBe(false);
+      expect(
+        core.getMessages().find(({ id }) => id === "answer")?.status,
+      ).toMatchObject(status);
+    },
+  );
+
   it("streams a nested assistant that starts after a snapshot", async () => {
     const core = createController([
       start,
@@ -911,4 +965,145 @@ describe("AgUiThreadController HTTP stream closure", () => {
       });
     },
   );
+
+  it.each(
+    [
+      {
+        ending: "EOF",
+        runFinished: false,
+        expectedStatus: { type: "incomplete", reason: "error" },
+      },
+      {
+        ending: "RUN_FINISHED",
+        runFinished: true,
+        expectedStatus: { type: "requires-action", reason: "tool-calls" },
+      },
+    ].flatMap((row) => [
+      { ...row, shape: "with text", otherContent: "Let me check." },
+      { ...row, shape: "with no text", otherContent: "" },
+    ]),
+  )(
+    "keeps a moved-to assistant $shape running when a later snapshot adds a tool call through $ending",
+    async ({ runFinished, expectedStatus, otherContent }) => {
+      let releaseStream!: () => void;
+      const streamCanClose = new Promise<void>((resolve) => {
+        releaseStream = resolve;
+      });
+      const answer = { id: "answer", role: "assistant", content: "a1" };
+      const other = (withCall: boolean) => ({
+        id: "other",
+        role: "assistant",
+        content: otherContent,
+        ...(withCall && {
+          toolCalls: [
+            {
+              id: "ocall",
+              type: "function",
+              function: { name: "lookup", arguments: "{}" },
+            },
+          ],
+        }),
+      });
+      const agent = new HttpAgent({
+        url: "https://example.invalid",
+        fetch: async (_url: unknown, init: RequestInit) => {
+          const messages = JSON.parse(init.body as string).messages ?? [];
+          const events = [
+            start,
+            {
+              type: "TEXT_MESSAGE_START",
+              messageId: "answer",
+              role: "assistant",
+            },
+            {
+              type: "TEXT_MESSAGE_CONTENT",
+              messageId: "answer",
+              delta: "a1",
+            },
+            { type: "TEXT_MESSAGE_END", messageId: "answer" },
+            {
+              type: "MESSAGES_SNAPSHOT",
+              messages: [...messages, answer, other(false)],
+            },
+            {
+              type: "MESSAGES_SNAPSHOT",
+              messages: [...messages, answer, other(true)],
+            },
+            {
+              type: "STATE_SNAPSHOT",
+              snapshot: { marker: "snapshots-imported" },
+            },
+          ];
+          return new Response(
+            new ReadableStream<Uint8Array>({
+              async start(controller) {
+                controller.enqueue(
+                  new TextEncoder().encode(
+                    events
+                      .map((event) => `data: ${JSON.stringify(event)}\n\n`)
+                      .join(""),
+                  ),
+                );
+                await streamCanClose;
+                if (runFinished) {
+                  controller.enqueue(
+                    new TextEncoder().encode(
+                      `data: ${JSON.stringify({
+                        type: "RUN_FINISHED",
+                        threadId: "thread",
+                        runId: "run",
+                      })}\n\n`,
+                    ),
+                  );
+                }
+                controller.close();
+              },
+            }),
+            { headers: { "Content-Type": "text/event-stream" } },
+          );
+        },
+      });
+      const core = new AgUiThreadController({
+        agent,
+        logger: makeLogger(),
+        showThinking: false,
+      });
+      const snapshotsImported = new Promise<void>((resolve) => {
+        const unsubscribe = core.subscribe(() => {
+          if (
+            JSON.stringify(core.getState()) ===
+            JSON.stringify({ marker: "snapshots-imported" })
+          ) {
+            unsubscribe();
+            resolve();
+          }
+        });
+      });
+
+      const run = core.append(userMessage);
+      await snapshotsImported;
+      const otherDuringRun = core
+        .getMessages()
+        .find(({ id }) => id === "other");
+      const statusDuringRun = {
+        isRunning: core.isRunning(),
+        status: otherDuringRun?.status,
+      };
+
+      releaseStream();
+      await run;
+
+      expect(statusDuringRun).toMatchObject({
+        isRunning: true,
+        status: { type: "running" },
+      });
+      expect(core.isRunning()).toBe(false);
+      expect(
+        core.getMessages().find(({ id }) => id === "other")?.status,
+      ).toMatchObject(expectedStatus);
+      expect(core.getPendingToolCalls()).toEqual(
+        runFinished ? { messageId: "other", toolCallIds: ["ocall"] } : null,
+      );
+    },
+  );
 });

package/src/runtime/AgUiThreadController.replay.test.ts

@@ -3371,6 +3371,63 @@ describe("AgUiThreadController selected-turn replay", () => {
     expect(attempts).toBe(1);
   });
 
+  it("leaves an off-branch answer alone when the first message collides with its ID", async () => {
+    const core = await loaded(
+      [
+        user("u"),
+        assistant("off-branch", [{ type: "text", text: "old off branch" }]),
+        assistant("selected", [{ type: "text", text: "selected branch" }]),
+      ],
+      async (input, subscriber) => {
+        text(subscriber, "off-branch", "answer A");
+        streamToolCall(subscriber, "t", "{}");
+        text(subscriber, "fresh-b", "answer B");
+        emit(subscriber, {
+          type: "TOOL_CALL_RESULT",
+          toolCallId: "t",
+          messageId: "result-t",
+          content: "server result",
+        });
+        emit(subscriber, {
+          type: "TEXT_MESSAGE_END",
+          messageId: "off-branch",
+        });
+        emit(subscriber, { type: "RUN_FINISHED", runId: input.runId });
+      },
+      [null, "u", "u"],
+    );
+
+    await core.append({
+      role: "user",
+      content: [{ type: "text", text: "new prompt" }],
+      attachments: [],
+      createdAt: new Date(1),
+      metadata: { custom: {} },
+      parentId: core.getMessageRepository().headId ?? null,
+      sourceId: null,
+      runConfig: {},
+      startRun: true,
+    });
+
+    expect(
+      core
+        .getMessageRepository()
+        .messages.find(({ message }) => message.id === "off-branch")?.message,
+    ).toMatchObject({
+      content: [{ type: "text", text: "old off branch" }],
+      status: { type: "complete" },
+    });
+    expect(
+      core
+        .getMessageRepository()
+        .messages.find(({ message }) => message.id === "fresh-b")?.message,
+    ).toMatchObject({
+      content: expect.arrayContaining([
+        expect.objectContaining({ toolCallId: "t", result: "server result" }),
+      ]),
+    });
+  });
+
   it.each(["frontend", "server"] as const)(
     "continues a running replay only for a %s result on its retained call",
     async (source) => {

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

@@ -20,7 +20,7 @@ const userMessage = {
   startRun: true,
 } as unknown as AppendMessage;
 
-type Row = { id: string; status: string | undefined };
+type Row = { id: string; status: string | undefined; text: string };
 
 // Follows the ThreadHistoryAdapter contract: both writes upsert by message
 // id, and only append moves the stored head.
@@ -36,7 +36,13 @@ function createHistory(writes: string[]) {
       new Map<string, Row>(
         [...items.values()].map((item) => [
           item.message.id,
-          { id: item.message.id, status: item.message.status?.type },
+          {
+            id: item.message.id,
+            status: item.message.status?.type,
+            text: item.message.content
+              .flatMap((part: any) => (part.type === "text" ? [part.text] : []))
+              .join(""),
+          },
         ]),
       ),
     load: async () => ({ headId, messages: [...items.values()] }),
@@ -202,6 +208,21 @@ describe("ag-ui: a snapshot's new head and the history adapter", () => {
     expect(reloaded.slice(1), writes.join(", ")).toEqual(shown.slice(1));
   });
 
+  it("stores the new head when the moved-off answer settles after the move", async () => {
+    const { rows, writes, shown, reloaded } = await run((req) => [
+      start,
+      [moveSnapshot(req)],
+      [{ type: "TEXT_MESSAGE_CONTENT", messageId: "answer", delta: "a2" }],
+      [end, { type: "RUN_FINISHED", threadId: "t", runId: "r" }],
+    ]);
+    expect(shown.slice(1)).toEqual(["answer:complete", "other:complete"]);
+    expect(rows.get("answer"), writes.join(", ")).toMatchObject({
+      status: "complete",
+      text: "a1a2",
+    });
+    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,
@@ -215,6 +236,17 @@ describe("ag-ui: a snapshot's new head and the history adapter", () => {
     expect(reloaded.slice(1), writes.join(", ")).toEqual(shown.slice(1));
   });
 
+  it("stores the new head when the stream ends after the moved-off answer's frame closes", async () => {
+    const { rows, writes, shown, reloaded } = await run((req) => [
+      start,
+      [moveSnapshot(req)],
+      [end],
+    ]);
+    expect(shown.slice(1)).toEqual(["answer:complete", "other:incomplete"]);
+    expect(rows.get("other")?.status, writes.join(", ")).toBe("incomplete");
+    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,