@assistant-ui/ai-sdk 0.0.6 → 0.0.8

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

@@ -0,0 +1,319 @@
+// @vitest-environment jsdom
+
+import { act, renderHook, waitFor } from "@testing-library/react";
+import {
+  getExternalStoreMessages,
+  type RealtimeVoiceAdapter,
+  type ThreadHistoryAdapter,
+  type ThreadMessage,
+} from "@assistant-ui/core";
+import type { UIMessage } from "ai";
+import { describe, expect, it, vi } from "vitest";
+
+import { useAISDKRuntime } from "./useAISDKRuntime";
+
+const createChatHelpers = (initialMessages: UIMessage[] = []) => {
+  let currentMessages: UIMessage[] = initialMessages;
+  const chatHelpers: any = {
+    id: "chat-1",
+    status: "ready",
+    error: null,
+    messages: currentMessages,
+    setMessages: vi.fn(
+      (next: UIMessage[] | ((current: UIMessage[]) => UIMessage[])) => {
+        currentMessages =
+          typeof next === "function" ? next(currentMessages) : [...next];
+        chatHelpers.messages = currentMessages;
+        return currentMessages;
+      },
+    ),
+    sendMessage: vi.fn().mockResolvedValue(undefined),
+    regenerate: vi.fn().mockResolvedValue(undefined),
+    addToolResult: vi.fn(),
+    addToolOutput: vi.fn(),
+    stop: vi.fn(),
+  };
+  return chatHelpers;
+};
+
+const createVoiceAdapter = ({
+  sendText,
+}: { sendText?: RealtimeVoiceAdapter.Session["sendText"] } = {}) => {
+  let transcriptCallback:
+    | ((transcript: RealtimeVoiceAdapter.TranscriptItem) => void)
+    | undefined;
+  const session: RealtimeVoiceAdapter.Session = {
+    status: { type: "running" },
+    isMuted: false,
+    disconnect: vi.fn(),
+    mute: vi.fn(),
+    unmute: vi.fn(),
+    ...(sendText && { sendText }),
+    onStatusChange: () => () => {},
+    onTranscript: (callback) => {
+      transcriptCallback = callback;
+      return () => {
+        transcriptCallback = undefined;
+      };
+    },
+    onModeChange: () => () => {},
+    onVolumeChange: () => () => {},
+  };
+  return {
+    adapter: { connect: () => session } satisfies RealtimeVoiceAdapter,
+    emitTranscript: (transcript: RealtimeVoiceAdapter.TranscriptItem) =>
+      transcriptCallback?.(transcript),
+  };
+};
+
+const createHistoryAdapter = () => {
+  const append = vi.fn(async () => {});
+  const adapter: ThreadHistoryAdapter = {
+    load: vi.fn(),
+    append: vi.fn(),
+    withFormat: vi.fn().mockReturnValue({
+      load: vi.fn().mockResolvedValue({ headId: null, messages: [] }),
+      append,
+    }),
+  };
+  return { adapter, append };
+};
+
+const textOf = (message: ThreadMessage) =>
+  message.content
+    .map((part) => (part.type === "text" ? part.text : ""))
+    .join("");
+
+const renderVoiceRuntime = async (
+  chat: ReturnType<typeof createChatHelpers>,
+  adapters: { voice: RealtimeVoiceAdapter; history?: ThreadHistoryAdapter },
+) => {
+  const rendered = renderHook(() => useAISDKRuntime(chat, { adapters }));
+  await waitFor(() => {
+    expect(rendered.result.current.thread.getState().capabilities.voice).toBe(
+      true,
+    );
+  });
+  return rendered;
+};
+
+describe("useAISDKRuntime voice transcripts", () => {
+  it("persists finalized transcripts through the useChat messages", async () => {
+    const chat = createChatHelpers();
+    const voice = createVoiceAdapter();
+    const { result, rerender } = await renderVoiceRuntime(chat, {
+      voice: voice.adapter,
+    });
+
+    act(() => {
+      result.current.thread.connectVoice();
+      voice.emitTranscript({
+        role: "assistant",
+        text: "Spoken reply",
+        isFinal: true,
+      });
+    });
+
+    const transcript = chat.messages[0];
+    expect(transcript).toEqual({
+      id: expect.any(String),
+      role: "assistant",
+      parts: [{ type: "text", text: "Spoken reply" }],
+      metadata: { modality: "voice" },
+    });
+
+    rerender();
+
+    await waitFor(() => {
+      const messages = result.current.thread
+        .getState()
+        .messages.filter((message) => message.id === transcript.id);
+      expect(messages).toHaveLength(1);
+      expect(messages[0]?.metadata.modality).toBe("voice");
+      expect(getExternalStoreMessages(messages[0]!)).toEqual([transcript]);
+    });
+  });
+
+  it("keeps a transcript spoken after a typed reply as its own message", async () => {
+    const chat = createChatHelpers([
+      {
+        id: "typed-user",
+        role: "user",
+        parts: [{ type: "text", text: "Hi" }],
+      },
+      {
+        id: "typed-assistant",
+        role: "assistant",
+        parts: [{ type: "text", text: "Typed reply" }],
+      },
+    ]);
+    const voice = createVoiceAdapter();
+    const { result, rerender } = await renderVoiceRuntime(chat, {
+      voice: voice.adapter,
+    });
+
+    act(() => {
+      result.current.thread.connectVoice();
+      voice.emitTranscript({
+        role: "assistant",
+        text: "Spoken reply",
+        isFinal: true,
+      });
+    });
+    const transcript = chat.messages[2];
+    rerender();
+
+    await waitFor(() => {
+      const spoken = result.current.thread.getState().messages.at(-1)!;
+      expect(getExternalStoreMessages(spoken)).toEqual([transcript]);
+    });
+    const messages = result.current.thread.getState().messages;
+    expect(messages.map((message) => message.id)).toEqual([
+      "typed-user",
+      "typed-assistant",
+      transcript.id,
+    ]);
+    expect(textOf(messages[1]!)).toBe("Typed reply");
+    expect(messages[1]?.metadata.modality).toBeUndefined();
+    expect(textOf(messages[2]!)).toBe("Spoken reply");
+    expect(messages[2]?.metadata.modality).toBe("voice");
+  });
+
+  it("keeps consecutive assistant transcripts as separate messages", async () => {
+    const chat = createChatHelpers();
+    const voice = createVoiceAdapter();
+    const { result, rerender } = await renderVoiceRuntime(chat, {
+      voice: voice.adapter,
+    });
+
+    act(() => {
+      result.current.thread.connectVoice();
+      voice.emitTranscript({ role: "assistant", text: "First", isFinal: true });
+      voice.emitTranscript({
+        role: "assistant",
+        text: "Second",
+        isFinal: true,
+      });
+    });
+    const [first, second] = chat.messages;
+    rerender();
+
+    await waitFor(() => {
+      const messages = result.current.thread.getState().messages;
+      expect(
+        messages.map((message) => getExternalStoreMessages(message)),
+      ).toEqual([[first], [second]]);
+    });
+    expect(result.current.thread.getState().messages.map(textOf)).toEqual([
+      "First",
+      "Second",
+    ]);
+  });
+
+  it("keeps the transcript in the thread after the session disconnects", async () => {
+    const chat = createChatHelpers();
+    const voice = createVoiceAdapter();
+    const { result, rerender } = await renderVoiceRuntime(chat, {
+      voice: voice.adapter,
+    });
+
+    act(() => {
+      result.current.thread.connectVoice();
+      voice.emitTranscript({
+        role: "assistant",
+        text: "Spoken reply",
+        isFinal: true,
+      });
+    });
+    const transcript = chat.messages[0];
+    rerender();
+    await waitFor(() => {
+      const messages = result.current.thread.getState().messages;
+      expect(getExternalStoreMessages(messages[0]!)).toEqual([transcript]);
+    });
+
+    act(() => {
+      result.current.thread.disconnectVoice();
+    });
+    rerender();
+
+    const messages = result.current.thread.getState().messages;
+    expect(messages).toHaveLength(1);
+    expect(messages[0]?.id).toBe(transcript.id);
+    expect(messages[0]?.metadata.modality).toBe("voice");
+    expect(result.current.thread.getState().voice).toBeUndefined();
+  });
+
+  it("appends a finalized transcript to the history adapter without a text run", async () => {
+    const chat = createChatHelpers();
+    const voice = createVoiceAdapter();
+    const history = createHistoryAdapter();
+    const { result, rerender } = await renderVoiceRuntime(chat, {
+      voice: voice.adapter,
+      history: history.adapter,
+    });
+
+    act(() => {
+      result.current.thread.connectVoice();
+      voice.emitTranscript({
+        role: "assistant",
+        text: "Spoken reply",
+        isFinal: true,
+      });
+    });
+    const transcript = chat.messages[0];
+    rerender();
+
+    await waitFor(() => expect(history.append).toHaveBeenCalledTimes(1));
+    expect(history.append).toHaveBeenCalledWith({
+      parentId: null,
+      message: transcript,
+    });
+  });
+
+  it("persists a message typed into the session as a typed turn", async () => {
+    const chat = createChatHelpers();
+    const sendText = vi.fn(async (_text: string) => {});
+    const voice = createVoiceAdapter({ sendText });
+    const history = createHistoryAdapter();
+    const { result, rerender } = await renderVoiceRuntime(chat, {
+      voice: voice.adapter,
+      history: history.adapter,
+    });
+
+    act(() => {
+      result.current.thread.connectVoice();
+    });
+    await act(async () => {
+      await result.current.thread.append({
+        role: "user",
+        content: [{ type: "text", text: "Typed" }],
+      });
+    });
+
+    expect(sendText).toHaveBeenCalledExactlyOnceWith("Typed");
+    expect(chat.sendMessage).not.toHaveBeenCalled();
+    const typed = chat.messages[0];
+    expect(typed).toEqual({
+      id: expect.any(String),
+      role: "user",
+      parts: [{ type: "text", text: "Typed" }],
+      metadata: {},
+    });
+    rerender();
+
+    await waitFor(() => {
+      const messages = result.current.thread
+        .getState()
+        .messages.filter((message) => message.id === typed.id);
+      expect(messages).toHaveLength(1);
+      expect(messages[0]?.metadata.modality).toBeUndefined();
+      expect(textOf(messages[0]!)).toBe("Typed");
+    });
+    await waitFor(() => expect(history.append).toHaveBeenCalledTimes(1));
+    expect(history.append).toHaveBeenCalledWith({
+      parentId: null,
+      message: typed,
+    });
+  });
+});

