@assistant-ui/ai-sdk 0.0.7 → 0.0.9

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

@@ -36,7 +36,9 @@ const createChatHelpers = (initialMessages: UIMessage[] = []) => {
   return chatHelpers;
 };
 
-const createVoiceAdapter = () => {
+const createVoiceAdapter = ({
+  sendText,
+}: { sendText?: RealtimeVoiceAdapter.Session["sendText"] } = {}) => {
   let transcriptCallback:
     | ((transcript: RealtimeVoiceAdapter.TranscriptItem) => void)
     | undefined;
@@ -46,6 +48,7 @@ const createVoiceAdapter = () => {
     disconnect: vi.fn(),
     mute: vi.fn(),
     unmute: vi.fn(),
+    ...(sendText && { sendText }),
     onStatusChange: () => () => {},
     onTranscript: (callback) => {
       transcriptCallback = callback;
@@ -267,4 +270,50 @@ describe("useAISDKRuntime voice transcripts", () => {
       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.fast-refresh.test.tsx

@@ -0,0 +1,94 @@
+// @vitest-environment jsdom
+
+import { act, type ReactNode } from "react";
+import { afterAll, afterEach, expect, it, vi } from "vitest";
+
+type Family = { current: unknown };
+type RendererInternals = {
+  setRefreshHandler: (resolve: (type: unknown) => Family | undefined) => void;
+  scheduleRefresh: (
+    root: unknown,
+    update: { staleFamilies: Set<Family>; updatedFamilies: Set<Family> },
+  ) => void;
+};
+
+// Fast Refresh drives React through the renderer internals handed to the DevTools hook, so the hook has to exist before react-dom loads.
+let renderer: RendererInternals | undefined;
+const fiberRoots = new Set<unknown>();
+vi.stubGlobal("IS_REACT_ACT_ENVIRONMENT", true);
+vi.stubGlobal("__REACT_DEVTOOLS_GLOBAL_HOOK__", {
+  supportsFiber: true,
+  inject: (internals: RendererInternals) => {
+    renderer = internals;
+    return 1;
+  },
+  onScheduleFiberRoot: () => {},
+  onCommitFiberRoot: (_id: number, root: unknown) => fiberRoots.add(root),
+  onCommitFiberUnmount: () => {},
+});
+const { cleanup, render, waitFor } = await import("@testing-library/react");
+const { AssistantRuntimeProvider } = await import("@assistant-ui/core/react");
+const { useChatRuntime } = await import("./useChatRuntime");
+const { createCancellableTransport, createStreamHarness } =
+  await import("./__tests__/controlled-transport");
+
+afterEach(cleanup);
+afterAll(() => vi.unstubAllGlobals());
+
+it("keeps a running thread across a Fast Refresh of its host and aborts it on unmount", async () => {
+  const { transport, getCancelCount } = createCancellableTransport();
+  const { Probe, send, isRunning, client } = createStreamHarness();
+  const runtimes: unknown[] = [];
+  let rendered: string | undefined;
+  const createHost =
+    (name: string) =>
+    ({ children }: { children: ReactNode }) => {
+      rendered = name;
+      const runtime = useChatRuntime({ transport });
+      runtimes.push(runtime);
+      return (
+        <AssistantRuntimeProvider runtime={runtime}>
+          {children}
+        </AssistantRuntimeProvider>
+      );
+    };
+  const Before = createHost("before");
+  const After = createHost("after");
+
+  const view = render(
+    <Before>
+      <Probe />
+    </Before>,
+  );
+  await act(async () => send());
+  await waitFor(() => expect(isRunning()).toBe(true));
+  const runtime = runtimes.at(-1);
+  const threadId = client().threads.getState().mainThreadId;
+  const messages = client().thread.getState().messages;
+  const error = vi.spyOn(console, "error");
+
+  const family: Family = { current: After };
+  renderer!.setRefreshHandler((type) =>
+    type === Before || type === After ? family : undefined,
+  );
+  await act(async () => {
+    for (const fiberRoot of fiberRoots) {
+      renderer!.scheduleRefresh(fiberRoot, {
+        staleFamilies: new Set(),
+        updatedFamilies: new Set([family]),
+      });
+    }
+  });
+  await act(async () => {});
+
+  expect(rendered).toBe("after");
+  expect(runtimes.at(-1)).toBe(runtime);
+  expect(isRunning()).toBe(true);
+  expect(getCancelCount()).toBe(0);
+  expect(client().threads.getState().mainThreadId).toBe(threadId);
+  expect(client().thread.getState().messages).toEqual(messages);
+  expect(error).not.toHaveBeenCalled();
+
+  view.unmount();
+  await waitFor(() => expect(getCancelCount()).toBe(1));
+});

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

@@ -1,12 +1,14 @@
 // @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";
-import { Activity, StrictMode, useState, type ReactNode } from "react";
+import { Activity, StrictMode, useState, version, type ReactNode } from "react";
 import { describe, expect, it } from "vitest";
 import { AssistantChatTransport } from "../transport/AssistantChatTransport";
 import {
@@ -17,6 +19,8 @@ import {
 import { useChatRuntime } from "./useChatRuntime";
 import { useThreadTokenUsage } from "../usage";
 
+const onReact18 = version.startsWith("18.");
+
 const messages: UIMessage[] = [
   {
     id: "initial-user-message",
@@ -86,7 +90,12 @@ describe("useChatRuntime integration", () => {
     await act(async () => send());
     await waitFor(() => expect(isRunning()).toBe(true));
 
-    await act(async () => client().threadListItem.delete());
+    // React 18's act holds the render that settles the delete until its callback returns, so the delete is awaited outside it.
+    let deletion: unknown;
+    await act(async () => {
+      deletion = client().threadListItem.delete();
+    });
+    await deletion;
 
     await waitFor(() => expect(getCancelCount()).toBe(1));
     view.unmount();
@@ -111,78 +120,147 @@ describe("useChatRuntime integration", () => {
     await waitFor(() => expect(getCancelCount()).toBe(1));
   });
 
-  it("keeps streaming while hidden and aborts when the hidden host unmounts", async () => {
-    const { transport, getCancelCount } = createCancellableTransport();
-    const { Probe, send, isRunning } = createStreamHarness();
-
-    let setMode: ((mode: "visible" | "hidden") => void) | undefined;
-    const Shell = () => {
-      const [mode, set] = useState<"visible" | "hidden">("visible");
-      setMode = set;
-      return (
-        <Activity mode={mode}>
-          <StreamingApp transport={transport} probe={<Probe />} />
-        </Activity>
+  // Activity is React 19 only.
+  it.skipIf(onReact18)(
+    "keeps streaming while hidden and aborts when the hidden host unmounts",
+    async () => {
+      const { transport, getCancelCount } = createCancellableTransport();
+      const { Probe, send, isRunning } = createStreamHarness();
+
+      let setMode: ((mode: "visible" | "hidden") => void) | undefined;
+      const Shell = () => {
+        const [mode, set] = useState<"visible" | "hidden">("visible");
+        setMode = set;
+        return (
+          <Activity mode={mode}>
+            <StreamingApp transport={transport} probe={<Probe />} />
+          </Activity>
+        );
+      };
+
+      const view = render(
+        <StrictMode>
+          <Shell />
+        </StrictMode>,
       );
-    };
 
-    const view = render(
-      <StrictMode>
-        <Shell />
-      </StrictMode>,
-    );
-
-    await act(async () => send());
-    await waitFor(() => expect(isRunning()).toBe(true));
+      await act(async () => send());
+      await waitFor(() => expect(isRunning()).toBe(true));
 
-    await act(async () => setMode?.("hidden"));
-    await act(nextTask);
-    expect(getCancelCount()).toBe(0);
-    expect(isRunning()).toBe(true);
-
-    await act(async () => setMode?.("visible"));
-    await act(nextTask);
-    expect(getCancelCount()).toBe(0);
-    expect(isRunning()).toBe(true);
+      await act(async () => setMode?.("hidden"));
+      await act(nextTask);
+      expect(getCancelCount()).toBe(0);
+      expect(isRunning()).toBe(true);
 
-    await act(async () => setMode?.("hidden"));
-    view.unmount();
-    await waitFor(() => expect(getCancelCount()).toBe(1));
-  });
+      await act(async () => setMode?.("visible"));
+      await act(nextTask);
+      expect(getCancelCount()).toBe(0);
+      expect(isRunning()).toBe(true);
 
-  it("aborts a nested runtime's stream when the provider above it unmounts", async () => {
-    const outer = createCancellableTransport();
-    const { transport, getCancelCount } = createCancellableTransport();
-    let nested: AssistantRuntime | undefined;
+      await act(async () => setMode?.("hidden"));
+      view.unmount();
+      await waitFor(() => expect(getCancelCount()).toBe(1));
+    },
+  );
 
-    // allowNesting: the inner useChatRuntime runs its thread hook directly, as
-    // a plain React hook under the provider rather than inside a tap resource.
-    const NestedChat = () => {
-      nested = useChatRuntime({ transport });
-      return null;
-    };
+  // Fails on React 18: TypeError: ReactRuntime.use is not a function. Shipped React 18 incompatibility.
+  it.skipIf(onReact18)(
+    "aborts a nested runtime's stream when the provider above it unmounts",
+    async () => {
+      const outer = createCancellableTransport();
+      const { transport, getCancelCount } = createCancellableTransport();
+      let nested: AssistantRuntime | undefined;
+
+      // allowNesting: the inner useChatRuntime runs its thread hook directly, as
+      // a plain React hook under the provider rather than inside a tap resource.
+      const NestedChat = () => {
+        nested = useChatRuntime({ transport });
+        return null;
+      };
+
+      const view = render(
+        <StrictMode>
+          <AuiProvider
+            config={AuiConfig({
+              threads: AISDKChat({ transport: outer.transport }),
+            })}
+          >
+            <NestedChat />
+          </AuiProvider>
+        </StrictMode>,
+      );
 
-    const view = render(
-      <StrictMode>
-        <AuiProvider
-          config={AuiConfig({
-            threads: AISDKChat({ transport: outer.transport }),
-          })}
-        >
-          <NestedChat />
-        </AuiProvider>
-      </StrictMode>,
-    );
+      await waitFor(() => expect(nested).toBeDefined());
+      await act(async () => {
+        await nested!.thread.append("keep streaming");
+      });
+      await waitFor(() =>
+        expect(nested!.thread.getState().isRunning).toBe(true),
+      );
 
-    await waitFor(() => expect(nested).toBeDefined());
-    await act(async () => {
-      await nested!.thread.append("keep streaming");
-    });
-    await waitFor(() => expect(nested!.thread.getState().isRunning).toBe(true));
+      view.unmount();
+      await waitFor(() => expect(getCancelCount()).toBe(1));
+    },
+  );
 
-    view.unmount();
-    await waitFor(() => expect(getCancelCount()).toBe(1));
-  });
+  // Fails on React 18: TypeError: ReactRuntime.use is not a function. Shipped React 18 incompatibility.
+  it.skipIf(onReact18)(
+    "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

@@ -1,6 +1,7 @@
 // @vitest-environment jsdom
 
 import { act, renderHook, waitFor } from "@testing-library/react";
+import { version } from "react";
 import { beforeEach, describe, expect, it, vi } from "vitest";
 
 const mocks = vi.hoisted(() => {
@@ -51,12 +52,12 @@ const mocks = vi.hoisted(() => {
   };
 });
 
-vi.mock("@ai-sdk/react", () => ({
-  useChat: mocks.useChat,
-  Chat: class MockChat {
-    constructor(config: unknown) {
-      Object.assign(this, config);
-    }
+vi.mock("@ai-sdk/react", async (importOriginal) => ({
+  ...(await importOriginal<typeof import("@ai-sdk/react")>()),
+  useChat: (...args: unknown[]) => {
+    const chat = mocks.useChat(...args);
+    if (chat) chat.stop ??= vi.fn(async () => {});
+    return chat;
   },
 }));
 
@@ -83,6 +84,8 @@ import {
 } from "../transport/resumable";
 import { useChatRuntime } from "./useChatRuntime";
 
+const onReact18 = version.startsWith("18.");
+
 const sendMessagesOptions = {
   trigger: "submit-message" as const,
   chatId: "thread-id",
@@ -91,7 +94,8 @@ const sendMessagesOptions = {
   abortSignal: undefined,
 };
 
-describe("useChatRuntime", () => {
+// Fails on React 18: TypeError: ReactRuntime.use is not a function. Shipped React 18 incompatibility.
+describe.skipIf(onReact18)("useChatRuntime", () => {
   beforeEach(() => {
     mocks.state.isLoadingHistory = false;
     mocks.state.threadId = "thread-id";

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,