@assistant-ui/ai-sdk 0.0.11 → 0.0.12

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

@@ -0,0 +1,143 @@
+// @vitest-environment jsdom
+
+import { Suspense, version } from "react";
+import { render } from "@testing-library/react";
+import type { AssistantRuntime } from "@assistant-ui/core";
+import type { ChatTransport, UIMessage } from "ai";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import { AssistantChatTransport } from "../transport/AssistantChatTransport";
+import { DynamicChatTransport } from "./DynamicChatTransport";
+
+const mocks = vi.hoisted(() => ({
+  system: "",
+  useAISDKRuntime: vi.fn(),
+  useChat: vi.fn(),
+}));
+
+vi.mock("@ai-sdk/react", async (importOriginal) => ({
+  ...(await importOriginal<typeof import("@ai-sdk/react")>()),
+  Chat: class MockChat {
+    constructor(config: unknown) {
+      Object.assign(this, config);
+    }
+  },
+  useChat: mocks.useChat,
+}));
+
+vi.mock("./useAISDKRuntime", async (importOriginal) => ({
+  ...(await importOriginal<typeof import("./useAISDKRuntime")>()),
+  useAISDKRuntime: mocks.useAISDKRuntime,
+}));
+
+import { useChatThread } from "./useChatThread";
+
+const onReact18 = version.startsWith("18.");
+
+const createRuntime = (system: string) =>
+  ({
+    thread: {
+      getModelContext: () => ({ system }),
+      getState: () => ({ isLoading: false }),
+      subscribe: () => () => {},
+    },
+  }) as unknown as AssistantRuntime;
+
+describe("useChatThread transport binding", () => {
+  beforeEach(() => {
+    mocks.system = "";
+    mocks.useAISDKRuntime.mockReset();
+    mocks.useChat.mockReset();
+    mocks.useAISDKRuntime.mockImplementation(() => createRuntime(mocks.system));
+    mocks.useChat.mockImplementation(({ chat }) => ({
+      ...chat,
+      messages: [],
+      status: "ready",
+      error: undefined,
+      stop: vi.fn(async () => {}),
+    }));
+  });
+
+  it.skipIf(onReact18)(
+    "keeps fallback wiring on committed values after an abandoned render",
+    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 pending = new Promise<never>(() => {});
+      let committedProxy: ChatTransport<UIMessage> | undefined;
+
+      const App = ({
+        system,
+        remoteId,
+        suspend = false,
+      }: {
+        system: string;
+        remoteId: string;
+        suspend?: boolean;
+      }) => {
+        mocks.system = system;
+        useChatThread(
+          { transport },
+          {
+            id: "thread-id",
+            isMainThread: true,
+            getThreadListItem: () => ({
+              initialize: async () => ({
+                remoteId,
+                externalId: undefined,
+              }),
+            }),
+          },
+        );
+        const proxy = mocks.useChat.mock.lastCall?.[0].chat.transport as
+          | ChatTransport<UIMessage>
+          | undefined;
+        if (!suspend) committedProxy = proxy;
+        if (suspend) throw pending;
+        return null;
+      };
+
+      const view = render(
+        <Suspense fallback={null}>
+          <App system="committed-system" remoteId="committed-remote" />
+        </Suspense>,
+      );
+
+      transport.registerThread("thread-id", {});
+      view.rerender(
+        <Suspense fallback={null}>
+          <App system="discarded-system" remoteId="discarded-remote" suspend />
+        </Suspense>,
+      );
+
+      await committedProxy!.sendMessages({
+        trigger: "submit-message",
+        chatId: "thread-id",
+        messageId: undefined,
+        messages: [
+          {
+            id: "message-id",
+            role: "user",
+            parts: [{ type: "text", text: "hello" }],
+          },
+        ],
+        abortSignal: undefined,
+      });
+
+      expect(bodies).toEqual([
+        expect.objectContaining({
+          id: "committed-remote",
+          system: "committed-system",
+        }),
+      ]);
+    },
+  );
+});

