@assistant-ui/ai-sdk 0.0.7 → 0.0.8

package/src/runtime/useAISDKRuntime.approval.test.tsx

@@ -1,11 +1,18 @@
 // @vitest-environment jsdom
 
-import { act, renderHook } from "@testing-library/react";
-import type { ExternalStoreAdapter } from "@assistant-ui/core";
-import { describe, expect, it, vi } from "vitest";
+import { act, renderHook, waitFor } from "@testing-library/react";
+import {
+  bindExternalStoreMessage,
+  type ExternalStoreAdapter,
+  type ThreadAssistantMessage,
+  type ThreadMessage,
+} from "@assistant-ui/core";
+import { beforeEach, describe, expect, it, vi } from "vitest";
 
 const mocks = vi.hoisted(() => ({
   adapter: undefined as ExternalStoreAdapter | undefined,
+  persistToolApprovalResponses: vi.fn(),
+  threadMessages: [] as ThreadMessage[],
 }));
 
 vi.mock("@assistant-ui/core/react", async (importOriginal) => {
@@ -15,7 +22,11 @@ vi.mock("@assistant-ui/core/react", async (importOriginal) => {
     ...original,
     useExternalStoreRuntime: vi.fn((adapter: ExternalStoreAdapter) => {
       mocks.adapter = adapter;
-      return {};
+      return {
+        thread: {
+          getState: () => ({ messages: mocks.threadMessages }),
+        },
+      } as never;
     }),
     useRuntimeAdapters: vi.fn(() => ({})),
   };
@@ -29,13 +40,30 @@ vi.mock("./useExternalHistory", async (importOriginal) => {
     useExternalHistory: vi.fn(() => ({
       isLoading: false,
       deleteMessage: vi.fn().mockResolvedValue(undefined),
+      persistToolInteractions: vi.fn().mockResolvedValue(undefined),
+      persistToolApprovalResponses: mocks.persistToolApprovalResponses,
     })),
   };
 });
 
 import { useAISDKRuntime } from "./useAISDKRuntime";
+import { useExternalHistory } from "./useExternalHistory";
 
 describe("useAISDKRuntime tool approvals", () => {
+  beforeEach(() => {
+    mocks.adapter = undefined;
+    mocks.threadMessages = [];
+    mocks.persistToolApprovalResponses.mockReset().mockResolvedValue(undefined);
+    vi.mocked(useExternalHistory)
+      .mockReset()
+      .mockImplementation(() => ({
+        isLoading: false,
+        deleteMessage: vi.fn().mockResolvedValue(undefined),
+        persistToolInteractions: vi.fn().mockResolvedValue(undefined),
+        persistToolApprovalResponses: mocks.persistToolApprovalResponses,
+      }));
+  });
+
   it("forwards the AI SDK approval promise to the external-store adapter", () => {
     const approvalPromise = Promise.resolve();
     const addToolApprovalResponse = vi.fn(() => approvalPromise);
@@ -88,6 +116,7 @@ describe("useAISDKRuntime tool approvals", () => {
       },
     ];
     const setMessages = vi.fn();
+    const sendMessage = vi.fn();
     const addToolApprovalResponse = vi.fn();
     const chat = {
       id: "chat-1",
@@ -95,7 +124,7 @@ describe("useAISDKRuntime tool approvals", () => {
       error: undefined,
       messages,
       setMessages,
-      sendMessage: vi.fn(),
+      sendMessage,
       regenerate: vi.fn(),
       addToolOutput: vi.fn(),
       addToolApprovalResponse,
@@ -116,17 +145,29 @@ describe("useAISDKRuntime tool approvals", () => {
       }) => mocks.adapter?.onRespondToToolApproval?.(response),
       setMessages,
       addToolApprovalResponse,
+      sendMessage,
+      messages,
       getApproval: () =>
         mocks.adapter?.messages?.[0]?.content.find(
           (part) => part.type === "tool-call",
         )?.approval,
+      getToolCall: () =>
+        mocks.adapter?.messages?.[0]?.content.find(
+          (part) => part.type === "tool-call",
+        ),
     };
   };
 
-  it("hands the complete response to a custom handler and applies the answer", async () => {
+  it("stores a host answer after the handler resolves without starting a run", async () => {
     const onRespondToToolApproval = vi.fn(async () => {});
-    const { respond, setMessages, addToolApprovalResponse, getApproval } =
-      setupPendingApproval(onRespondToToolApproval);
+    const {
+      respond,
+      setMessages,
+      addToolApprovalResponse,
+      sendMessage,
+      messages,
+      getApproval,
+    } = setupPendingApproval(onRespondToToolApproval);
 
     const response = {
       approvalId: "approval-1",
@@ -146,6 +187,13 @@ describe("useAISDKRuntime tool approvals", () => {
     });
     expect(addToolApprovalResponse).not.toHaveBeenCalled();
     expect(setMessages).not.toHaveBeenCalled();
+    expect(sendMessage).not.toHaveBeenCalled();
+    expect(mocks.persistToolApprovalResponses).toHaveBeenCalledExactlyOnceWith(
+      "message-1",
+    );
+    expect(messages[0]).not.toHaveProperty(
+      "metadata.__aui_toolApprovalResponses",
+    );
     expect(getApproval()).toEqual({
       id: "approval-1",
       approved: true,
@@ -155,7 +203,7 @@ describe("useAISDKRuntime tool approvals", () => {
     });
   });
 
-  it("sends a request the handler hands back through the AI SDK", async () => {
+  it("does not store a request the handler hands back through the AI SDK", async () => {
     const { respond, addToolApprovalResponse } = setupPendingApproval(
       (_response, { respondViaAISDK }) => respondViaAISDK(),
     );
@@ -175,6 +223,27 @@ describe("useAISDKRuntime tool approvals", () => {
       reason: "Not now",
       options: { metadata: undefined },
     });
+    expect(mocks.persistToolApprovalResponses).not.toHaveBeenCalled();
+  });
+
+  it("does not store a host answer when the handler rejects", async () => {
+    const { respond, messages, getApproval } = setupPendingApproval(
+      async () => {
+        throw new Error("resume failed");
+      },
+    );
+
+    await expect(
+      act(async () => {
+        await respond({ approvalId: "approval-1", approved: true });
+      }),
+    ).rejects.toThrow("resume failed");
+
+    expect(mocks.persistToolApprovalResponses).not.toHaveBeenCalled();
+    expect(getApproval()).toEqual({ id: "approval-1" });
+    expect(messages[0]).not.toHaveProperty(
+      "metadata.__aui_toolApprovalResponses",
+    );
   });
 
   it("reopens a request when a handed-back AI SDK response fails inside the handler", async () => {
@@ -254,6 +323,153 @@ describe("useAISDKRuntime tool approvals", () => {
     expect(onRespondToToolApproval).toHaveBeenCalledTimes(2);
   });
 
+  it("restores a host answer and refuses a second response", async () => {
+    const onRespondToToolApproval = vi.fn(async () => {});
+    const { respond, getApproval, messages } = setupPendingApproval(
+      onRespondToToolApproval,
+    );
+    const historyCall = vi.mocked(useExternalHistory).mock.calls.at(-1)!;
+    const toolApprovalResponses = historyCall[9] as Map<
+      string,
+      { approvalId: string; approved: boolean; reason?: string }
+    >;
+    const onToolApprovalResponsesRestored = historyCall[10] as () => void;
+
+    await act(async () => {
+      toolApprovalResponses.set("approval-1", {
+        approvalId: "approval-1",
+        approved: true,
+        reason: "Approved by operator",
+      });
+      onToolApprovalResponsesRestored();
+    });
+
+    await waitFor(() =>
+      expect(getApproval()).toEqual({
+        id: "approval-1",
+        approved: true,
+        reason: "Approved by operator",
+      }),
+    );
+    await expect(
+      respond({ approvalId: "approval-1", approved: true }),
+    ).rejects.toThrow(
+      "Tool approval approval-1 is not waiting for a response.",
+    );
+    expect(onRespondToToolApproval).not.toHaveBeenCalled();
+    expect(messages[0]).not.toHaveProperty(
+      "metadata.__aui_toolApprovalResponses",
+    );
+  });
+
+  it("reopens a restored approval when history clears its response map", async () => {
+    const onRespondToToolApproval = vi.fn(async () => {});
+    const { respond, getApproval } = setupPendingApproval(
+      onRespondToToolApproval,
+    );
+    const historyCall = vi.mocked(useExternalHistory).mock.calls.at(-1)!;
+    const toolApprovalResponses = historyCall[9] as Map<
+      string,
+      { approvalId: string; approved: boolean }
+    >;
+    const onToolApprovalResponsesRestored = historyCall[10] as () => void;
+
+    await act(async () => {
+      toolApprovalResponses.set("approval-1", {
+        approvalId: "approval-1",
+        approved: true,
+      });
+      onToolApprovalResponsesRestored();
+    });
+    await waitFor(() =>
+      expect(getApproval()).toEqual({ id: "approval-1", approved: true }),
+    );
+
+    await act(async () => {
+      toolApprovalResponses.clear();
+      onToolApprovalResponsesRestored();
+    });
+    await waitFor(() => expect(getApproval()).toEqual({ id: "approval-1" }));
+
+    await act(async () => {
+      await respond({ approvalId: "approval-1", approved: true });
+    });
+    expect(onRespondToToolApproval).toHaveBeenCalledOnce();
+  });
+
+  it("clears deleted tool sidecars and its host approval reservation", async () => {
+    const onRespondToToolApproval = vi.fn(async () => {});
+    const { getApproval, getToolCall, messages, respond } =
+      setupPendingApproval(onRespondToToolApproval);
+    const threadMessage: ThreadAssistantMessage = {
+      id: "message-1",
+      role: "assistant",
+      content: [
+        {
+          type: "tool-call",
+          toolCallId: "tool-1",
+          toolName: "deploy",
+          args: {},
+          argsText: "{}",
+          result: undefined,
+          isError: false,
+          approval: { id: "approval-1" },
+        },
+      ],
+      createdAt: new Date(),
+      status: { type: "requires-action", reason: "tool-calls" },
+      metadata: {
+        unstable_state: null,
+        unstable_annotations: [],
+        unstable_data: [],
+        steps: [],
+        custom: {},
+      },
+    };
+    bindExternalStoreMessage(threadMessage, messages[0]!);
+    mocks.threadMessages = [threadMessage];
+
+    await act(async () => {
+      await respond({ approvalId: "approval-1", approved: true });
+      await mocks.adapter?.onAddToolResult?.({
+        messageId: "message-1",
+        toolCallId: "tool-1",
+        toolName: "deploy",
+        result: "deployed",
+        artifact: { preview: "deployment complete" },
+        isError: false,
+      });
+      await mocks.adapter?.unstable_onRecordToolInteraction?.({
+        messageId: "message-1",
+        toolCallId: "tool-1",
+        interaction: {
+          type: "action",
+          occurredAt: 1,
+          payload: { copied: true },
+        },
+      });
+    });
+    expect(getToolCall()).toMatchObject({
+      artifact: { preview: "deployment complete" },
+      unstable_interactions: {
+        entries: [{ type: "action", occurredAt: 1, payload: { copied: true } }],
+      },
+      approval: { id: "approval-1", approved: true },
+    });
+
+    await act(async () => {
+      await mocks.adapter?.onDelete?.("message-1");
+    });
+    expect(getToolCall()).not.toHaveProperty("artifact");
+    expect(getToolCall()).not.toHaveProperty("unstable_interactions");
+    expect(getApproval()).toEqual({ id: "approval-1" });
+
+    await act(async () => {
+      await respond({ approvalId: "approval-1", approved: true });
+    });
+    expect(onRespondToToolApproval).toHaveBeenCalledTimes(2);
+  });
+
   it("rejects an approval that is not waiting for a response", async () => {
     const onRespondToToolApproval = vi.fn();
     const { respond } = setupPendingApproval(onRespondToToolApproval);

package/src/runtime/useAISDKRuntime.test.ts

@@ -3,6 +3,7 @@
 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";
 
@@ -17,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),
     })),
   };
 });
@@ -82,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),
     });
   });
 
@@ -121,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" } },
       });
     });
   });
@@ -571,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",
@@ -587,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(
@@ -599,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",
@@ -614,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 () => {
@@ -704,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" }] },
@@ -738,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" }] },