@assistant-ui/ai-sdk 0.0.10 → 0.0.12

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

@@ -1,15 +1,37 @@
+import { createRoot } from "react-dom/client";
+import { useTapHost } from "@assistant-ui/tap";
+import { DynamicChatTransport } from "./DynamicChatTransport";
+import { useChatThread } from "./useChatThread";
 // @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 {
+  act,
+  render,
+  renderHook,
+  screen,
+  waitFor,
+} from "@testing-library/react";
+import {
+  AssistantRuntimeProvider,
+  RuntimeAdapterProvider,
+} 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, version, type ReactNode } from "react";
-import { describe, expect, it } from "vitest";
+import {
+  Activity,
+  StrictMode,
+  Suspense,
+  useEffect,
+  useLayoutEffect,
+  useState,
+  version,
+  type ReactNode,
+} from "react";
+import { describe, expect, it, vi } from "vitest";
 import { AssistantChatTransport } from "../transport/AssistantChatTransport";
 import {
   createCancellableTransport,
@@ -322,3 +344,296 @@ describe("useThreadTokenUsage through useChatRuntime", () => {
     });
   });
 });
+
+describe("replacement transports", () => {
+  it.skipIf(onReact18)(
+    "does not loop when an inline transport recreates resumable storage",
+    async () => {
+      const fetch = vi.fn(
+        async () =>
+          new Response(
+            new ReadableStream({ start: (controller) => controller.close() }),
+            { headers: { "content-type": "text/event-stream" } },
+          ),
+      );
+      let renders = 0;
+      let sent = false;
+      const SendOnLayout = ({ runtime }: { runtime: AssistantRuntime }) => {
+        useLayoutEffect(() => {
+          if (sent) return;
+          sent = true;
+          void runtime.thread.append("hello");
+        }, [runtime]);
+        return null;
+      };
+      const NestedChat = () => {
+        renders += 1;
+        const runtime = useChatRuntime({
+          transport: new AssistantChatTransport({
+            fetch,
+            resumable: {
+              storage: {
+                getStreamId: () => null,
+                setStreamId: vi.fn(),
+                clear: vi.fn(),
+              },
+              resumeApi: "/api/chat/resume",
+            },
+          }),
+        });
+        return <SendOnLayout runtime={runtime} />;
+      };
+
+      render(
+        <AuiProvider
+          config={AuiConfig({
+            threads: AISDKChat({ transport: new AssistantChatTransport() }),
+          })}
+        >
+          <NestedChat />
+        </AuiProvider>,
+      );
+      await waitFor(() => expect(fetch).toHaveBeenCalledOnce());
+
+      expect(renders).toBeLessThan(10);
+    },
+  );
+
+  it("routes sends through a replacement transport", async () => {
+    const createEmptyStream = () =>
+      new ReadableStream({ start: (controller) => controller.close() });
+    const sendA = vi.fn(async () => createEmptyStream());
+    const sendB = vi.fn(async () => createEmptyStream());
+    const transportA: ChatTransport<UIMessage> = {
+      sendMessages: sendA,
+      reconnectToStream: vi.fn(),
+    };
+    const transportB: ChatTransport<UIMessage> = {
+      sendMessages: sendB,
+      reconnectToStream: vi.fn(),
+    };
+    const SendOnLayout = ({ runtime }: { runtime: AssistantRuntime }) => {
+      useLayoutEffect(() => {
+        void runtime.thread.append({
+          role: "user",
+          content: [{ type: "text", text: "hello" }],
+        });
+      }, [runtime]);
+      return null;
+    };
+    const App = ({
+      transport,
+      send = false,
+    }: {
+      transport: ChatTransport<UIMessage>;
+      send?: boolean;
+    }) => {
+      const runtime = useChatRuntime({ transport });
+      return (
+        <AssistantRuntimeProvider runtime={runtime}>
+          {send && <SendOnLayout runtime={runtime} />}
+        </AssistantRuntimeProvider>
+      );
+    };
+
+    const view = render(<App transport={transportA} />);
+    view.rerender(<App transport={transportB} send />);
+
+    await waitFor(() => expect(sendB).toHaveBeenCalledOnce());
+    expect(sendA).not.toHaveBeenCalled();
+  });
+
+  it("routes sends triggered during the initial layout commit", async () => {
+    const send = vi.fn(
+      async () =>
+        new ReadableStream({ start: (controller) => controller.close() }),
+    );
+    const transport: ChatTransport<UIMessage> = {
+      sendMessages: send,
+      reconnectToStream: vi.fn(),
+    };
+    let resolveLayout!: () => void;
+    const layoutCommitted = new Promise<void>((resolve) => {
+      resolveLayout = resolve;
+    });
+    const SendOnLayout = ({ runtime }: { runtime: AssistantRuntime }) => {
+      useLayoutEffect(() => {
+        runtime.thread.append({
+          role: "user",
+          content: [{ type: "text", text: "hello" }],
+        });
+        resolveLayout();
+      }, [runtime]);
+      return null;
+    };
+    const App = () => {
+      const runtime = useChatRuntime({ transport });
+      return (
+        <AssistantRuntimeProvider runtime={runtime}>
+          <SendOnLayout runtime={runtime} />
+        </AssistantRuntimeProvider>
+      );
+    };
+    const actEnvironment = globalThis as typeof globalThis & {
+      IS_REACT_ACT_ENVIRONMENT?: boolean;
+    };
+    const previousActEnvironment = actEnvironment.IS_REACT_ACT_ENVIRONMENT;
+    const container = document.createElement("div");
+    const root = createRoot(container);
+
+    actEnvironment.IS_REACT_ACT_ENVIRONMENT = false;
+    document.body.append(container);
+    try {
+      root.render(<App />);
+      await layoutCommitted;
+      await new Promise((resolve) => setTimeout(resolve, 0));
+
+      expect(send).toHaveBeenCalledOnce();
+    } finally {
+      root.unmount();
+      container.remove();
+      if (previousActEnvironment === undefined) {
+        delete actEnvironment.IS_REACT_ACT_ENVIRONMENT;
+      } else {
+        actEnvironment.IS_REACT_ACT_ENVIRONMENT = previousActEnvironment;
+      }
+    }
+  });
+
+  it("keeps AssistantChatTransport wired through StrictMode effect replay", () => {
+    const sourceTransport = new AssistantChatTransport<UIMessage>({
+      api: "/api/chat",
+    });
+    const transport = new DynamicChatTransport(sourceTransport);
+    const wiredDuringEffectSetup: boolean[] = [];
+    const Probe = ({ effects }: { effects: () => void }) => {
+      useEffect(effects);
+      useEffect(() => {
+        try {
+          wiredDuringEffectSetup.push(
+            transport.getCurrentTransport("strict-mode-thread") !==
+              sourceTransport,
+          );
+        } catch {
+          wiredDuringEffectSetup.push(false);
+        }
+      }, []);
+      return null;
+    };
+
+    const App = () => {
+      const { effects } = useTapHost(function ChatThreadResource() {
+        return useChatThread(
+          { transport },
+          {
+            id: "strict-mode-thread",
+            isMainThread: true,
+            getThreadListItem: () => undefined,
+          },
+        );
+      });
+      return <Probe effects={effects} />;
+    };
+
+    render(
+      <StrictMode>
+        <App />
+      </StrictMode>,
+    );
+
+    expect(wiredDuringEffectSetup).toEqual([true, true]);
+  });
+
+  it("does not publish transport contexts from suspended renders", () => {
+    const sourceTransport = new AssistantChatTransport<UIMessage>({
+      api: "/api/chat",
+    });
+    const transport = new DynamicChatTransport(sourceTransport);
+    const pending = new Promise<never>(() => {});
+    const SuspendedThread = () => {
+      useTapHost(function ChatThreadResource() {
+        return useChatThread(
+          { transport },
+          {
+            id: "suspended-thread",
+            isMainThread: true,
+            getThreadListItem: () => undefined,
+          },
+        );
+      });
+      throw pending;
+    };
+
+    render(
+      <Suspense fallback={null}>
+        <SuspendedThread />
+      </Suspense>,
+    );
+
+    expect(transport.getCurrentTransport("suspended-thread")).toBe(
+      sourceTransport,
+    );
+  });
+
+  it.skipIf(onReact18)(
+    "routes runtime sends through a wired clone with the latest thread item",
+    async () => {
+      const bodies: Array<{ id: string; system: string }> = [];
+      const sourceTransport = new AssistantChatTransport<UIMessage>({
+        fetch: vi.fn(async (_input, init) => {
+          bodies.push(JSON.parse(String(init?.body)));
+          return new Response(
+            new ReadableStream({ start: (controller) => controller.close() }),
+            { headers: { "content-type": "text/event-stream" } },
+          );
+        }),
+      });
+      const transport = new DynamicChatTransport(sourceTransport);
+      const wrapper = ({ children }: { children: ReactNode }) => (
+        <RuntimeAdapterProvider
+          adapters={{
+            modelContext: {
+              getModelContext: () => ({ system: "system prompt" }),
+            },
+          }}
+        >
+          {children}
+        </RuntimeAdapterProvider>
+      );
+      const { result, rerender } = renderHook(
+        ({ remoteId }: { remoteId: string }) =>
+          useChatThread(
+            { transport },
+            {
+              id: "stable-thread",
+              isMainThread: true,
+              getThreadListItem: () => ({
+                initialize: async () => ({
+                  remoteId,
+                  externalId: undefined,
+                }),
+              }),
+            },
+          ),
+        { initialProps: { remoteId: "remote-a" }, wrapper },
+      );
+      const send = async () => {
+        await act(async () => {
+          await result.current.thread.append({
+            role: "user",
+            content: [{ type: "text", text: "hello" }],
+          });
+        });
+      };
+
+      await send();
+      rerender({ remoteId: "remote-b" });
+      await send();
+
+      expect(bodies).toEqual([
+        expect.objectContaining({ id: "remote-a", system: "system prompt" }),
+        expect.objectContaining({ id: "remote-b", system: "system prompt" }),
+      ]);
+    },
+  );
+});

