@assistant-ui/ai-sdk 0.0.6 → 0.0.8

package/src/runtime/useAISDKRuntime.test.ts

@@ -3,6 +3,9 @@
 import { act, renderHook, waitFor } from "@testing-library/react";
 import { beforeEach, describe, expect, it, vi } from "vitest";
 import { validateUIMessages } from "ai";
+import { ToolResponse } from "assistant-stream";
+import type { UIMessage } from "@ai-sdk/react";
+import type { MessageFormatRepository } from "@assistant-ui/core";
 
 // Mock only the sibling module that requires AUI store context (not available
 // in isolation). Every other dependency — useExternalStoreRuntime,
@@ -15,6 +18,8 @@ vi.mock("./useExternalHistory", async (importOriginal) => {
     useExternalHistory: vi.fn(() => ({
       isLoading: false,
       deleteMessage: vi.fn().mockResolvedValue(undefined),
+      persistToolInteractions: vi.fn().mockResolvedValue(undefined),
+      persistToolApprovalResponses: vi.fn().mockResolvedValue(undefined),
     })),
   };
 });
@@ -80,6 +85,8 @@ describe("useAISDKRuntime", () => {
     vi.mocked(useExternalHistory).mockReturnValue({
       isLoading: false,
       deleteMessage: vi.fn().mockResolvedValue(undefined),
+      persistToolInteractions: vi.fn().mockResolvedValue(undefined),
+      persistToolApprovalResponses: vi.fn().mockResolvedValue(undefined),
     });
   });
 
@@ -119,13 +126,13 @@ describe("useAISDKRuntime", () => {
       result.current.thread.append({
         role: "user",
         content: [{ type: "text", text: "hello" }],
-        runConfig: { custom: { model: "gpt-5.6-luna" } },
+        runConfig: { custom: { model: "gpt-6-luna" } },
       });
     });
 
     await waitFor(() => {
       expect(chat.sendMessage).toHaveBeenCalledWith(expect.anything(), {
-        metadata: { custom: { model: "gpt-5.6-luna" } },
+        metadata: { custom: { model: "gpt-6-luna" } },
       });
     });
   });
@@ -569,7 +576,65 @@ describe("useAISDKRuntime", () => {
     ).resolves.toBeDefined();
   });
 