package/src/runtime/useChatThread.ts

@@ -12,24 +12,28 @@ import {
   type AISDKRuntimeAdapter,
   type CustomToCreateMessageFunction,
 } from "./useAISDKRuntime";
-import type { ChatInit, ChatTransport } from "ai";
+import type { ChatInit } from "ai";
 import {
   AssistantChatTransport,
   type InitializableThreadListItem,
 } from "../transport/AssistantChatTransport";
-import type {
-  AssistantChatResumableOptions,
-  ResumableClientStorage,
-} from "../transport/resumable";
 import {
   useCallback,
   useEffect,
+  useEffectEvent,
+  useInsertionEffect,
   useMemo,
   useRef,
   useState,
   useSyncExternalStore,
 } from "react";
 import { useResourceCleanup } from "./useResourceCleanup";
+import {
+  DynamicChatTransport,
+  getResumedStreamIds,
+} from "./DynamicChatTransport";
+import { getResumableAdapter } from "./getResumableAdapter";
+import { useDynamicChatTransport } from "./useDynamicChatTransport";
 
 export type ChatThreadOptions<UI_MESSAGE extends UIMessage = UIMessage> =
   ChatInit<UI_MESSAGE> &
@@ -49,6 +53,7 @@ export type ChatThreadOptions<UI_MESSAGE extends UIMessage = UIMessage> =
       onResumeError?: ((error: unknown) => void) | undefined;
       joinStrategy?: AISDKRuntimeAdapter["joinStrategy"];
       messageRepository?: AISDKRuntimeAdapter<UI_MESSAGE>["messageRepository"];
+      /** @deprecated Experimental since 2026-06-23. Not scheduled for removal; the API may change in any release. */
       unstable_onBranchChange?: AISDKRuntimeAdapter["unstable_onBranchChange"];
     };
 
@@ -77,58 +82,13 @@ export type ChatThreadEnvironment<UI_MESSAGE extends UIMessage = UIMessage> = {
   messageRepositoryInstance?: MessageRepository | undefined;
 };
 
