@assistant-ui/ai-sdk 0.0.7 → 0.0.9

package/src/runtime/useChatThread.test.ts

@@ -3,6 +3,7 @@
 import { describe, expect, it, vi } from "vitest";
 import { resource, useResource, flushTapSync } from "@assistant-ui/tap";
 import { useState } from "react";
+import type { ChatTransport, UIMessage } from "ai";
 import {
   RuntimeAdapter,
   runtimeAdapterTransformScopes,
@@ -12,14 +13,21 @@ import {
   AuiConfig,
   createAssistantClient,
 } from "@assistant-ui/store/client";
-import { useChatThread, type ChatThreadEnvironment } from "./useChatThread";
+import {
+  createChat,
+  useChatThread,
+  type ChatThreadEnvironment,
+} from "./useChatThread";
+import { AssistantChatTransport } from "../transport/AssistantChatTransport";
+import { createResumableSessionStorage } from "../transport/resumable";
 import {
   createCancellableTransport,
   nextTask,
 } from "./__tests__/controlled-transport";
 
 const createHost = (
-  env: Pick<ChatThreadEnvironment, "stopOnClientDestroy">,
+  env: Pick<ChatThreadEnvironment, "stopOnClientDestroy" | "chat">,
+  onRuntime?: (runtime: ReturnType<typeof useChatThread>) => void,
 ) => {
   const useHost = (options: Parameters<typeof useChatThread>[0]) => {
     const [threadListItem] = useState(() => ({
@@ -31,6 +39,7 @@ const createHost = (
       getThreadListItem: () => threadListItem,
       ...env,
     });
+    onRuntime?.(runtime);
     return useResource(RuntimeAdapter(runtime));
   };
   attachTransformScopes(useHost, runtimeAdapterTransformScopes);
@@ -64,6 +73,161 @@ const streamThenDestroy = async (
 };
 
 describe("useChatThread", () => {
+  it.each([
+    { adapters: undefined, threadId: "main" },
+    {
+      adapters: { threadList: { threadId: "caller-id" } },
+      threadId: "caller-id",
+    },
+  ])("uses $threadId as the runtime thread id", ({ adapters, threadId }) => {
+    let runtime: ReturnType<typeof useChatThread> | undefined;
+    const Host = createHost({}, (value) => {
+      runtime = value;
+    });
+    const handle = createAssistantClient(
+      AuiConfig({ threads: Host({ adapters }) }),
+    );
+    try {
+      handle.subscribe(() => {});
+      expect(runtime?.thread.getState().threadId).toBe(threadId);
+    } finally {
+      handle.destroy();
+    }
+  });
+
+  it.each(["send", "resume"])(
+    "ignores errors from a subsequent %s when a resume finishes",
+    async (nextRequest) => {
+      const storage = createResumableSessionStorage({
+        key: `resume-${nextRequest}-race`,
+      });
+      storage.setStreamId("stream-1", "main");
+      const error = new Error(`${nextRequest} offline`);
+      const onError = vi.fn();
+      const onResumeError = vi.fn();
+      const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+      const transport = {
+        getResumableAdapter: () => ({ storage }),
+        reconnectToStream: vi
+          .fn<ChatTransport<UIMessage>["reconnectToStream"]>()
+          .mockRejectedValue(error)
+          .mockImplementationOnce(
+            async () =>
+              new ReadableStream({
+                start(controller) {
+                  controller.enqueue({ type: "start", messageId: "resumed" });
+                  controller.close();
+                },
+              }),
+          ),
+        sendMessages: vi.fn(async () => {
+          throw error;
+        }),
+      } satisfies ChatTransport<UIMessage> & {
+        getResumableAdapter: () => { storage: typeof storage };
+      };
+      let sent = false;
+      const callbacks = {
+        onError,
+        onFinish: () => {
+          if (!sent) {
+            sent = true;
+            if (nextRequest === "resume") {
+              void chat.resumeStream();
+            } else {
+              void chat.sendMessage({
+                role: "user",
+                parts: [{ type: "text", text: "next" }],
+              });
+            }
+          }
+        },
+      };
+      const chat = createChat(
+        {
+          id: "main",
+          transport,
+        },
+        { current: callbacks },
+      );
+      const Host = createHost({ chat });
+      const handle = createAssistantClient(
+        AuiConfig({ threads: Host({ transport, onResumeError }) }),
+      );
+      handle.subscribe(() => {});
+      try {
+        await vi.waitFor(() => expect(onError).toHaveBeenCalledWith(error));
+        await nextTask();
+        expect(onResumeError).not.toHaveBeenCalled();
+        expect(storage.getStreamId("main")).toBe("stream-1");
+      } finally {
+        handle.destroy();
+        storage.clear();
+        warn.mockRestore();
+      }
+    },
+  );
+
+  it.each([false, true])(
+    "reports SDK reconnect failures and preserves replacement checkpoints: %s",
+    async (replaceCheckpoint) => {
+      const storage = createResumableSessionStorage({
+        key: `automatic-resume-error-${replaceCheckpoint}`,
+      });
+      storage.setStreamId("failed-stream", "main");
+      const error = new Error("resume offline");
+      const onError = vi.fn();
+      const onResumeError = vi.fn(() => {
+        if (replaceCheckpoint) storage.setStreamId("replacement", "main");
+      });
+      const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+      let finishReplacement: (() => void) | undefined;
+      const fetch = vi
+        .fn<typeof globalThis.fetch>()
+        .mockRejectedValueOnce(error)
+        .mockImplementation(
+          () =>
+            new Promise<Response>((resolve) => {
+              finishReplacement = () =>
+                resolve(new Response(null, { status: 204 }));
+            }),
+        );
+      const Host = createHost({});
+      const handle = createAssistantClient(
+        AuiConfig({
+          threads: Host({
+            transport: new AssistantChatTransport({
+              fetch,
+              resumable: { storage, resumeApi: (id) => `/api/resume/${id}` },
+            }),
+            onError,
+            onResumeError,
+          }),
+        }),
+      );
+      handle.subscribe(() => {});
+      try {
+        await vi.waitFor(() => {
+          expect(onError).toHaveBeenCalledWith(error);
+          expect(onResumeError).toHaveBeenCalledOnce();
+        });
+        expect(onResumeError).toHaveBeenCalledWith(error);
+        expect(warn).toHaveBeenCalledWith(
+          "[assistant-ui] resumable: resume failed",
+          error,
+        );
+        expect(storage.getStreamId("main")).toBe(
+          replaceCheckpoint ? "replacement" : null,
+        );
+      } finally {
+        handle.destroy();
+        finishReplacement?.();
+        storage.clear();
+        warn.mockRestore();
+      }
+    },
+  );
+
   it("stops an in-flight chat on client destroy when stopOnClientDestroy is omitted", async () => {
     expect(await streamThenDestroy({})).toBe(1);
   });

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

@@ -0,0 +1,152 @@
+// @vitest-environment jsdom
+
+import { act, renderHook, waitFor } from "@testing-library/react";
+import { useEffect, version } 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 onReact18 = version.startsWith("18.");
+
+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" } },
+  );
+
+// Fails on React 18: TypeError: ReactRuntime.use is not a function. Shipped React 18 incompatibility.
+describe.skipIf(onReact18)("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();
+  });
+});

package/src/runtime/useChatThread.ts

@@ -57,6 +57,12 @@ export type ChatThreadEnvironment<UI_MESSAGE extends UIMessage = UIMessage> = {
   isMainThread: boolean;
   getThreadListItem: () => InitializableThreadListItem | undefined;
   stopOnClientDestroy?: boolean;
+  /**
+   * Aborts when the React component hosting the runtime is deleted. A nested
+   * runtime resolves the destroy signal of the provider above it, which
+   * outlives the nested component, so this stops the chat on its own unmount.
+   */
+  hostDestroySignal?: AbortSignal | undefined;
   /**
    * An externally owned chat instance. State lives on the instance, so it
    * survives the hosting resource unmounting; construction options are read
@@ -175,6 +181,8 @@ type ChatCallbacks<UI_MESSAGE extends UIMessage> = Pick<
   "onToolCall" | "onData" | "onFinish" | "onError" | "sendAutomaticallyWhen"
 >;
 
+const requestsByChat = new WeakMap<object, symbol>();
+
 /**
  * Constructs a `Chat` whose callbacks read the latest options through
  * `callbacksRef`, the forwarding `useChat` applies only to a chat it
@@ -183,9 +191,22 @@ type ChatCallbacks<UI_MESSAGE extends UIMessage> = Pick<
 export const createChat = <UI_MESSAGE extends UIMessage>(
   init: ChatInit<UI_MESSAGE>,
   callbacksRef: { readonly current: ChatCallbacks<UI_MESSAGE> | undefined },
-): Chat<UI_MESSAGE> =>
-  new Chat<UI_MESSAGE>({
+): Chat<UI_MESSAGE> => {
+  const transport = init.transport;
+  const chat = new Chat<UI_MESSAGE>({
     ...init,
+    ...(transport && {
+      transport: {
+        sendMessages: (options) => {
+          requestsByChat.set(chat, Symbol());
+          return transport.sendMessages(options);
+        },
+        reconnectToStream: (options) => {
+          requestsByChat.set(chat, Symbol());
+          return transport.reconnectToStream(options);
+        },
+      },
+    }),
     onToolCall: (arg) => callbacksRef.current?.onToolCall?.(arg),
     onData: (arg) => callbacksRef.current?.onData?.(arg),
     onFinish: (arg) => callbacksRef.current?.onFinish?.(arg),
@@ -193,6 +214,8 @@ export const createChat = <UI_MESSAGE extends UIMessage>(
     sendAutomaticallyWhen: (arg) =>
       callbacksRef.current?.sendAutomaticallyWhen?.(arg) ?? false,
   });
+  return chat;
+};
 
 export const useChatThread = <UI_MESSAGE extends UIMessage = UIMessage>(
   options: ChatThreadOptions<UI_MESSAGE> | undefined,
@@ -218,12 +241,25 @@ export const useChatThread = <UI_MESSAGE extends UIMessage = UIMessage>(
     isMainThread,
     getThreadListItem,
     stopOnClientDestroy = true,
+    hostDestroySignal,
     chat: externalChat,
     messageRepositoryInstance,
   } = env;
 
-  const defaultTransport = useMemo(() => new AssistantChatTransport(), []);
-  const sourceTransport = transportOptions ?? defaultTransport;
+  // Wiring below is per thread and mutated on the instance, so a transport
+  // shared across simultaneously mounted threads is last-writer-wins. A
+  // caller-owned chat is already bound to its own clone, so cloning again
+  // here would wire a copy the chat never sends through.
+  const sourceTransport = useMemo(
+    () =>
+      transportOptions === undefined
+        ? new AssistantChatTransport()
+        : externalChat === undefined &&
+            transportOptions instanceof AssistantChatTransport
+          ? transportOptions.__internal_clone()
+          : transportOptions,
+    [transportOptions, externalChat],
+  );
   const transport = useDynamicChatTransport(sourceTransport);
 
   const latestChatOptionsRef = useRef(chatOptions);
@@ -243,12 +279,19 @@ export const useChatThread = <UI_MESSAGE extends UIMessage = UIMessage>(
     ...(throttle !== undefined && { throttle }),
   });
 
-  useResourceCleanup(stopOnClientDestroy, () => {
-    void chat.stop().catch(() => {});
-  });
+  useResourceCleanup(
+    stopOnClientDestroy,
+    () => {
+      void chat.stop().catch(() => {});
+    },
+    hostDestroySignal,
+  );
 
   const runtime = useAISDKRuntime(chat, {
-    adapters,
+    adapters: {
+      ...adapters,
+      threadList: { threadId: id, ...adapters?.threadList },
+    },
     ...pickExternalStoreSharedOptions(options ?? {}),
     ...(toCreateMessage && { toCreateMessage }),
     ...(onResume && { onResume }),
@@ -259,9 +302,19 @@ export const useChatThread = <UI_MESSAGE extends UIMessage = UIMessage>(
     ...(messageRepositoryInstance && {
       unstable_messageRepositoryInstance: messageRepositoryInstance,
     }),
+    // The chat outlives this runtime when a host mounts only the visible
+    // thread, so a host approval answer is kept with it. This is the Chat
+    // instance, not the useChat helpers, which are re-minted every render and
+    // would be a dead WeakMap key by the next one.
+    unstable_hostApprovalOwner: externalChat ?? ownedChat,
     ...(unstable_onBranchChange && { unstable_onBranchChange }),
   });
 
+  // Wire in render, not an effect: a send from a descendant's mount effect
+  // runs before this hook's effect would (effects fire child-first), and must
+  // see a wired transport. The clone is per thread, so a discarded render's
+  // wiring is discarded with it and the committed render re-wires the same
+  // instance.
   if (sourceTransport instanceof AssistantChatTransport) {
     sourceTransport.setRuntime(runtime);
     sourceTransport.__internal_setGetThreadListItem(getThreadListItem);
@@ -322,27 +375,41 @@ export const useChatThread = <UI_MESSAGE extends UIMessage = UIMessage>(
     }
     if (isLoadingHistory) return;
     resumedStreamIds.add(pendingStreamId);
-    chat.resumeStream().catch((err: unknown) => {
-      console.warn("[assistant-ui] resumable: resume failed", err);
-      try {
-        onResumeErrorRef.current?.(err);
-      } catch (callbackError) {
-        console.error(
-          "[assistant-ui] resumable: onResumeError callback failed",
-          callbackError,
-        );
-      } finally {
-        if (resumableStorage?.getStreamId(id) === pendingStreamId) {
-          resumableStorage.clear(id);
+    const activeChat = externalChat ?? ownedChat;
+    activeChat.clearError();
+    const pending = chat.resumeStream();
+    const request = requestsByChat.get(activeChat);
+    pending
+      .then(() => {
+        // Chat.error is shared with sends and resumes that can start before
+        // this promise settles, including inside the caller's onFinish.
+        if (requestsByChat.get(activeChat) === request && activeChat.error) {
+          throw activeChat.error;
+        }
+      })
+      .catch((err: unknown) => {
+        console.warn("[assistant-ui] resumable: resume failed", err);
+        try {
+          onResumeErrorRef.current?.(err);
+        } catch (callbackError) {
+          console.error(
+            "[assistant-ui] resumable: onResumeError callback failed",
+            callbackError,
+          );
+        } finally {
+          if (resumableStorage?.getStreamId(id) === pendingStreamId) {
+            resumableStorage.clear(id);
+          }
         }
-      }
-    });
+      });
   }, [
     chat,
+    externalChat,
     id,
     isChatRunning,
     isLoadingHistory,
     pendingStreamId,
+    ownedChat,
     resumableStorage,
     resumedStreamIds,
   ]);