package/src/runtime/useChatRuntime.local-storage.test.tsx

@@ -0,0 +1,123 @@
+// @vitest-environment jsdom
+
+import { act, render, waitFor } from "@testing-library/react";
+import {
+  AssistantRuntimeProvider,
+  createLocalStorageAdapter,
+  useRemoteThreadListRuntime,
+  type AsyncStorageLike,
+} from "@assistant-ui/core/react";
+import type { ChatTransport, UIMessage } from "ai";
+import { useState, type ReactNode } from "react";
+import { describe, expect, it } from "vitest";
+import {
+  createControlledTransport,
+  createStreamHarness,
+} from "./__tests__/controlled-transport";
+import { useChatRuntime } from "./useChatRuntime";
+
+const createStorage = (): AsyncStorageLike => {
+  const values = new Map<string, string>();
+  return {
+    getItem: async (key) => values.get(key) ?? null,
+    setItem: async (key, value) => {
+      values.set(key, value);
+    },
+    removeItem: async (key) => {
+      values.delete(key);
+    },
+  };
+};
+
+const LocalStorageApp = ({
+  storage,
+  transport,
+  threadId,
+  probe,
+}: {
+  storage: AsyncStorageLike;
+  transport: ChatTransport<UIMessage>;
+  threadId?: string;
+  probe: ReactNode;
+}) => {
+  const [adapter] = useState(() => createLocalStorageAdapter({ storage }));
+  const runtime = useRemoteThreadListRuntime({
+    runtimeHook: function RuntimeHook() {
+      return useChatRuntime({ transport });
+    },
+    adapter,
+    ...(threadId !== undefined ? { threadId } : undefined),
+  });
+  return (
+    <AssistantRuntimeProvider runtime={runtime}>
+      {probe}
+    </AssistantRuntimeProvider>
+  );
+};
+
+const threadText = (
+  harness: ReturnType<typeof createStreamHarness>,
+): string[] =>
+  harness
+    .client()
+    .thread.getState()
+    .messages.map((message) =>
+      message.content
+        .map((part) => (part.type === "text" ? part.text : ""))
+        .join(""),
+    );
+
+describe("useChatRuntime with createLocalStorageAdapter", () => {
+  it("persists AI SDK messages and restores them after a remount", async () => {
+    const storage = createStorage();
+    const { transport, emit, close } = createControlledTransport();
+    const harness = createStreamHarness();
+
+    const first = render(
+      <LocalStorageApp
+        storage={storage}
+        transport={transport}
+        probe={<harness.Probe />}
+      />,
+    );
+
+    await act(async () => harness.send());
+    await waitFor(() => expect(harness.isRunning()).toBe(true));
+    await act(async () => {
+      emit(
+        { type: "start" },
+        { type: "text-start", id: "t1" },
+        { type: "text-delta", id: "t1", delta: "stored answer" },
+        { type: "text-end", id: "t1" },
+        { type: "finish" },
+      );
+      close();
+    });
+    await waitFor(() => expect(harness.isRunning()).toBe(false));
+
+    const remoteId = harness.client().threadListItem.getState().remoteId;
+    expect(remoteId).toBeDefined();
+    await waitFor(async () =>
+      expect(
+        await storage.getItem(
+          `@assistant-ui:formatted-messages:${JSON.stringify([remoteId, "ai-sdk/v6"])}`,
+        ),
+      ).toContain("stored answer"),
+    );
+    first.unmount();
+
+    const restored = createStreamHarness();
+    render(
+      <LocalStorageApp
+        storage={storage}
+        transport={createControlledTransport().transport}
+        threadId={remoteId!}
+        probe={<restored.Probe />}
+      />,
+    );
+
+    await waitFor(() =>
+      expect(threadText(restored)).toEqual(["keep streaming", "stored answer"]),
+    );
+  });
+});