-const useDynamicChatTransport = <UI_MESSAGE extends UIMessage = UIMessage>(
-  transport: ChatTransport<UI_MESSAGE>,
-): ChatTransport<UI_MESSAGE> => {
-  const transportRef = useRef<ChatTransport<UI_MESSAGE>>(transport);
-  useEffect(() => {
-    transportRef.current = transport;
-  });
-  const dynamicTransport = useMemo(
-    () =>
-      new Proxy(transportRef.current, {
-        get(_, prop) {
-          const res =
-            transportRef.current[prop as keyof ChatTransport<UI_MESSAGE>];
-          return typeof res === "function"
-            ? res.bind(transportRef.current)
-            : res;
-        },
-      }),
-    [],
-  );
-  return dynamicTransport;
-};
-
-const getResumableAdapter = <UI_MESSAGE extends UIMessage>(
-  transport: ChatTransport<UI_MESSAGE>,
-): AssistantChatResumableOptions | undefined => {
-  if (transport instanceof AssistantChatTransport) {
-    return transport.getResumableAdapter();
-  }
-  const candidate = (transport as { getResumableAdapter?: () => unknown })
-    .getResumableAdapter;
-  if (typeof candidate !== "function") return undefined;
-  return candidate.call(transport) as AssistantChatResumableOptions | undefined;
+type ChatThreadTransportBinding = {
+  runtime: AssistantRuntime;
+  getThreadListItem: () => InitializableThreadListItem | undefined;
 };
 
 const getNoPendingStreamId = () => null;
 
-const resumedStreamIdsByStorage = new WeakMap<
-  ResumableClientStorage,
-  Set<string>
->();
-
-const getResumedStreamIds = (storage: ResumableClientStorage | undefined) => {
-  if (!storage) return new Set<string>();
-  let resumedStreamIds = resumedStreamIdsByStorage.get(storage);
-  if (!resumedStreamIds) {
-    resumedStreamIds = new Set();
-    resumedStreamIdsByStorage.set(storage, resumedStreamIds);
-  }
-  return resumedStreamIds;
-};
-
 /**
  * Splits the combined options into the assistant-ui side and the `ChatInit`
  * remainder the AI SDK consumes, so external `Chat` construction forwards the
@@ -246,21 +206,47 @@ export const useChatThread = <UI_MESSAGE extends UIMessage = UIMessage>(
     messageRepositoryInstance,
   } = env;
 
-  // 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(
+  const defaultTransport = useMemo(() => new AssistantChatTransport(), []);
+  const configuredTransport = transportOptions ?? defaultTransport;
+  const transport = useDynamicChatTransport(
+    configuredTransport,
+    externalChat === undefined,
+  );
+  const transportContextOwner = useMemo(() => ({}), []);
+  const getThreadListItemRef = useRef(getThreadListItem);
+  useInsertionEffect(() => {
+    getThreadListItemRef.current = getThreadListItem;
+  }, [getThreadListItem]);
+  const getCurrentThreadListItem = useCallback(
+    () => getThreadListItemRef.current(),
+    [],
+  );
+  const resumableStorage = useMemo(
     () =>
-      transportOptions === undefined
-        ? new AssistantChatTransport()
-        : externalChat === undefined &&
-            transportOptions instanceof AssistantChatTransport
-          ? transportOptions.__internal_clone()
-          : transportOptions,
-    [transportOptions, externalChat],
+      transport instanceof DynamicChatTransport
+        ? transport.getCurrentResumableStorage()
+        : getResumableAdapter(transport)?.storage,
+    [transport],
+  );
+
+  const transportBindingRef = useRef<ChatThreadTransportBinding | null>(null);
+  // This is null when useMemo runs, so the proxy pins its mount render's binding. The fallback
+  // serves only pre-commit sends; after insertion effects, transportBindingRef is authoritative.
+  let initialTransportBinding: ChatThreadTransportBinding | null = null;
+  const chatTransport = useMemo(
+    () =>
+      transport instanceof DynamicChatTransport
+        ? transport.createThreadProxy(transportContextOwner, () => {
+            const binding =
+              transportBindingRef.current ?? initialTransportBinding;
+            if (!binding) {
+              throw new Error("Chat transport used before runtime setup");
+            }
+            return binding;
+          })
+        : transport,
+    [initialTransportBinding, transport, transportContextOwner],
   );
-  const transport = useDynamicChatTransport(sourceTransport);
 
   const latestChatOptionsRef = useRef(chatOptions);
   useEffect(() => {
@@ -271,7 +257,10 @@ export const useChatThread = <UI_MESSAGE extends UIMessage = UIMessage>(
   const [ownedChat] = useState(
     () =>
       externalChat ??
-      createChat({ ...chatOptions, id, transport }, latestChatOptionsRef),
+      createChat(
+        { ...chatOptions, id, transport: chatTransport },
+        latestChatOptionsRef,
+      ),
   );
 
   const chat = useChat({
@@ -309,15 +298,42 @@ export const useChatThread = <UI_MESSAGE extends UIMessage = UIMessage>(
     unstable_hostApprovalOwner: externalChat ?? ownedChat,
     ...(unstable_onBranchChange && { unstable_onBranchChange }),
   });
+  initialTransportBinding = {
+    runtime,
+    getThreadListItem: getCurrentThreadListItem,
+  };
+  useInsertionEffect(() => {
+    transportBindingRef.current = {
+      runtime,
+      getThreadListItem: getCurrentThreadListItem,
+    };
+  }, [runtime, getCurrentThreadListItem]);
+
+  const registerTransportContext = useEffectEvent(
+    (dynamicTransport: DynamicChatTransport<UI_MESSAGE>, chatId: string) => {
+      dynamicTransport.setThreadContext(
+        chatId,
+        transportContextOwner,
+        runtime,
+        getCurrentThreadListItem,
+      );
+      return () =>
+        dynamicTransport.unregisterThread(chatId, transportContextOwner);
+    },
+  );
+  useInsertionEffect(() => {
+    if (!(transport instanceof DynamicChatTransport)) return undefined;
+    return registerTransportContext(transport, id);
+  }, [id, runtime, transport, transportContextOwner]);
 
-  // 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);
+  if (
+    !(transport instanceof DynamicChatTransport) &&
+    configuredTransport instanceof AssistantChatTransport
+  ) {
+    configuredTransport.setRuntime(runtime);
+    configuredTransport.__internal_setGetThreadListItem(
+      getCurrentThreadListItem,
+    );
   }
 
   const subscribeToRuntime = useCallback(
@@ -334,10 +350,6 @@ export const useChatThread = <UI_MESSAGE extends UIMessage = UIMessage>(
     getHistoryLoadingSnapshot,
   );
 
-  const resumableStorage = useMemo(
-    () => getResumableAdapter(sourceTransport)?.storage,
-    [sourceTransport],
-  );
   const subscribeToResumableStorage = useCallback(
     (callback: () => void) =>
       isMainThread
@@ -357,7 +369,7 @@ export const useChatThread = <UI_MESSAGE extends UIMessage = UIMessage>(
   const isChatRunning =
     chat.status === "submitted" || chat.status === "streaming";
 
-  const resumedStreamIds = useMemo(
+  const staticResumedStreamIds = useMemo(
     () => getResumedStreamIds(resumableStorage),
     [resumableStorage],
   );
@@ -366,6 +378,10 @@ export const useChatThread = <UI_MESSAGE extends UIMessage = UIMessage>(
     onResumeErrorRef.current = onResumeError;
   });
   useEffect(() => {
+    const resumedStreamIds =
+      transport instanceof DynamicChatTransport
+        ? transport.getResumedStreamIds()
+        : staticResumedStreamIds;
     if (!pendingStreamId || resumedStreamIds.has(pendingStreamId)) {
       return;
     }
@@ -411,7 +427,8 @@ export const useChatThread = <UI_MESSAGE extends UIMessage = UIMessage>(
     pendingStreamId,
     ownedChat,
     resumableStorage,
-    resumedStreamIds,
+    staticResumedStreamIds,
+    transport,
   ]);
 
   return runtime;

package/src/runtime/useDynamicChatTransport.ts

@@ -0,0 +1,26 @@
+import type { UIMessage } from "@ai-sdk/react";
+import type { ChatTransport } from "ai";
+import { useEffect, useInsertionEffect, useState } from "react";
+import { DynamicChatTransport } from "./DynamicChatTransport";
+
+export const useDynamicChatTransport = <UI_MESSAGE extends UIMessage>(
+  transport: ChatTransport<UI_MESSAGE>,
+  enabled = true,
+): ChatTransport<UI_MESSAGE> => {
+  const [dynamicTransport] = useState(() =>
+    transport instanceof DynamicChatTransport
+      ? transport
+      : new DynamicChatTransport(transport),
+  );
+
+  useInsertionEffect(() => {
+    if (!enabled || dynamicTransport === transport) return;
+    dynamicTransport.setTransport(transport);
+  }, [dynamicTransport, enabled, transport]);
+  useEffect(() => {
+    if (!enabled || dynamicTransport === transport) return;
+    dynamicTransport.flushTransportChange();
+  }, [dynamicTransport, enabled, transport]);
+
+  return enabled ? dynamicTransport : transport;
+};

package/src/runtime/useExternalHistory.test.ts

@@ -206,6 +206,53 @@ describe("useExternalHistory withFormat contract", () => {
     await waitFor(() => expect(load).toHaveBeenCalledTimes(1));
   });
 
+  it("re-enters loading when remoteId arrives before history settles", async () => {
+    mocks.hasThreadListItem = true;
+    let resolveLoad!: (repo: MessageFormatRepository<unknown>) => void;
+    const load = vi.fn(
+      () =>
+        new Promise<MessageFormatRepository<unknown>>((resolve) => {
+          resolveLoad = resolve;
+        }),
+    );
+    const adapter: ThreadHistoryAdapter = {
+      load: vi.fn(),
+      append: vi.fn(),
+      withFormat: vi.fn().mockReturnValue({
+        load,
+        append: vi.fn().mockResolvedValue(undefined),
+      }),
+    };
+
+    const { result } = renderHook(() =>
+      useExternalHistory(
+        runtimeRef,
+        adapter,
+        toThreadMessages,
+        storageFormat,
+        onSetMessages,
+      ),
+    );
+
+    await act(async () => {});
+    expect(result.current.isLoading).toBe(false);
+    expect(load).not.toHaveBeenCalled();
+
+    mocks.remoteId = "remote-thread";
+    await act(async () => {
+      for (const listener of mocks.listeners) listener();
+    });
+
+    await waitFor(() => expect(load).toHaveBeenCalledTimes(1));
+    expect(result.current.isLoading).toBe(true);
+
+    await act(async () => {
+      resolveLoad({ headId: null, messages: [] });
+    });
+
+    await waitFor(() => expect(result.current.isLoading).toBe(false));
+  });
+
   it("does not load history when remoteId appears during an active run", async () => {
     mocks.hasThreadListItem = true;
     const load = vi.fn().mockResolvedValue({ headId: null, messages: [] });
@@ -288,6 +335,70 @@ describe("useExternalHistory withFormat contract", () => {
 
     await waitFor(() => expect(result.current.isLoading).toBe(false));
   });
+
+  it("does not replace live messages when history resolves after a send starts", async () => {
+    mocks.hasThreadListItem = true;
+    mocks.remoteId = "remote-thread";
+    let resolveLoad!: (repo: MessageFormatRepository<unknown>) => void;
+    const load = vi.fn(
+      () =>
+        new Promise<MessageFormatRepository<unknown>>((resolve) => {
+          resolveLoad = resolve;
+        }),
+    );
+    const threadState = {
+      isRunning: false,
+      messages: [] as ThreadMessage[],
+    };
+    const threadImport = vi.fn();
+    const setMessages = vi.fn();
+    const thread = {
+      subscribe: () => () => {},
+      getState: () => threadState,
+      import: threadImport,
+      export: () => ({ headId: null, messages: [] }),
+    } as unknown as AssistantRuntime["thread"];
+    const historyRuntimeRef = {
+      current: { thread } as AssistantRuntime,
+    };
+    const adapter: ThreadHistoryAdapter = {
+      load: vi.fn(),
+      append: vi.fn(),
+      withFormat: vi.fn().mockReturnValue({
+        load,
+        append: vi.fn().mockResolvedValue(undefined),
+      }),
+    };
+
+    const { result } = renderHook(() =>
+      useExternalHistory(
+        historyRuntimeRef,
+        adapter,
+        (messages) => messages.map((message) => message as ThreadMessage),
+        storageFormat,
+        setMessages,
+      ),
+    );
+
+    await waitFor(() => expect(load).toHaveBeenCalledTimes(1));
+    const liveMessages = [
+      { id: "live-user" },
+      { id: "live-assistant" },
+    ] as ThreadMessage[];
+    await act(async () => {
+      threadState.isRunning = true;
+      threadState.messages = liveMessages;
+      resolveLoad({
+        headId: "stored",
+        messages: [{ parentId: null, message: { id: "stored" } }],
+      });
+    });
+
+    await waitFor(() => expect(result.current.isLoading).toBe(false));
+    expect(threadState.messages).toBe(liveMessages);
+    expect(threadImport).not.toHaveBeenCalled();
+    expect(setMessages).not.toHaveBeenCalled();
+  });
 });
 
 describe("toExportedMessageRepository", () => {