@assistant-ui/ai-sdk 0.0.2 → 0.0.4

package/src/runtime/useAISDKRuntime.test.ts

@@ -7,13 +7,17 @@ import { validateUIMessages } from "ai";
 // Mock only the sibling module that requires AUI store context (not available
 // in isolation). Every other dependency — useExternalStoreRuntime,
 // useToolInvocations, the message converter — runs for real.
-vi.mock("./useExternalHistory", () => ({
-  useExternalHistory: vi.fn(() => ({
-    isLoading: false,
-    deleteMessage: vi.fn().mockResolvedValue(undefined),
-  })),
-  toExportedMessageRepository: vi.fn(),
-}));
+vi.mock("./useExternalHistory", async (importOriginal) => {
+  const original =
+    await importOriginal<typeof import("./useExternalHistory")>();
+  return {
+    ...original,
+    useExternalHistory: vi.fn(() => ({
+      isLoading: false,
+      deleteMessage: vi.fn().mockResolvedValue(undefined),
+    })),
+  };
+});
 
 import { useExternalHistory } from "./useExternalHistory";
 import { useAISDKRuntime } from "./useAISDKRuntime";
@@ -23,6 +27,7 @@ const createChatHelpers = (messages: any[] = []) => {
   let currentMessages = [...messages];
 
   const chatHelpers: any = {
+    id: "chat-1",
     status: "ready",
     error: null,
     messages: currentMessages,
@@ -131,10 +136,11 @@ describe("useAISDKRuntime", () => {
     abortError.name = "AbortError";
     const chat = createChatHelpers();
     let stopCalls = 0;
+    let rejectStop!: (error: unknown) => void;
     chat.stop = () => {
       stopCalls += 1;
       return new Promise((_, reject) => {
-        setTimeout(() => reject(abortError), 5);
+        rejectStop = reject;
       });
     };
     const consoleError = vi
@@ -143,8 +149,11 @@ describe("useAISDKRuntime", () => {
 
     try {
       const { result } = renderHook(() => useAISDKRuntime(chat));
-      const unhandledRejections = await captureUnhandledRejections(() => {
-        result.current.thread.cancelRun();
+      const unhandledRejections = await captureUnhandledRejections(async () => {
+        await act(async () => {
+          result.current.thread.cancelRun();
+          rejectStop(abortError);
+        });
       });
 
       expect(stopCalls).toBe(1);
@@ -155,6 +164,286 @@ describe("useAISDKRuntime", () => {
     }
   });
 
+  it("marks only the stopped output cancelled", async () => {
+    let resolveStop!: () => void;
+    const chat = createChatHelpers([
+      {
+        id: "assistant-1",
+        role: "assistant",
+        parts: [{ type: "text", text: "partial", state: "streaming" }],
+      },
+    ]);
+    chat.status = "streaming";
+    chat.stop = vi.fn(
+      () =>
+        new Promise<void>((resolve) => {
+          resolveStop = resolve;
+        }),
+    );
+
+    const { result, rerender } = renderHook(() => useAISDKRuntime(chat));
+
+    act(() => {
+      result.current.thread.cancelRun();
+    });
+    rerender();
+
+    expect(
+      result.current.thread.getState().messages.at(-1)?.status,
+    ).toMatchObject({ type: "running" });
+
+    act(() => {
+      chat.status = "ready";
+      rerender();
+    });
+
+    await waitFor(() => {
+      expect(
+        result.current.thread.getState().messages.at(-1)?.status,
+      ).toMatchObject({
+        type: "incomplete",
+        reason: "cancelled",
+      });
+    });
+
+    await act(async () => {
+      resolveStop();
+      await Promise.resolve();
+    });
+
+    act(() => {
+      chat.setMessages([
+        {
+          id: "assistant-2",
+          role: "assistant",
+          parts: [{ type: "text", text: "replacement" }],
+        },
+      ]);
+      rerender();
+    });
+    await waitFor(() => {
+      expect(
+        result.current.thread.getState().messages.at(-1)?.status,
+      ).toMatchObject({ type: "complete", reason: "unknown" });
+    });
+
+    act(() => {
+      chat.setMessages([
+        {
+          id: "assistant-1",
+          role: "assistant",
+          parts: [{ type: "text", text: "partial" }],
+        },
+      ]);
+      chat.status = "streaming";
+      rerender();
+    });
+
+    expect(
+      result.current.thread.getState().messages.at(-1)?.status,
+    ).toMatchObject({ type: "running" });
+
+    act(() => {
+      chat.status = "ready";
+      rerender();
+    });
+
+    await waitFor(() => {
+      expect(
+        result.current.thread.getState().messages.at(-1)?.status,
+      ).toMatchObject({
+        type: "complete",
+        reason: "unknown",
+      });
+    });
+  });
+
+  it("keeps the stopped output cancelled through the next turn", async () => {
+    const chat = createChatHelpers([
+      { id: "u1", role: "user", parts: [{ type: "text", text: "hi" }] },
+      {
+        id: "assistant-1",
+        role: "assistant",
+        parts: [{ type: "text", text: "partial", state: "streaming" }],
+      },
+    ]);
+    chat.status = "streaming";
+
+    const { result, rerender } = renderHook(() => useAISDKRuntime(chat));
+
+    await act(async () => {
+      result.current.thread.cancelRun();
+    });
+    act(() => {
+      chat.status = "ready";
+      rerender();
+    });
+
+    await waitFor(() => {
+      expect(
+        result.current.thread.getState().messages.at(-1)?.status,
+      ).toMatchObject({ type: "incomplete", reason: "cancelled" });
+    });
+
+    act(() => {
+      chat.setMessages([
+        ...chat.messages,
+        { id: "u2", role: "user", parts: [{ type: "text", text: "next" }] },
+        {
+          id: "assistant-2",
+          role: "assistant",
+          parts: [{ type: "text", text: "answer" }],
+        },
+      ]);
+      rerender();
+    });
+
+    await waitFor(() => {
+      const messages = result.current.thread.getState().messages;
+      expect(
+        messages.find((message) => message.id === "assistant-1")?.status,
+      ).toMatchObject({ type: "incomplete", reason: "cancelled" });
+      expect(messages.at(-1)?.status).toMatchObject({
+        type: "complete",
+        reason: "unknown",
+      });
+    });
+  });
+
+  it("retracts the cancellation when the provider picks the message back up", async () => {
+    const chat = createChatHelpers([
+      {
+        id: "assistant-1",
+        role: "assistant",
+        parts: [{ type: "text", text: "partial", state: "streaming" }],
+      },
+    ]);
+    chat.status = "streaming";
+
+    const { result, rerender } = renderHook(() => useAISDKRuntime(chat));
+
+    await act(async () => {
+      result.current.thread.cancelRun();
+    });
+    act(() => {
+      chat.status = "ready";
+      rerender();
+    });
+
+    await waitFor(() => {
+      expect(
+        result.current.thread.getState().messages.at(-1)?.status,
+      ).toMatchObject({ type: "incomplete", reason: "cancelled" });
+    });
+
+    act(() => {
+      chat.status = "streaming";
+      rerender();
+    });
+    act(() => {
+      chat.status = "ready";
+      rerender();
+    });
+
+    await waitFor(() => {
+      expect(
+        result.current.thread.getState().messages.at(-1)?.status,
+      ).toMatchObject({ type: "complete", reason: "unknown" });
+    });
+  });
+
+  it("does not mark completed output cancelled when already idle", async () => {
+    const chat = createChatHelpers([
+      {
+        id: "assistant-1",
+        role: "assistant",
+        parts: [{ type: "text", text: "finished" }],
+      },
+    ]);
+
+    const { result, rerender } = renderHook(() => useAISDKRuntime(chat));
+
+    act(() => {
+      result.current.thread.cancelRun();
+      rerender();
+    });
+
+    await waitFor(() => {
+      expect(
+        result.current.thread.getState().messages.at(-1)?.status,
+      ).toMatchObject({ type: "complete", reason: "unknown" });
+    });
+  });
+
+  it("marks output cancelled while a client tool is still executing", async () => {
+    let resolveTool!: (value: string) => void;
+    const execute = vi.fn(
+      () =>
+        new Promise<string>((resolve) => {
+          resolveTool = resolve;
+        }),
+    );
+    const chat = createChatHelpers();
+
+    const { result, rerender } = renderHook(() => useAISDKRuntime(chat));
+    const unregister = result.current.registerModelContextProvider({
+      getModelContext: () => ({
+        tools: {
+          weather: {
+            parameters: { type: "object", properties: {} },
+            execute,
+          },
+        },
+      }),
+    });
+
+    try {
+      act(() => {
+        chat.setMessages([
+          {
+            id: "assistant-1",
+            role: "assistant",
+            parts: [
+              {
+                type: "tool-weather",
+                toolCallId: "tool-1",
+                state: "input-available",
+                input: { city: "London" },
+              },
+            ],
+          },
+        ]);
+        rerender();
+      });
+
+      await waitFor(() => {
+        expect(execute).toHaveBeenCalledOnce();
+        expect(result.current.thread.getState().isRunning).toBe(true);
+      });
+
+      act(() => {
+        result.current.thread.cancelRun();
+        chat.setMessages([
+          {
+            id: "assistant-1",
+            role: "assistant",
+            parts: [{ type: "text", text: "stopped" }],
+          },
+        ]);
+        resolveTool("sunny");
+        rerender();
+      });
+
+      await waitFor(() => {
+        expect(
+          result.current.thread.getState().messages.at(-1)?.status,
+        ).toMatchObject({ type: "incomplete", reason: "cancelled" });
+      });
+    } finally {
+      unregister();
+    }
+  });
+
   it("reports non-AbortError cancellation failures", async () => {
     const stopError = new Error("stop failed");
     const chat = createChatHelpers();
@@ -690,6 +979,186 @@ describe("useAISDKRuntime", () => {
     expect(result.current.thread.getState().suggestions).toEqual(suggestions);
   });
 
+  it("imports a message tree without replacing the chat feed", async () => {
+    const chat = createChatHelpers();
+    const onBranchChange = vi.fn();
+    const messageRepository = {
+      headId: "a2",
+      messages: [
+        {
+          parentId: null,
+          message: {
+            id: "u1",
+            role: "user" as const,
+            parts: [{ type: "text" as const, text: "question" }],
+          },
+        },
+        {
+          parentId: "u1",
+          message: {
+            id: "a1",
+            role: "assistant" as const,
+            parts: [{ type: "text" as const, text: "first" }],
+          },
+        },
+        {
+          parentId: "u1",
+          message: {
+            id: "a2",
+            role: "assistant" as const,
+            parts: [{ type: "text" as const, text: "second" }],
+          },
+        },
+      ],
+    };
+
+    const { result, rerender } = renderHook(() =>
+      useAISDKRuntime(chat, {
+        messageRepository,
+        unstable_onBranchChange: onBranchChange,
+      }),
+    );
+
+    await waitFor(() => {
+      expect(
+        result.current.thread.getState().messages.map((message) => message.id),
+      ).toEqual(["u1", "a2"]);
+    });
+    expect(result.current.thread.getMessageById("a2").getState()).toMatchObject(
+      {
+        branchNumber: 2,
+        branchCount: 2,
+      },
+    );
+    expect(chat.messages.map((message: { id: string }) => message.id)).toEqual([
+      "u1",
+      "a2",
+    ]);
+
+    act(() => {
+      result.current.thread
+        .getMessageById("a2")
+        .switchToBranch({ branchId: "a1" });
+    });
+
+    expect(onBranchChange).toHaveBeenCalledWith({
+      headId: "a1",
+      visibleMessageIds: ["u1", "a1"],
+    });
+    expect(chat.messages.map((message: { id: string }) => message.id)).toEqual([
+      "u1",
+      "a1",
+    ]);
+    expect(chat.messages).not.toEqual([]);
+
+    chat.messages = [
+      { id: "u1", role: "user", parts: [{ type: "text", text: "question" }] },
+      {
+        id: "a1",
+        role: "assistant",
+        parts: [{ type: "text", text: "first token" }],
+      },
+    ];
+    rerender();
+
+    await waitFor(() => {
+      expect(textOf(result.current.thread.getState().messages.at(-1))).toBe(
+        "first token",
+      );
+    });
+    expect(result.current.thread.getMessageById("a1").getState()).toMatchObject(
+      {
+        branchNumber: 1,
+        branchCount: 2,
+      },
+    );
+  });
+
+  it("does not overwrite a nonempty chat with the repository", async () => {
+    const chat = createChatHelpers([
+      { id: "live", role: "user", parts: [{ type: "text", text: "keep me" }] },
+    ]);
+    const messageRepository = {
+      headId: "a1",
+      messages: [
+        {
+          parentId: null,
+          message: {
+            id: "u1",
+            role: "user" as const,
+            parts: [{ type: "text" as const, text: "question" }],
+          },
+        },
+        {
+          parentId: "u1",
+          message: {
+            id: "a1",
+            role: "assistant" as const,
+            parts: [{ type: "text" as const, text: "first" }],
+          },
+        },
+      ],
+    };
+
+    const { result } = renderHook(() =>
+      useAISDKRuntime(chat, { messageRepository }),
+    );
+
+    await waitFor(() => {
+      expect(
+        result.current.thread.getState().messages.map((message) => message.id),
+      ).toEqual(["live"]);
+    });
+    expect(chat.messages.map((message: { id: string }) => message.id)).toEqual([
+      "live",
+    ]);
+  });
+
+  it("does not reseed when the repository object identity changes", async () => {
+    const chat = createChatHelpers();
+    const makeRepository = (text: string) => ({
+      headId: "a1",
+      messages: [
+        {
+          parentId: null,
+          message: {
+            id: "u1",
+            role: "user" as const,
+            parts: [{ type: "text" as const, text: "question" }],
+          },
+        },
+        {
+          parentId: "u1",
+          message: {
+            id: "a1",
+            role: "assistant" as const,
+            parts: [{ type: "text" as const, text }],
+          },
+        },
+      ],
+    });
+
+    const { result, rerender } = renderHook(
+      ({ repository }) =>
+        useAISDKRuntime(chat, { messageRepository: repository }),
+      { initialProps: { repository: makeRepository("first") } },
+    );
+
+    await waitFor(() => {
+      expect(textOf(result.current.thread.getState().messages.at(-1))).toBe(
+        "first",
+      );
+    });
+    const setMessagesCalls = chat.setMessages.mock.calls.length;
+
+    rerender({ repository: makeRepository("second") });
+
+    expect(textOf(result.current.thread.getState().messages.at(-1))).toBe(
+      "first",
+    );
+    expect(chat.setMessages.mock.calls.length).toBe(setMessagesCalls);
+  });
+
   it("calls adapters.suggestion after settle with messages and signal", async () => {
     const generate = vi.fn().mockResolvedValue([{ prompt: "next" }]);
     const chat = createChatHelpers([