@assistant-ui/ai-sdk 0.0.11 → 0.0.12

package/src/runtime/useAISDKRuntime.test.ts

@@ -5,7 +5,10 @@ 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";
+import {
+  createSuggestionAdapter,
+  type MessageFormatRepository,
+} from "@assistant-ui/core";
 
 // Mock only the sibling module that requires AUI store context (not available
 // in isolation). Every other dependency — useExternalStoreRuntime,
@@ -381,6 +384,31 @@ describe("useAISDKRuntime", () => {
     });
   });
 
+  it("keeps rendering through a run when the last assistant message has malformed parts", () => {
+    const chat = createChatHelpers([
+      { id: "user-1", role: "user", parts: [{ type: "text", text: "hi" }] },
+      {
+        id: "assistant-1",
+        role: "assistant",
+        parts: [
+          null,
+          { text: "no type" },
+          { type: "text" },
+          { type: "text", text: "yo" },
+        ],
+      },
+    ]);
+    chat.status = "submitted";
+
+    const { result, rerender } = renderHook(() => useAISDKRuntime(chat));
+    act(() => {
+      chat.status = "ready";
+      rerender();
+    });
+
+    expect(textOf(result.current.thread.getState().messages.at(-1))).toBe("yo");
+  });
+
   it("marks output cancelled while a client tool is still executing", async () => {
     let resolveTool!: (value: string) => void;
     const execute = vi.fn(
@@ -528,6 +556,86 @@ describe("useAISDKRuntime", () => {
     expect(chat.messages[0].parts[1].state).toBe("output-available");
   });
 
+  it("sends past an assistant message with malformed parts", async () => {
+    const chat = createChatHelpers([
+      {
+        id: "a1",
+        role: "assistant",
+        parts: [
+          null,
+          { text: "no type" },
+          {
+            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.length).toBeGreaterThan(
+        0,
+      );
+    });
+
+    await act(async () => {
+      result.current.thread.append({
+        role: "user",
+        content: [{ type: "text", text: "continue" }],
+      });
+    });
+
+    await waitFor(() => {
+      expect(chat.sendMessage).toHaveBeenCalledTimes(1);
+    });
+    expect(chat.messages[0].parts.slice(0, 2)).toEqual([
+      null,
+      { text: "no type" },
+    ]);
+    expect(chat.messages[0].parts[2].state).toBe("output-error");
+  });
+
+  it("answers a tool call in an earlier assistant message with malformed parts", async () => {
+    const chat = createChatHelpers([
+      {
+        id: "a1",
+        role: "assistant",
+        parts: [
+          null,
+          {
+            type: "tool-weather",
+            toolCallId: "tc-1",
+            state: "input-available",
+            input: { city: "NYC" },
+          },
+        ],
+      },
+      { id: "u2", role: "user", parts: [{ type: "text", text: "and?" }] },
+      { id: "a2", role: "assistant", parts: [{ type: "text", text: "ok" }] },
+    ]);
+
+    const { result } = renderHook(() => useAISDKRuntime(chat));
+    await waitFor(() => {
+      expect(result.current.thread.getState().messages.length).toBe(3);
+    });
+
+    await act(async () => {
+      result.current.thread
+        .getMessageById("a1")
+        .getMessagePartByToolCallId("tc-1")
+        .addToolResult({ temp: 72 });
+    });
+
+    await waitFor(() => {
+      expect(chat.messages[0].parts[1].state).toBe("output-available");
+    });
+    expect(chat.messages[0].parts[0]).toBeNull();
+    expect(chat.messages[0].parts[1].output).toEqual({ temp: 72 });
+  });
+
   it("strips stale approval when cancelling a tool pending approval so history stays valid", async () => {
     const chat = createChatHelpers([
       {
@@ -845,10 +953,9 @@ describe("useAISDKRuntime", () => {
     expect(chat.sendMessage).not.toHaveBeenCalled();
     expect(chat.messages.map((m: any) => m.id)).toEqual([
       "u1",
-      "a1",
       expect.any(String),
     ]);
-    expect(chat.messages[2]).toEqual(
+    expect(chat.messages[1]).toEqual(
       expect.objectContaining({
         role: "user",
         parts: expect.arrayContaining([
@@ -1044,8 +1151,7 @@ describe("useAISDKRuntime", () => {
       expect(chat.sendMessage).toHaveBeenCalledTimes(1);
     });
 
-    // sliceMessagesUntil("u1") keeps u1 + following assistant messages (a1)
-    expect(chat.messages.map((m: any) => m.id)).toEqual(["u1", "a1"]);
+    expect(chat.messages.map((m: any) => m.id)).toEqual(["u1"]);
     expect(chat.sendMessage).toHaveBeenCalledWith(
       expect.objectContaining({ role: "user" }),
       { metadata: { custom: { temperature: 0.2 } } },
@@ -1697,6 +1803,307 @@ describe("useAISDKRuntime", () => {
     expect(result.current.thread.getState().suggestions).toEqual([]);
   });
 
+  it("regenerates suggestions when the same adapter is restored", async () => {
+    const generate = vi
+      .fn()
+      .mockResolvedValueOnce([{ prompt: "first" }])
+      .mockResolvedValueOnce([{ prompt: "second" }]);
+    const chat = createChatHelpers([
+      { id: "u1", role: "user", parts: [{ type: "text", text: "hi" }] },
+      {
+        id: "a1",
+        role: "assistant",
+        parts: [{ type: "text", text: "hello" }],
+      },
+    ]);
+
+    const { result, rerender } = renderHook(
+      ({ status, withAdapter }) => {
+        chat.status = status;
+        return useAISDKRuntime(
+          chat,
+          withAdapter ? { adapters: { suggestion: { generate } } } : {},
+        );
+      },
+      {
+        initialProps: {
+          status: "submitted" as string,
+          withAdapter: true,
+        },
+      },
+    );
+
+    rerender({ status: "ready", withAdapter: true });
+    await waitFor(() => {
+      expect(result.current.thread.getState().suggestions).toEqual([
+        { prompt: "first" },
+      ]);
+    });
+
+    rerender({ status: "ready", withAdapter: false });
+    expect(result.current.thread.getState().suggestions).toEqual([]);
+
+    rerender({ status: "ready", withAdapter: true });
+    await waitFor(() => expect(generate).toHaveBeenCalledTimes(2));
+    await waitFor(() => {
+      expect(result.current.thread.getState().suggestions).toEqual([
+        { prompt: "second" },
+      ]);
+    });
+  });
+
+  it("preserves the suggestion adapter receiver", async () => {
+    const adapter = {
+      prompt: "bound",
+      async generate() {
+        return [{ prompt: this.prompt }];
+      },
+    };
+    const chat = createChatHelpers([
+      { id: "u1", role: "user", parts: [{ type: "text", text: "hi" }] },
+      {
+        id: "a1",
+        role: "assistant",
+        parts: [{ type: "text", text: "hello" }],
+      },
+    ]);
+
+    const { result, rerender } = renderHook(
+      ({ status }) => {
+        chat.status = status;
+        return useAISDKRuntime(chat, {
+          adapters: { suggestion: adapter },
+        });
+      },
+      { initialProps: { status: "submitted" as string } },
+    );
+
+    rerender({ status: "ready" });
+    await waitFor(() => {
+      expect(result.current.thread.getState().suggestions).toEqual([
+        { prompt: "bound" },
+      ]);
+    });
+  });
+
+  it("keeps a keyless inline suggestion adapter stable across rerenders", async () => {
+    let resolveGenerate!: (value: readonly { prompt: string }[]) => void;
+    const generate = vi.fn(
+      ({ signal: _signal }: { signal?: AbortSignal }) =>
+        new Promise<readonly { prompt: string }[]>((resolve) => {
+          resolveGenerate = resolve;
+        }),
+    );
+    const chat = createChatHelpers([
+      { id: "u1", role: "user", parts: [{ type: "text", text: "hi" }] },
+      {
+        id: "a1",
+        role: "assistant",
+        parts: [{ type: "text", text: "hello" }],
+      },
+    ]);
+
+    const { result, rerender } = renderHook(
+      ({ status }) => {
+        chat.status = status;
+        return useAISDKRuntime(chat, {
+          adapters: {
+            suggestion: {
+              async generate({ signal }) {
+                return generate({
+                  ...(signal !== undefined && { signal }),
+                });
+              },
+            },
+          },
+        });
+      },
+      { initialProps: { status: "submitted" as string } },
+    );
+
+    rerender({ status: "ready" });
+    await waitFor(() => expect(generate).toHaveBeenCalledOnce());
+    const signal = generate.mock.calls[0]![0].signal;
+
+    rerender({ status: "ready" });
+    expect(generate).toHaveBeenCalledOnce();
+    expect(signal?.aborted).toBe(false);
+
+    resolveGenerate([{ prompt: "next" }]);
+    await waitFor(() => {
+      expect(result.current.thread.getState().suggestions).toEqual([
+        { prompt: "next" },
+      ]);
+    });
+    await act(async () => {
+      await new Promise((resolve) => setTimeout(resolve, 20));
+    });
+    expect(generate).toHaveBeenCalledOnce();
+  });
+
+  it("does not generate for an inline adapter on mount or parent rerender", async () => {
+    const complete = vi.fn().mockResolvedValue(["unexpected"]);
+    const chat = createChatHelpers([
+      { id: "u1", role: "user", parts: [{ type: "text", text: "hi" }] },
+      {
+        id: "a1",
+        role: "assistant",
+        parts: [{ type: "text", text: "hello" }],
+      },
+    ]);
+
+    const { rerender } = renderHook(() =>
+      useAISDKRuntime(chat, {
+        adapters: {
+          suggestion: createSuggestionAdapter({ complete }),
+        },
+      }),
+    );
+
+    rerender();
+    await act(async () => {
+      await new Promise((resolve) => setTimeout(resolve, 20));
+    });
+    expect(complete).not.toHaveBeenCalled();
+  });
+
+  it("restarts an inline suggestion adapter when its key changes", async () => {
+    const firstComplete = vi.fn().mockResolvedValue(["first"]);
+    const secondComplete = vi.fn().mockResolvedValue(["second"]);
+    const chat = createChatHelpers([
+      { id: "u1", role: "user", parts: [{ type: "text", text: "hi" }] },
+      {
+        id: "a1",
+        role: "assistant",
+        parts: [{ type: "text", text: "hello" }],
+      },
+    ]);
+
+    const { result, rerender } = renderHook(
+      ({ status, adapterKey, complete }) => {
+        chat.status = status;
+        return useAISDKRuntime(chat, {
+          adapters: {
+            suggestion: createSuggestionAdapter({
+              key: adapterKey,
+              complete,
+            }),
+          },
+        });
+      },
+      {
+        initialProps: {
+          status: "submitted" as string,
+          adapterKey: "first",
+          complete: firstComplete,
+        },
+      },
+    );
+
+    rerender({
+      status: "ready",
+      adapterKey: "first",
+      complete: firstComplete,
+    });
+    await waitFor(() => {
+      expect(result.current.thread.getState().suggestions).toEqual([
+        { prompt: "first" },
+      ]);
+    });
+
+    rerender({
+      status: "ready",
+      adapterKey: "first",
+      complete: firstComplete,
+    });
+    expect(firstComplete).toHaveBeenCalledOnce();
+
+    rerender({
+      status: "ready",
+      adapterKey: "second",
+      complete: secondComplete,
+    });
+    await waitFor(() => expect(secondComplete).toHaveBeenCalledOnce());
+    await waitFor(() => {
+      expect(result.current.thread.getState().suggestions).toEqual([
+        { prompt: "second" },
+      ]);
+    });
+  });
+
+  it("restarts suggestion generation across stable adapter replacements", async () => {
+    let resolveFirst!: (value: readonly { prompt: string }[]) => void;
+    const firstGenerate = vi.fn().mockImplementation(
+      () =>
+        new Promise<readonly { prompt: string }[]>((resolve) => {
+          resolveFirst = resolve;
+        }),
+    );
+    let resolveSecond!: (value: readonly { prompt: string }[]) => void;
+    const secondGenerate = vi.fn().mockImplementation(
+      () =>
+        new Promise<readonly { prompt: string }[]>((resolve) => {
+          resolveSecond = resolve;
+        }),
+    );
+    const thirdGenerate = vi.fn().mockResolvedValue([{ prompt: "third" }]);
+    const firstAdapter = { key: "first", generate: firstGenerate };
+    const secondAdapter = { key: "second", generate: secondGenerate };
+    const thirdAdapter = { key: "third", generate: thirdGenerate };
+    const chat = createChatHelpers([
+      { id: "u1", role: "user", parts: [{ type: "text", text: "hi" }] },
+      {
+        id: "a1",
+        role: "assistant",
+        parts: [{ type: "text", text: "hello" }],
+      },
+    ]);
+
+    const { result, rerender } = renderHook(
+      ({ status, adapter }) => {
+        chat.status = status;
+        return useAISDKRuntime(chat, {
+          adapters: { suggestion: adapter },
+        });
+      },
+      {
+        initialProps: {
+          status: "submitted" as string,
+          adapter: firstAdapter,
+        },
+      },
+    );
+
+    rerender({ status: "ready", adapter: firstAdapter });
+    await waitFor(() => expect(firstGenerate).toHaveBeenCalledOnce());
+    const firstSignal = firstGenerate.mock.calls[0]![0].signal as AbortSignal;
+
+    rerender({ status: "ready", adapter: secondAdapter });
+
+    expect(firstSignal.aborted).toBe(true);
+    await waitFor(() => expect(secondGenerate).toHaveBeenCalledOnce());
+    const secondSignal = secondGenerate.mock.calls[0]![0].signal as AbortSignal;
+
+    rerender({ status: "ready", adapter: thirdAdapter });
+
+    expect(secondSignal.aborted).toBe(true);
+    await waitFor(() => expect(thirdGenerate).toHaveBeenCalledOnce());
+    await waitFor(() => {
+      expect(result.current.thread.getState().suggestions).toEqual([
+        { prompt: "third" },
+      ]);
+    });
+
+    resolveFirst([{ prompt: "stale" }]);
+    resolveSecond([{ prompt: "also stale" }]);
+    await act(async () => {
+      await new Promise((resolve) => setTimeout(resolve, 20));
+    });
+    expect(result.current.thread.getState().suggestions).toEqual([
+      { prompt: "third" },
+    ]);
+  });
+
   it("merges consecutive assistant messages into one turn by default", async () => {
     const chat = createChatHelpers([
       { id: "u1", role: "user", parts: [{ type: "text", text: "hi" }] },
@@ -1799,4 +2206,130 @@ describe("useAISDKRuntime", () => {
       error: { code: "AI_APICallError", message: "upstream failed" },
     });
   });
+
+  it("keeps a cancelled output when the runtime remounts over its owner", 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 owner = {};
+
+    const first = renderHook(() =>
+      useAISDKRuntime(chat, { unstable_hostApprovalOwner: owner }),
+    );
+
+    act(() => {
+      first.result.current.thread.cancelRun();
+      chat.status = "ready";
+      first.rerender();
+    });
+
+    await waitFor(() => {
+      expect(
+        first.result.current.thread.getState().messages.at(-1)?.status,
+      ).toMatchObject({
+        type: "incomplete",
+        reason: "cancelled",
+      });
+    });
+
+    await act(async () => {
+      resolveStop();
+      await Promise.resolve();
+    });
+    first.unmount();
+
+    const second = renderHook(() =>
+      useAISDKRuntime(chat, { unstable_hostApprovalOwner: owner }),
+    );
+
+    expect(
+      second.result.current.thread.getState().messages.at(-1)?.status,
+    ).toMatchObject({
+      type: "incomplete",
+      reason: "cancelled",
+    });
+
+    act(() => {
+      chat.status = "streaming";
+      second.rerender();
+    });
+    await waitFor(() => {
+      expect(
+        second.result.current.thread.getState().messages.at(-1)?.status,
+      ).toMatchObject({ type: "running" });
+    });
+
+    act(() => {
+      chat.status = "ready";
+      second.rerender();
+    });
+    await waitFor(() => {
+      expect(
+        second.result.current.thread.getState().messages.at(-1)?.status,
+      ).toMatchObject({ type: "complete", reason: "unknown" });
+    });
+    second.unmount();
+
+    const third = renderHook(() =>
+      useAISDKRuntime(chat, { unstable_hostApprovalOwner: owner }),
+    );
+    expect(
+      third.result.current.thread.getState().messages.at(-1)?.status,
+    ).toMatchObject({ type: "complete", reason: "unknown" });
+    third.unmount();
+  });
+
+  it("keeps a cancellation when the runtime unmounts before rendering it", 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 owner = {};
+    const first = renderHook(() =>
+      useAISDKRuntime(chat, { unstable_hostApprovalOwner: owner }),
+    );
+
+    act(() => {
+      first.result.current.thread.cancelRun();
+      first.unmount();
+    });
+    expect(chat.stop).toHaveBeenCalledOnce();
+
+    chat.status = "ready";
+    await act(async () => {
+      resolveStop();
+      await Promise.resolve();
+    });
+
+    const second = renderHook(() =>
+      useAISDKRuntime(chat, { unstable_hostApprovalOwner: owner }),
+    );
+    expect(
+      second.result.current.thread.getState().messages.at(-1)?.status,
+    ).toMatchObject({ type: "incomplete", reason: "cancelled" });
+    second.unmount();
+  });
 });