package/src/runtime/useChatRuntime.integration.test.tsx

@@ -1,8 +1,10 @@
 // @vitest-environment jsdom
 
+import { getEventListeners } from "node:events";
 import { act, render, screen, waitFor } from "@testing-library/react";
 import { AssistantRuntimeProvider } from "@assistant-ui/core/react";
 import { AuiConfig, AuiProvider, useAuiState } from "@assistant-ui/store";
+import { useAssistantClientDestroySignal } from "@assistant-ui/store/internal";
 import type { AssistantRuntime } from "@assistant-ui/core";
 import { AISDKChat } from "./AISDKChat";
 import type { ChatTransport, UIMessage } from "ai";
@@ -183,6 +185,60 @@ describe("useChatRuntime integration", () => {
     view.unmount();
     await waitFor(() => expect(getCancelCount()).toBe(1));
   });
+
+  it("aborts a nested runtime when only its own component unmounts", async () => {
+    const outer = createCancellableTransport();
+    const { transport, getCancelCount } = createCancellableTransport();
+    let nested: AssistantRuntime | undefined;
+    let providerSignal: AbortSignal | undefined;
+    let setVisible: ((visible: boolean) => void) | undefined;
+
+    const NestedChat = () => {
+      providerSignal = useAssistantClientDestroySignal();
+      nested = useChatRuntime({ transport });
+      return null;
+    };
+    const Shell = () => {
+      const [visible, set] = useState(true);
+      setVisible = set;
+      return (
+        <AuiProvider
+          config={AuiConfig({
+            threads: AISDKChat({ transport: outer.transport }),
+          })}
+        >
+          {visible && <NestedChat />}
+        </AuiProvider>
+      );
+    };
+    const listeners = () => getEventListeners(providerSignal!, "abort").length;
+
+    const view = render(<Shell />);
+    await waitFor(() => expect(nested).toBeDefined());
+    const mounted = listeners();
+
+    for (const cycle of [1, 2]) {
+      await act(async () => {
+        await nested!.thread.append(`stream ${cycle}`);
+      });
+      await waitFor(() =>
+        expect(nested!.thread.getState().isRunning).toBe(true),
+      );
+
+      await act(async () => setVisible?.(false));
+      await waitFor(() => expect(getCancelCount()).toBe(cycle));
+      expect(listeners()).toBe(mounted - 1);
+
+      nested = undefined;
+      await act(async () => setVisible?.(true));
+      await waitFor(() => expect(nested).toBeDefined());
+      expect(listeners()).toBe(mounted);
+    }
+
+    view.unmount();
+    await act(nextTask);
+    expect(getCancelCount()).toBe(2);
+  });
 });
 
 const StreamingApp = ({

package/src/runtime/useChatRuntime.test.ts

@@ -52,7 +52,16 @@ const mocks = vi.hoisted(() => {
 });
 
 vi.mock("@ai-sdk/react", () => ({
-  useChat: mocks.useChat,
+  useChat: (...args: unknown[]) => {
+    const chat = mocks.useChat(...args);
+    if (chat) chat.stop ??= vi.fn(async () => {});
+    return chat;
+  },
+  Chat: class MockChat {
+    constructor(config: unknown) {
+      Object.assign(this, config);
+    }
+  },
 }));
 
 vi.mock("@assistant-ui/core/react", async (importOriginal) => ({
@@ -95,6 +104,51 @@ describe("useChatRuntime", () => {
     window.sessionStorage.clear();
   });
 
+  it("forwards a callback through a ref, so a later render's callback fires instead of the mounted one", () => {
+    mocks.useChat.mockReturnValue({
+      resumeStream: vi.fn(),
+      status: "ready",
+    });
+
+    const onToolCallA = vi.fn();
+    const onToolCallB = vi.fn();
+
+    const { rerender } = renderHook(
+      ({ onToolCall }: { onToolCall: typeof onToolCallA }) =>
+        useChatRuntime({ onToolCall }),
+      { initialProps: { onToolCall: onToolCallA } },
+    );
+
+    const chat = mocks.useChat.mock.calls[0]?.[0]?.chat as {
+      onToolCall?: (arg: unknown) => void;
+      sendAutomaticallyWhen?: (arg: unknown) => boolean;
+    };
+
+    chat.onToolCall?.("first");
+    expect(onToolCallA).toHaveBeenCalledExactlyOnceWith("first");
+
+    rerender({ onToolCall: onToolCallB });
+    chat.onToolCall?.("second");
+
+    expect(onToolCallB).toHaveBeenCalledExactlyOnceWith("second");
+    expect(onToolCallA).toHaveBeenCalledOnce();
+  });
+
+  it("coerces an unset sendAutomaticallyWhen to false, matching useChat's own default", () => {
+    mocks.useChat.mockReturnValue({
+      resumeStream: vi.fn(),
+      status: "ready",
+    });
+
+    renderHook(() => useChatRuntime());
+
+    const chat = mocks.useChat.mock.calls[0]?.[0]?.chat as {
+      sendAutomaticallyWhen?: (arg: unknown) => boolean;
+    };
+
+    expect(chat.sendAutomaticallyWhen?.({})).toBe(false);
+  });
+
   it("forwards a defined chat update throttle to useChat", () => {
     mocks.useChat.mockReturnValue({
       resumeStream: vi.fn(),
@@ -111,6 +165,24 @@ describe("useChatRuntime", () => {
     expect(mocks.useChat.mock.calls[1]?.[0]).not.toHaveProperty("throttle");
   });
 
+  it("forwards a custom approval handler to the runtime only", () => {
+    const onRespondToToolApproval = vi.fn();
+    mocks.useChat.mockReturnValue({
+      resumeStream: vi.fn(),
+      status: "ready",
+    });
+
+    renderHook(() => useChatRuntime({ onRespondToToolApproval }));
+
+    expect(mocks.useAISDKRuntime).toHaveBeenCalledWith(
+      expect.anything(),
+      expect.objectContaining({ onRespondToToolApproval }),
+    );
+    expect(mocks.useChat.mock.calls[0]?.[0]).not.toHaveProperty(
+      "onRespondToToolApproval",
+    );
+  });
+
   it("waits for external history to load before resuming a stream", async () => {
     mocks.state.isLoadingHistory = true;
     const resumeStream = vi.fn().mockResolvedValue(undefined);
@@ -283,8 +355,8 @@ describe("useChatRuntime", () => {
       resumeStream: vi.fn().mockResolvedValue(undefined),
       status: "streaming",
     };
-    mocks.useChat.mockImplementation(({ id }: { id: string }) =>
-      id === "thread-a" ? threadA : threadB,
+    mocks.useChat.mockImplementation(({ chat }: { chat: { id: string } }) =>
+      chat.id === "thread-a" ? threadA : threadB,
     );
 
     mocks.state.threadId = "thread-a";
@@ -321,8 +393,8 @@ describe("useChatRuntime", () => {
       resumeStream: vi.fn().mockResolvedValue(undefined),
       status: "ready",
     };
-    mocks.useChat.mockImplementation(({ id }: { id: string }) =>
-      id === "__LOCALID_background" ? backgroundThread : mainThread,
+    mocks.useChat.mockImplementation(({ chat }: { chat: { id: string } }) =>
+      chat.id === "__LOCALID_background" ? backgroundThread : mainThread,
     );
     const transport = {
       getResumableAdapter: () => ({

package/src/runtime/useChatRuntime.ts

@@ -8,6 +8,7 @@ import {
   useRemoteThreadListRuntime,
 } from "@assistant-ui/core/react";
 import { useAui, useAuiState } from "@assistant-ui/store";
+import { useHostDestroySignal } from "@assistant-ui/store/internal";
 import { useChatThread, type ChatThreadOptions } from "./useChatThread";
 import { AI_SDK_SDK } from "./sdkIdentity";
 
@@ -18,7 +19,8 @@ export type UseChatRuntimeOptions<UI_MESSAGE extends UIMessage = UIMessage> =
   };
 
 const useChatThreadRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
-  options?: ChatThreadOptions<UI_MESSAGE>,
+  options: ChatThreadOptions<UI_MESSAGE> | undefined,
+  hostDestroySignal: AbortSignal,
 ): AssistantRuntime => {
   const id = useAuiState((s) => s.threadListItem.id);
   const isMainThread = useAuiState(
@@ -31,6 +33,7 @@ const useChatThreadRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
     getThreadListItem: () =>
       aui.threadListItem.source ? aui.threadListItem : undefined,
     stopOnClientDestroy: true,
+    hostDestroySignal,
   });
 };
 
@@ -39,10 +42,11 @@ export const useChatRuntime = <UI_MESSAGE extends UIMessage = UIMessage>({
   onThreadIdChange,
   ...options
 }: UseChatRuntimeOptions<UI_MESSAGE> = {}): AssistantRuntime => {
+  const hostDestroySignal = useHostDestroySignal();
   const cloudAdapter = useCloudThreadListAdapter({ cloud, sdk: AI_SDK_SDK });
   return useRemoteThreadListRuntime({
     runtimeHook: function RuntimeHook() {
-      return useChatThreadRuntime(options);
+      return useChatThreadRuntime(options, hostDestroySignal);
     },
     adapter: cloudAdapter,
     allowNesting: true,

package/src/runtime/useChatThread.transport.test.tsx

@@ -0,0 +1,149 @@
+// @vitest-environment jsdom
+
+import { act, renderHook, waitFor } from "@testing-library/react";
+import { useEffect } from "react";
+import { describe, expect, it, vi } from "vitest";
+import { Chat } from "@ai-sdk/react";
+import {
+  AssistantChatTransport,
+  type InitializableThreadListItem,
+} from "../transport/AssistantChatTransport";
+import { useChatThread } from "./useChatThread";
+
+const itemFor = (remoteId: string) => ({
+  initialize: async () => ({ remoteId, externalId: undefined }),
+});
+
+const finishedStream = () =>
+  new Response(
+    new ReadableStream<Uint8Array>({
+      start(controller) {
+        controller.enqueue(
+          new TextEncoder().encode(
+            'data: {"type":"start"}\n\ndata: {"type":"finish"}\n\ndata: [DONE]\n\n',
+          ),
+        );
+        controller.close();
+      },
+    }),
+    { headers: { "content-type": "text/event-stream" } },
+  );
+
+describe("useChatThread shared transport isolation", () => {
+  it("gives each thread its own clone wired to its own thread-list item", async () => {
+    const transport = new AssistantChatTransport({ api: "/api/chat" });
+    const setRuntime = vi.spyOn(transport, "setRuntime");
+    const setGetItem = vi.spyOn(transport, "__internal_setGetThreadListItem");
+    const clones: AssistantChatTransport<never>[] = [];
+    const getters: (() => InitializableThreadListItem | undefined)[] = [];
+    const realClone = transport.__internal_clone.bind(transport);
+    vi.spyOn(transport, "__internal_clone").mockImplementation(() => {
+      const clone = realClone();
+      const index = clones.push(clone as AssistantChatTransport<never>) - 1;
+      vi.spyOn(clone, "__internal_setGetThreadListItem").mockImplementation(
+        (getter) => {
+          getters[index] = getter;
+        },
+      );
+      return clone;
+    });
+
+    renderHook(() =>
+      useChatThread(
+        { transport },
+        {
+          id: "thread-a",
+          isMainThread: true,
+          getThreadListItem: () => itemFor("remote-a"),
+        },
+      ),
+    );
+    renderHook(() =>
+      useChatThread(
+        { transport },
+        {
+          id: "thread-b",
+          isMainThread: false,
+          getThreadListItem: () => itemFor("remote-b"),
+        },
+      ),
+    );
+
+    expect(setRuntime).not.toHaveBeenCalled();
+    expect(setGetItem).not.toHaveBeenCalled();
+    expect(clones).toHaveLength(2);
+    expect(clones[0]).not.toBe(clones[1]);
+
+    const remoteIdOf = async (
+      getter?: () => InitializableThreadListItem | undefined,
+    ) => (await getter?.()?.initialize())?.remoteId;
+    expect(await remoteIdOf(getters[0])).toBe("remote-a");
+    expect(await remoteIdOf(getters[1])).toBe("remote-b");
+  });
+
+  it("sends each thread's request with that thread's remoteId and model context", async () => {
+    const sent: { id: unknown; system: unknown }[] = [];
+    const transport = new AssistantChatTransport({
+      api: "/api/chat",
+      fetch: async (_input, init) => {
+        const body = JSON.parse(String(init?.body)) as Record<string, unknown>;
+        sent.push({ id: body["id"], system: body["system"] });
+        return finishedStream();
+      },
+    });
+
+    const mount = (id: string, remoteId: string, system: string) =>
+      renderHook(() => {
+        const runtime = useChatThread(
+          { transport },
+          {
+            id,
+            isMainThread: id === "thread-a",
+            getThreadListItem: () => itemFor(remoteId),
+          },
+        );
+        useEffect(
+          () =>
+            runtime.registerModelContextProvider({
+              getModelContext: () => ({ system }),
+            }),
+          [runtime],
+        );
+        return runtime;
+      });
+
+    const threadA = mount("thread-a", "remote-a", "system-a");
+    const threadB = mount("thread-b", "remote-b", "system-b");
+
+    threadA.rerender();
+
+    await act(async () => {
+      threadB.result.current.thread.append("hello");
+    });
+
+    await waitFor(() => expect(sent).toHaveLength(1));
+    expect(sent[0]).toEqual({ id: "remote-b", system: "system-b" });
+  });
+
+  it("uses the supplied instance when the caller owns the chat", () => {
+    const transport = new AssistantChatTransport({ api: "/api/chat" });
+    const clone = vi.spyOn(transport, "__internal_clone");
+    const chat = new Chat({ id: "thread-a", transport });
+
+    renderHook(() =>
+      useChatThread(
+        { transport },
+        {
+          id: "thread-a",
+          isMainThread: true,
+          getThreadListItem: () => itemFor("remote-a"),
+          chat,
+        },
+      ),
+    );
+
+    // A caller-owned chat is already bound to its transport, so cloning again
+    // would wire an instance the chat never sends through.
+    expect(clone).not.toHaveBeenCalled();
+  });
+});