package/src/runtime/useChatRuntime.test.ts

@@ -1,8 +1,10 @@
+import type { ChatTransport, UIMessage } from "ai";
 // @vitest-environment jsdom
 
 import { act, renderHook, waitFor } from "@testing-library/react";
 import { version } from "react";
 import { beforeEach, describe, expect, it, vi } from "vitest";
+import type { AssistantCloud } from "assistant-cloud";
 
 const mocks = vi.hoisted(() => {
   const state = {
@@ -17,7 +19,9 @@ const mocks = vi.hoisted(() => {
       getModelContext: () => ({}),
       subscribe: (callback: () => void) => {
         subscribers.add(callback);
-        return () => subscribers.delete(callback);
+        return () => {
+          subscribers.delete(callback);
+        };
       },
     },
     threads: {
@@ -82,6 +86,8 @@ import {
   createResumableSessionStorage,
   RESUMABLE_STREAM_ID_HEADER,
 } from "../transport/resumable";
+import { getResumedStreamIds } from "./DynamicChatTransport";
+import { AI_SDK_SDK } from "./sdkIdentity";
 import { useChatRuntime } from "./useChatRuntime";
 
 const onReact18 = version.startsWith("18.");
@@ -97,6 +103,7 @@ const sendMessagesOptions = {
 // Fails on React 18: TypeError: ReactRuntime.use is not a function. Shipped React 18 incompatibility.
 describe.skipIf(onReact18)("useChatRuntime", () => {
   beforeEach(() => {
+    mocks.useAISDKRuntime.mockImplementation(() => mocks.runtime);
     mocks.state.isLoadingHistory = false;
     mocks.state.threadId = "thread-id";
     mocks.state.mainThreadId = "thread-id";
@@ -104,6 +111,63 @@ describe.skipIf(onReact18)("useChatRuntime", () => {
     window.sessionStorage.clear();
   });
 
+  it("refreshes AssistantChatTransport wiring when the runtime changes", async () => {
+    const bodies: Array<{ system: string }> = [];
+    const transport = new AssistantChatTransport({
+      fetch: vi.fn(async (_input, init) => {
+        bodies.push(JSON.parse(String(init?.body)));
+        return new Response(
+          new ReadableStream({ start: (controller) => controller.close() }),
+          { headers: { "content-type": "text/event-stream" } },
+        );
+      }),
+    });
+    const createRuntime = (system: string) => ({
+      thread: {
+        getState: () => ({ isLoading: false }),
+        getModelContext: () => ({ system }),
+        subscribe: () => () => {},
+      },
+      threads: { mainItem: undefined },
+    });
+    let currentRuntime = createRuntime("system-a");
+    mocks.useAISDKRuntime.mockImplementation(() => currentRuntime);
+    mocks.useChat.mockReturnValue({
+      resumeStream: vi.fn(),
+      status: "ready",
+    });
+    const { rerender } = renderHook(() => useChatRuntime({ transport }));
+    const dynamicTransport = mocks.useChat.mock.lastCall?.[0].chat
+      .transport as ChatTransport<UIMessage>;
+
+    await dynamicTransport.sendMessages(sendMessagesOptions as never);
+    currentRuntime = createRuntime("system-b");
+    rerender();
+    await dynamicTransport.sendMessages(sendMessagesOptions as never);
+
+    expect(bodies).toEqual([
+      expect.objectContaining({ system: "system-a" }),
+      expect.objectContaining({ system: "system-b" }),
+    ]);
+  });
+
+  it("forwards the Cloud scope to the thread-list adapter", () => {
+    const cloud = {} as AssistantCloud;
+    mocks.useCloudThreadListAdapter.mockClear();
+    mocks.useChat.mockReturnValue({
+      resumeStream: vi.fn(),
+      status: "ready",
+    });
+
+    renderHook(() => useChatRuntime({ cloud, scopeId: "workspace-1" }));
+
+    expect(mocks.useCloudThreadListAdapter).toHaveBeenCalledWith({
+      cloud,
+      scopeId: "workspace-1",
+      sdk: AI_SDK_SDK,
+    });
+  });
+
   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(),
@@ -455,6 +519,48 @@ describe.skipIf(onReact18)("useChatRuntime", () => {
     await waitFor(() => expect(resumeStream).toHaveBeenCalledTimes(1));
   });
 
+  it("does not resume an unchanged stream id already handled by replacement storage", () => {
+    mocks.state.isLoadingHistory = true;
+    const resumeStream = vi.fn().mockResolvedValue(undefined);
+    mocks.useChat.mockReturnValue({
+      resumeStream,
+      status: "ready",
+    });
+    const storageA = {
+      getStreamId: (): string | null => "stream-1",
+      setStreamId: vi.fn(),
+      clear: vi.fn(),
+    };
+    const storageB = {
+      getStreamId: (): string | null => "stream-1",
+      setStreamId: vi.fn(),
+      clear: vi.fn(),
+    };
+    getResumedStreamIds(storageB).add("stream-1");
+    const transportA = {
+      getResumableAdapter: () => ({
+        storage: storageA,
+        resumeApi: "/api/chat/resume",
+      }),
+    };
+    const transportB = {
+      getResumableAdapter: () => ({
+        storage: storageB,
+        resumeApi: "/api/chat/resume",
+      }),
+    };
+
+    const { rerender } = renderHook(
+      ({ transport }) => useChatRuntime({ transport: transport as never }),
+      { initialProps: { transport: transportA } },
+    );
+
+    mocks.state.isLoadingHistory = false;
+    rerender({ transport: transportB });
+
+    expect(resumeStream).not.toHaveBeenCalled();
+  });
+
   it("does not clear a newer stream id when an older resume fails", async () => {
     let streamId: string | null = "stream-1";
     const storageSubscribers = new Set<() => void>();

package/src/runtime/useChatRuntime.ts

@@ -8,13 +8,28 @@ import {
   useRemoteThreadListRuntime,
 } from "@assistant-ui/core/react";
 import { useAui, useAuiState } from "@assistant-ui/store";
-import { useHostDestroySignal } from "@assistant-ui/store/internal";
+import type { ChatTransport } from "ai";
+import { useMemo } from "react";
+import { AssistantChatTransport } from "../transport/AssistantChatTransport";
 import { useChatThread, type ChatThreadOptions } from "./useChatThread";
+import { useDynamicChatTransport } from "./useDynamicChatTransport";
+import { useHostDestroySignal } from "@assistant-ui/store/internal";
 import { AI_SDK_SDK } from "./sdkIdentity";
 
 export type UseChatRuntimeOptions<UI_MESSAGE extends UIMessage = UIMessage> =
-  ChatThreadOptions<UI_MESSAGE> & {
+  Omit<ChatThreadOptions<UI_MESSAGE>, "transport"> & {
+    /**
+     * The transport threads send through. `AssistantChatTransport` instances
+     * are cloned per thread through `__internal_clone()` so their assistant-ui
+     * wiring remains isolated. Other transport instances are shared as-is.
+     */
+    transport?: ChatTransport<UI_MESSAGE> | undefined;
     cloud?: AssistantCloud | undefined;
+    /**
+     * Stable identity for the account or workspace owning Cloud runtime state.
+     * Provide it from the first render and change it when that scope changes.
+     */
+    scopeId?: string | undefined;
     onThreadIdChange?: ((threadId: string | undefined) => void) | undefined;
   };
 
@@ -39,14 +54,21 @@ const useChatThreadRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
 
 export const useChatRuntime = <UI_MESSAGE extends UIMessage = UIMessage>({
   cloud,
+  scopeId,
   onThreadIdChange,
   ...options
 }: UseChatRuntimeOptions<UI_MESSAGE> = {}): AssistantRuntime => {
   const hostDestroySignal = useHostDestroySignal();
-  const cloudAdapter = useCloudThreadListAdapter({ cloud, sdk: AI_SDK_SDK });
+  const cloudAdapter = useCloudThreadListAdapter({
+    cloud,
+    scopeId,
+    sdk: AI_SDK_SDK,
+  });
+  const fallback = useMemo(() => new AssistantChatTransport<UI_MESSAGE>(), []);
+  const transport = useDynamicChatTransport(options.transport ?? fallback);
   return useRemoteThreadListRuntime({
     runtimeHook: function RuntimeHook() {
-      return useChatThreadRuntime(options, hostDestroySignal);
+      return useChatThreadRuntime({ ...options, transport }, hostDestroySignal);
     },
     adapter: cloudAdapter,
     allowNesting: true,