-  it("forwards a successful tool result through addToolOutput, not the deprecated addToolResult", async () => {
+  it("cancels a pending tool call left behind a staged message", async () => {
+    const chat = createChatHelpers([
+      {
+        id: "a1",
+        role: "assistant",
+        parts: [
+          {
+            type: "dynamic-tool",
+            toolName: "mcp_search",
+            toolCallId: "tc-1",
+            state: "approval-requested",
+            input: { q: "hi" },
+            approval: { id: "appr-1" },
+          },
+        ],
+      },
+    ]);
+
+    const { result } = renderHook(() => useAISDKRuntime(chat));
+
+    await waitFor(() => {
+      expect(result.current.thread.getState().messages.length).toBeGreaterThan(
+        0,
+      );
+    });
+
+    act(() => {
+      result.current.thread.append({
+        role: "user",
+        content: [{ type: "text", text: "context" }],
+        startRun: false,
+      });
+    });
+
+    await waitFor(() => {
+      expect(chat.messages).toHaveLength(2);
+    });
+
+    act(() => {
+      result.current.thread.append({
+        role: "user",
+        content: [{ type: "text", text: "what" }],
+      });
+    });
+
+    await waitFor(() => {
+      expect(chat.sendMessage).toHaveBeenCalledTimes(1);
+    });
+
+    const part = chat.messages[0].parts[0];
+    expect(part.state).toBe("output-error");
+    expect(part.approval).toBeUndefined();
+
+    await expect(
+      validateUIMessages({ messages: chat.messages }),
+    ).resolves.toBeDefined();
+  });
+
+  it("attaches a tool artifact to the live part and forwards its result through addToolOutput", async () => {
     const chat = createChatHelpers([
       {
         id: "a1",
@@ -585,7 +650,7 @@ describe("useAISDKRuntime", () => {
       },
     ]);
 
-    const { result } = renderHook(() => useAISDKRuntime(chat));
+    const { result, rerender } = renderHook(() => useAISDKRuntime(chat));
 
     await waitFor(() => {
       expect(result.current.thread.getState().messages.length).toBeGreaterThan(
@@ -597,13 +662,26 @@ describe("useAISDKRuntime", () => {
       result.current.thread
         .getMessageById("a1")
         .getMessagePartByToolCallId("tc-1")
-        .addToolResult({ temp: 72 });
+        .addToolResult(
+          new ToolResponse({
+            result: { temp: 72 },
+            artifact: { preview: "72°F and sunny" },
+          }),
+        );
     });
 
     await waitFor(() => {
       expect(chat.addToolOutput).toHaveBeenCalledTimes(1);
     });
 
+    const livePart = result.current.thread
+      .getMessageById("a1")
+      .getMessagePartByToolCallId("tc-1")
+      .getState();
+    expect(
+      livePart.type === "tool-call" ? livePart.artifact : undefined,
+    ).toEqual({ preview: "72°F and sunny" });
+
     expect(chat.addToolOutput).toHaveBeenCalledWith(
       expect.objectContaining({
         tool: "weather",
@@ -612,7 +690,90 @@ describe("useAISDKRuntime", () => {
         options: { metadata: undefined },
       }),
     );
+    expect(chat.addToolOutput.mock.calls[0]?.[0]).not.toHaveProperty(
+      "artifact",
+    );
     expect(chat.addToolResult).not.toHaveBeenCalled();
+    expect(
+      vi.mocked(useExternalHistory).mock.calls.at(-1)?.[5]?.get("tc-1"),
+    ).toEqual({ preview: "72°F and sunny" });
+
+    chat.messages = chat.messages.map((message: UIMessage) => ({
+      ...message,
+    }));
+    rerender();
+
+    await waitFor(() => {
+      const part = result.current.thread
+        .getMessageById("a1")
+        .getMessagePartByToolCallId("tc-1")
+        .getState();
+      expect(part.type === "tool-call" ? part.artifact : undefined).toEqual({
+        preview: "72°F and sunny",
+      });
+    });
+    expect(chat.messages[0]?.metadata).toBeUndefined();
+  });
+
+  it("shows recorded tool interactions without writing them to chat messages", async () => {
+    const chat = createChatHelpers([
+      {
+        id: "a1",
+        role: "assistant",
+        parts: [
+          {
+            type: "tool-weather",
+            toolCallId: "tc-1",
+            state: "input-available",
+            input: { city: "NYC" },
+          },
+        ],
+      },
+    ]);
+    const persistToolInteractions = vi.fn().mockResolvedValue(undefined);
+    vi.mocked(useExternalHistory).mockReturnValue({
+      isLoading: false,
+      deleteMessage: vi.fn().mockResolvedValue(undefined),
+      persistToolInteractions,
+      persistToolApprovalResponses: vi.fn().mockResolvedValue(undefined),
+    });
+
+    const { result } = renderHook(() => useAISDKRuntime(chat));
+
+    await waitFor(() => {
+      expect(result.current.thread.getState().messages).toHaveLength(1);
+    });
+
+    await act(async () => {
+      await result.current.thread
+        .getMessageById("a1")
+        .getMessagePartByToolCallId("tc-1").unstable_recordInteraction!({
+        type: "action",
+        payload: { refresh: true },
+      });
+    });
+
+    await waitFor(() => {
+      const part = result.current.thread
+        .getMessageById("a1")
+        .getMessagePartByToolCallId("tc-1")
+        .getState();
+      expect(
+        part.type === "tool-call" ? part.unstable_interactions : undefined,
+      ).toEqual({
+        entries: [
+          {
+            type: "action",
+            occurredAt: expect.any(Number),
+            payload: { refresh: true },
+          },
+        ],
+      });
+    });
+    expect(persistToolInteractions).toHaveBeenCalledExactlyOnceWith("a1");
+    expect(chat.messages[0]?.metadata).toBeUndefined();
+    expect(chat.addToolOutput).not.toHaveBeenCalled();
+    expect(chat.sendMessage).not.toHaveBeenCalled();
   });
 
   it("appends a new user message without sending when startRun is false", async () => {
@@ -702,6 +863,8 @@ describe("useAISDKRuntime", () => {
     vi.mocked(useExternalHistory).mockReturnValue({
       isLoading: false,
       deleteMessage,
+      persistToolInteractions: vi.fn().mockResolvedValue(undefined),
+      persistToolApprovalResponses: vi.fn().mockResolvedValue(undefined),
     });
     const chat = createChatHelpers([
       { id: "u1", role: "user", parts: [{ type: "text", text: "first" }] },
@@ -736,6 +899,115 @@ describe("useAISDKRuntime", () => {
     ]);
   });
 
+  it("removes tool artifacts with their deleted message", async () => {
+    const deleteMessage = vi.fn().mockResolvedValue(undefined);
+    vi.mocked(useExternalHistory).mockReturnValue({
+      isLoading: false,
+      deleteMessage,
+      persistToolInteractions: vi.fn().mockResolvedValue(undefined),
+      persistToolApprovalResponses: vi.fn().mockResolvedValue(undefined),
+    });
+    const chat = createChatHelpers([
+      {
+        id: "a1",
+        role: "assistant",
+        parts: [
+          {
+            type: "tool-weather",
+            toolCallId: "tc-1",
+            state: "input-available",
+            input: { city: "NYC" },
+          },
+        ],
+      },
+    ]);
+
+    const { result } = renderHook(() => useAISDKRuntime(chat));
+
+    await waitFor(() => {
+      expect(result.current.thread.getState().messages).toHaveLength(1);
+    });
+
+    act(() => {
+      result.current.thread
+        .getMessageById("a1")
+        .getMessagePartByToolCallId("tc-1")
+        .addToolResult(
+          new ToolResponse({
+            result: { temp: 72 },
+            artifact: { preview: "72°F and sunny" },
+          }),
+        );
+    });
+
+    await waitFor(() => {
+      expect(chat.addToolOutput).toHaveBeenCalledTimes(1);
+    });
+
+    const toolArtifacts = vi.mocked(useExternalHistory).mock.calls.at(-1)?.[5];
+    expect(toolArtifacts?.get("tc-1")).toEqual({
+      preview: "72°F and sunny",
+    });
+    await act(async () => {
+      await result.current.thread.getMessageById("a1").delete();
+    });
+
+    expect(deleteMessage).toHaveBeenCalledWith("a1");
+    expect(toolArtifacts?.has("tc-1")).toBe(false);
+  });
+
+  it("removes tool interactions with their deleted message", async () => {
+    const deleteMessage = vi.fn().mockResolvedValue(undefined);
+    vi.mocked(useExternalHistory).mockReturnValue({
+      isLoading: false,
+      deleteMessage,
+      persistToolInteractions: vi.fn().mockResolvedValue(undefined),
+      persistToolApprovalResponses: vi.fn().mockResolvedValue(undefined),
+    });
+    const chat = createChatHelpers([
+      {
+        id: "a1",
+        role: "assistant",
+        parts: [
+          {
+            type: "tool-weather",
+            toolCallId: "tc-1",
+            state: "output-available",
+            input: { city: "NYC" },
+            output: { temp: 72 },
+          },
+        ],
+      },
+    ]);
+
+    const { result } = renderHook(() => useAISDKRuntime(chat));
+
+    await waitFor(() => {
+      expect(result.current.thread.getState().messages).toHaveLength(1);
+    });
+
+    await act(async () => {
+      await result.current.thread
+        .getMessageById("a1")
+        .getMessagePartByToolCallId("tc-1").unstable_recordInteraction!({
+        type: "action",
+        payload: { type: "refresh" },
+      });
+    });
+
+    const toolInteractions = vi
+      .mocked(useExternalHistory)
+      .mock.calls.at(-1)?.[7];
+    expect(toolInteractions?.has("tc-1")).toBe(true);
+
+    await act(async () => {
+      await result.current.thread.getMessageById("a1").delete();
+    });
+
+    expect(deleteMessage).toHaveBeenCalledWith("a1");
+    expect(toolInteractions?.has("tc-1")).toBe(false);
+  });
+
   it("edit slices history to parentId and sends the edited message", async () => {
     const chat = createChatHelpers([
       { id: "u1", role: "user", parts: [{ type: "text", text: "first" }] },
@@ -981,7 +1253,7 @@ describe("useAISDKRuntime", () => {
   it("imports a message tree without replacing the chat feed", async () => {
     const chat = createChatHelpers();
     const onBranchChange = vi.fn();
-    const messageRepository = {
+    const messageRepository: MessageFormatRepository<UIMessage> = {
       headId: "a2",
       messages: [
         {
@@ -1077,7 +1349,7 @@ describe("useAISDKRuntime", () => {
     const chat = createChatHelpers([
       { id: "live", role: "user", parts: [{ type: "text", text: "keep me" }] },
     ]);
-    const messageRepository = {
+    const messageRepository: MessageFormatRepository<UIMessage> = {
       headId: "a1",
       messages: [
         {
@@ -1115,7 +1387,9 @@ describe("useAISDKRuntime", () => {
 
   it("does not reseed when the repository object identity changes", async () => {
     const chat = createChatHelpers();
-    const makeRepository = (text: string) => ({
+    const makeRepository = (
+      text: string,
+    ): MessageFormatRepository<UIMessage> => ({
       headId: "a1",
       messages: [
         {