@assistant-ui/ai-sdk 0.0.11 → 0.0.12

package/src/runtime/AISDKThreads.test.ts

@@ -55,7 +55,188 @@ const chatCallbacks = () => ({
   sendAutomaticallyWhen: vi.fn(() => false),
 });
 
+const thread = (id: string) => ({
+  id,
+  title: id,
+  is_archived: false,
+  last_message_at: null,
+  external_id: null,
+  metadata: null,
+});
+const makeCloud = (id: string) =>
+  ({
+    registerSdk: vi.fn(),
+    threads: {
+      list: vi.fn(async ({ is_archived }: { is_archived?: boolean } = {}) => ({
+        threads: is_archived ? [] : [thread(id)],
+      })),
+      get: vi.fn(async () => thread(id)),
+      messages: {
+        list: vi.fn(async () => ({ messages: [] })),
+        create: vi.fn(async () => ({ message_id: "stored" })),
+      },
+    },
+    runs: { report: vi.fn() },
+    telemetry: { enabled: false },
+  }) as unknown as AssistantCloud;
+
 describe("AISDKThreads", () => {
+  it.each(
+    (["client", "scope", "both", "rapid"] as const).flatMap((change) =>
+      (["uncontrolled", "changed", "unchanged"] as const).map((selection) => ({
+        change,
+        selection,
+      })),
+    ),
+  )(
+    "reloads Cloud threads when $change changes with selection=$selection",
+    async ({ change, selection }) => {
+      const cloudA = makeCloud("a-thread");
+      const cloudB = makeCloud("b-thread");
+      const intermediate = makeCloud("intermediate-thread");
+      const stream = createCancellableTransport();
+      let cloud = cloudA;
+      let scopeId = "workspace-a";
+      let threadId = selection === "uncontrolled" ? undefined : "a-thread";
+      const onThreadIdChange = vi.fn();
+      const { handle, rerender } = createLiveHandle(() => ({
+        cloud,
+        scopeId,
+        threadId,
+        onThreadIdChange,
+        transport: () => stream.transport,
+      }));
+      try {
+        await handle.getClient().threads.getLoadThreadsPromise();
+        await vi.waitFor(() =>
+          expect(handle.getClient().threads.getState().threadIds).toContain(
+            "a-thread",
+          ),
+        );
+        if (selection === "uncontrolled") {
+          flushTapSync(() =>
+            handle.getClient().threads.switchToThread("a-thread"),
+          );
+        }
+        await vi.waitFor(() =>
+          expect(handle.getClient().threads.getState().mainThreadId).toBe(
+            "a-thread",
+          ),
+        );
+        await vi.waitFor(() =>
+          expect(handle.getClient().thread.getState().isLoading).toBe(false),
+        );
+        flushTapSync(() => handle.getClient().composer.setText("A only"));
+        flushTapSync(() => handle.getClient().composer.send());
+        await vi.waitFor(() =>
+          expect(handle.getClient().thread.getState().isRunning).toBe(true),
+        );
+        stream.emit(
+          { type: "start" },
+          { type: "text-start", id: "partial" },
+          { type: "text-delta", id: "partial", delta: "A response" },
+        );
+        await vi.waitFor(() =>
+          expect(threadText(handle)).toContain("A response"),
+        );
+        const callsBeforeScopeChange = vi.mocked(cloudA.threads.list).mock.calls
+          .length;
+        rerender();
+        expect(cloudA.threads.list).toHaveBeenCalledTimes(
+          callsBeforeScopeChange,
+        );
+        const writesBeforeScopeChange = vi.mocked(
+          cloudA.threads.messages.create,
+        ).mock.calls.length;
+        const fetchesBeforeScopeChange = vi.mocked(cloudA.threads.get).mock
+          .calls.length;
+        if (change === "rapid") {
+          cloud = intermediate;
+          scopeId = "workspace-intermediate";
+          rerender();
+        }
+        if (change === "scope") {
+          vi.mocked(cloudA.threads.list).mockImplementation(
+            makeCloud("b-thread").threads.list,
+          );
+        } else {
+          cloud = cloudB;
+        }
+        if (change !== "client") scopeId = "workspace-b";
+        if (selection === "changed") threadId = "b-thread";
+        rerender();
+        rerender();
+        await vi.waitFor(() =>
+          expect(handle.getClient().threads.getState().threadIds).toContain(
+            "b-thread",
+          ),
+        );
+        expect(handle.getClient().threads.getState().threadIds).not.toContain(
+          "a-thread",
+        );
+        await vi.waitFor(() => expect(stream.getCancelCount()).toBe(1));
+        expect(cloudB.threads.messages.create).not.toHaveBeenCalled();
+        expect(cloudA.threads.messages.create).toHaveBeenCalledTimes(
+          writesBeforeScopeChange,
+        );
+        if (change === "rapid") {
+          expect(intermediate.threads.list).not.toHaveBeenCalled();
+          expect(handle.getClient().threads.getState().threadIds).not.toContain(
+            "intermediate-thread",
+          );
+        }
+        if (selection === "changed") {
+          await vi.waitFor(() =>
+            expect(handle.getClient().threads.getState().mainThreadId).toBe(
+              "b-thread",
+            ),
+          );
+          expect(onThreadIdChange).not.toHaveBeenCalled();
+        } else if (selection === "unchanged") {
+          expect(
+            handle.getClient().threads.item("main").getState().status,
+          ).toBe("new");
+          expect(onThreadIdChange).toHaveBeenCalledExactlyOnceWith(undefined);
+          const fetches = vi
+            .mocked(cloud.threads.get)
+            .mock.calls.slice(
+              change === "scope" ? fetchesBeforeScopeChange : 0,
+            );
+          expect(fetches).not.toContainEqual(["a-thread"]);
+        }
+      } finally {
+        handle.destroy();
+      }
+    },
+  );
+
+  it("loads once when entering Cloud mode", async () => {
+    let cloud: AssistantCloud | undefined;
+    const { handle, rerender } = createLiveHandle(() => ({ cloud }));
+    try {
+      cloud = makeCloud("a-thread");
+      rerender();
+      await vi.waitFor(() =>
+        expect(handle.getClient().threads.getState().threadIds).toContain(
+          "a-thread",
+        ),
+      );
+      expect(cloud.threads.list).toHaveBeenCalledTimes(2);
+      cloud = undefined;
+      rerender();
+      cloud = makeCloud("b-thread");
+      rerender();
+      await vi.waitFor(() =>
+        expect(handle.getClient().threads.getState().threadIds).toContain(
+          "b-thread",
+        ),
+      );
+      expect(cloud.threads.list).toHaveBeenCalledTimes(2);
+    } finally {
+      handle.destroy();
+    }
+  });
+
   it("emits the selected thread id when its composer sends", async () => {
     const { transport } = createControlledTransport();
     const handle = createAssistantClient(

package/src/runtime/AISDKThreads.ts

@@ -1,7 +1,7 @@
 "use client";
 
 import { resource, useResource, withKey } from "@assistant-ui/tap";
-import { useEffect, useMemo, useState } from "react";
+import { useEffect, useMemo, useRef, useState } from "react";
 import type { Chat, UIMessage } from "@ai-sdk/react";
 import type { ChatTransport } from "ai";
 import type { AssistantCloud } from "assistant-cloud";
@@ -49,6 +49,11 @@ export type AISDKThreadsOptions<UI_MESSAGE extends UIMessage = UIMessage> =
      * switch and stops on delete; per-thread history loads once per thread.
      */
     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;
     /**
      * Controlled thread id for the cloud list. Ignored without `cloud`.
      */
@@ -61,7 +66,7 @@ export type AISDKThreadsOptions<UI_MESSAGE extends UIMessage = UIMessage> =
 
 type AISDKThreadChatOptions<UI_MESSAGE extends UIMessage = UIMessage> = Omit<
   AISDKThreadsOptions<UI_MESSAGE>,
-  "cloud" | "threadId" | "onThreadIdChange"
+  "cloud" | "scopeId" | "threadId" | "onThreadIdChange"
 >;
 
 type ChatOptionsRef<UI_MESSAGE extends UIMessage> = {
@@ -182,7 +187,8 @@ const AISDKChatThread = resource(useAISDKChatThread);
 const useAISDKThreads = <UI_MESSAGE extends UIMessage = UIMessage>(
   options?: AISDKThreadsOptions<UI_MESSAGE>,
 ) => {
-  const { cloud, threadId, onThreadIdChange, ...threadOptions } = options ?? {};
+  const { cloud, scopeId, threadId, onThreadIdChange, ...threadOptions } =
+    options ?? {};
   const [chats] = useState(() => new Map<string, ChatEntry<UI_MESSAGE>>());
   const bindCloud = cloud !== undefined;
 
@@ -198,7 +204,28 @@ const useAISDKThreads = <UI_MESSAGE extends UIMessage = UIMessage>(
     }
   });
 
-  const cloudAdapter = useCloudThreadListAdapter({ cloud, sdk: AI_SDK_SDK });
+  const cloudAdapter = useCloudThreadListAdapter({
+    cloud,
+    scopeId,
+    sdk: AI_SDK_SDK,
+  });
+  const clientRef = useAssistantClientRef();
+  const previousCloudAdapter = useRef(bindCloud ? cloudAdapter : undefined);
+  useEffect(() => {
+    const previous = previousCloudAdapter.current;
+    previousCloudAdapter.current = bindCloud ? cloudAdapter : undefined;
+    if (!bindCloud || previous === undefined || previous === cloudAdapter)
+      return;
+    // Wait until RemoteThreadList commits the replacement adapter before
+    // requesting its explicit reload.
+    let cancelled = false;
+    queueMicrotask(() => {
+      if (!cancelled) void clientRef.current!.threads.reload();
+    });
+    return () => {
+      cancelled = true;
+    };
+  }, [bindCloud, clientRef, cloudAdapter]);
   const thread = (id: string) => {
     const element = AISDKChatThread({
       threadId: id,

package/src/runtime/DynamicChatTransport.test.ts

@@ -0,0 +1,203 @@
+import type { AssistantRuntime } from "@assistant-ui/core";
+import type { UIMessage } from "@ai-sdk/react";
+import { afterEach, describe, expect, it, vi } from "vitest";
+import { AssistantChatTransport } from "../transport/AssistantChatTransport";
+import { DynamicChatTransport } from "./DynamicChatTransport";
+
+const emptyStreamResponse = () =>
+  new Response(
+    new ReadableStream({ start: (controller) => controller.close() }),
+    {
+      headers: { "content-type": "text/event-stream" },
+    },
+  );
+
+const sendMessagesOptions = (chatId: string) => ({
+  trigger: "submit-message" as const,
+  chatId,
+  messageId: undefined,
+  messages: [
+    {
+      id: `message-${chatId}`,
+      role: "user" as const,
+      parts: [{ type: "text" as const, text: "hello" }],
+    },
+  ],
+  abortSignal: undefined,
+});
+
+const createRuntime = (system: string) =>
+  ({
+    thread: {
+      getModelContext: () => ({ system }),
+    },
+  }) as AssistantRuntime;
+
+describe("DynamicChatTransport", () => {
+  afterEach(() => {
+    vi.restoreAllMocks();
+  });
+
+  it("rejects sends without a registered thread context", () => {
+    const dynamicTransport = new DynamicChatTransport(
+      new AssistantChatTransport<UIMessage>(),
+    );
+
+    expect(() =>
+      dynamicTransport.sendMessages(sendMessagesOptions("missing-thread")),
+    ).toThrow(
+      'DynamicChatTransport has no registered context for chat "missing-thread"',
+    );
+  });
+
+  it("keeps AssistantChatTransport wiring scoped per thread", async () => {
+    const createTransport = () => {
+      const bodies: Array<{ id: string; system: string }> = [];
+      const fetch = vi.fn(async (_input, init) => {
+        bodies.push(JSON.parse(String(init?.body)));
+        return emptyStreamResponse();
+      });
+      return {
+        bodies,
+        fetch,
+        transport: new AssistantChatTransport<UIMessage>({
+          fetch,
+        }),
+      };
+    };
+    const initial = createTransport();
+    const dynamicTransport = new DynamicChatTransport(initial.transport);
+    const ownerA = {};
+    const ownerB = {};
+
+    dynamicTransport.setThreadContext(
+      "thread-a",
+      ownerA,
+      createRuntime("system-a"),
+      () => ({
+        initialize: async () => ({
+          remoteId: "remote-a",
+          externalId: undefined,
+        }),
+      }),
+    );
+    dynamicTransport.setThreadContext(
+      "thread-b",
+      ownerB,
+      createRuntime("system-b"),
+      () => ({
+        initialize: async () => ({
+          remoteId: "remote-b",
+          externalId: undefined,
+        }),
+      }),
+    );
+
+    await dynamicTransport.sendMessages(sendMessagesOptions("thread-a"));
+    await dynamicTransport.sendMessages(sendMessagesOptions("thread-b"));
+
+    const replacement = createTransport();
+    dynamicTransport.setTransport(replacement.transport);
+    await dynamicTransport.sendMessages(sendMessagesOptions("thread-a"));
+    await dynamicTransport.sendMessages(sendMessagesOptions("thread-b"));
+
+    expect(initial.fetch).toHaveBeenCalledTimes(2);
+    expect(replacement.fetch).toHaveBeenCalledTimes(2);
+    expect(initial.bodies).toEqual([
+      expect.objectContaining({ id: "remote-a", system: "system-a" }),
+      expect.objectContaining({ id: "remote-b", system: "system-b" }),
+    ]);
+    expect(replacement.bodies).toEqual([
+      expect.objectContaining({ id: "remote-a", system: "system-a" }),
+      expect.objectContaining({ id: "remote-b", system: "system-b" }),
+    ]);
+    expect(dynamicTransport.getCurrentTransport("thread-a")).not.toBe(
+      dynamicTransport.getCurrentTransport("thread-b"),
+    );
+  });
+
+  it("defers replacement clones until a thread uses the transport", async () => {
+    const clone = vi.spyOn(
+      AssistantChatTransport.prototype,
+      "__internal_clone",
+    );
+    const dynamicTransport = new DynamicChatTransport(
+      new AssistantChatTransport<UIMessage>(),
+    );
+    const owner = {};
+    dynamicTransport.setThreadContext(
+      "thread-a",
+      owner,
+      createRuntime("system-a"),
+      () => ({
+        initialize: async () => ({
+          remoteId: "remote-a",
+          externalId: undefined,
+        }),
+      }),
+    );
+    clone.mockClear();
+
+    for (let index = 0; index < 10; index++) {
+      dynamicTransport.setTransport(
+        new AssistantChatTransport<UIMessage>({
+          fetch: async () => emptyStreamResponse(),
+        }),
+      );
+    }
+    dynamicTransport.setThreadContext(
+      "thread-a",
+      owner,
+      createRuntime("updated-system-a"),
+      () => ({
+        initialize: async () => ({
+          remoteId: "updated-remote-a",
+          externalId: undefined,
+        }),
+      }),
+    );
+
+    expect(clone).not.toHaveBeenCalled();
+
+    await dynamicTransport.sendMessages(sendMessagesOptions("thread-a"));
+
+    expect(clone).toHaveBeenCalledOnce();
+  });
+
+  it("retries a replacement clone after cloning throws", async () => {
+    const initialFetch = vi.fn(async () => emptyStreamResponse());
+    const dynamicTransport = new DynamicChatTransport(
+      new AssistantChatTransport<UIMessage>({ fetch: initialFetch }),
+    );
+    dynamicTransport.setThreadContext(
+      "thread-a",
+      {},
+      createRuntime("system-a"),
+      () => ({
+        initialize: async () => ({
+          remoteId: "remote-a",
+          externalId: undefined,
+        }),
+      }),
+    );
+    await dynamicTransport.sendMessages(sendMessagesOptions("thread-a"));
+
+    const replacement = new AssistantChatTransport<UIMessage>();
+    const cloneError = new Error("clone failed");
+    const clone = vi
+      .spyOn(replacement, "__internal_clone")
+      .mockImplementation(() => {
+        throw cloneError;
+      });
+    dynamicTransport.setTransport(replacement);
+
+    expect(() =>
+      dynamicTransport.sendMessages(sendMessagesOptions("thread-a")),
+    ).toThrow(cloneError);
+    expect(() =>
+      dynamicTransport.sendMessages(sendMessagesOptions("thread-a")),
+    ).toThrow(cloneError);
+    expect(clone).toHaveBeenCalledTimes(2);
+    expect(initialFetch).toHaveBeenCalledOnce();
+  });
+});