@assistant-ui/ai-sdk 0.0.1

package/src/runtime/AISDKChat.test.ts

@@ -0,0 +1,193 @@
+import { afterEach, describe, expect, it, vi } from "vitest";
+import { flushTapSync } from "@assistant-ui/tap";
+import { AuiConfig, createAssistantClient } from "@assistant-ui/store/client";
+import { AISDKChat } from "./AISDKChat";
+import {
+  createCancellableTransport,
+  createControlledTransport,
+} from "./__tests__/controlled-transport";
+
+describe("AISDKChat as a standalone client config entry", () => {
+  afterEach(() => {
+    vi.unstubAllGlobals();
+  });
+
+  it("streams a chat round trip without React", async () => {
+    const { transport, emit, close } = createControlledTransport();
+    const handle = createAssistantClient(
+      AuiConfig({ threads: AISDKChat({ transport }) }),
+    );
+    try {
+      handle.subscribe(() => {});
+      const aui = handle.getClient();
+
+      expect(aui.thread.getState().messages).toHaveLength(0);
+
+      flushTapSync(() => aui.composer.setText("hi"));
+      flushTapSync(() => aui.composer.send());
+
+      await vi.waitFor(() => {
+        const state = aui.thread.getState();
+        expect(state.isRunning).toBe(true);
+        expect(state.messages[0]).toMatchObject({ role: "user" });
+      });
+
+      emit(
+        { type: "start" },
+        { type: "text-start", id: "t1" },
+        { type: "text-delta", id: "t1", delta: "hello " },
+      );
+      await vi.waitFor(() => {
+        const last = aui.thread.getState().messages.at(-1);
+        expect(last?.role).toBe("assistant");
+        expect(last?.content).toContainEqual(
+          expect.objectContaining({ type: "text", text: "hello " }),
+        );
+      });
+
+      emit(
+        { type: "text-delta", id: "t1", delta: "world" },
+        { type: "text-end", id: "t1" },
+        { type: "finish" },
+      );
+      close();
+
+      await vi.waitFor(() => {
+        const state = aui.thread.getState();
+        expect(state.isRunning).toBe(false);
+        expect(state.messages).toHaveLength(2);
+        expect(state.messages.at(-1)?.content).toContainEqual(
+          expect.objectContaining({ type: "text", text: "hello world" }),
+        );
+      });
+    } finally {
+      handle.destroy();
+    }
+  });
+
+  it("stops an in-flight chat when its client is destroyed", async () => {
+    const { transport, getCancelCount } = createCancellableTransport();
+    const handle = createAssistantClient(
+      AuiConfig({ threads: AISDKChat({ transport }) }),
+    );
+    handle.subscribe(() => {});
+    const aui = handle.getClient();
+
+    try {
+      flushTapSync(() => aui.composer.setText("stop me"));
+      flushTapSync(() => aui.composer.send());
+      await vi.waitFor(() => {
+        expect(aui.thread.getState().isRunning).toBe(true);
+      });
+    } finally {
+      handle.destroy();
+    }
+
+    await vi.waitFor(() => {
+      expect(getCancelCount()).toBe(1);
+    });
+  });
+
+  it("stops a soft-unmounted chat on later client destruction", async () => {
+    const { transport, getCancelCount } = createCancellableTransport();
+    const handle = createAssistantClient(
+      AuiConfig({ threads: AISDKChat({ transport }) }),
+    );
+    const release = handle.subscribe(() => {});
+    const aui = handle.getClient();
+
+    flushTapSync(() => aui.composer.setText("keep streaming"));
+    flushTapSync(() => aui.composer.send());
+    await vi.waitFor(() => {
+      expect(aui.thread.getState().isRunning).toBe(true);
+    });
+
+    release();
+    await new Promise((resolve) => setTimeout(resolve, 0));
+    expect(getCancelCount()).toBe(0);
+
+    handle.destroy();
+    await vi.waitFor(() => {
+      expect(getCancelCount()).toBe(1);
+    });
+  });
+
+  it("installs the RuntimeAdapter scope defaults", () => {
+    const { transport } = createControlledTransport();
+    const handle = createAssistantClient(
+      AuiConfig({ threads: AISDKChat({ transport }) }),
+    );
+    try {
+      handle.subscribe(() => {});
+      const aui = handle.getClient();
+
+      expect(aui.threads.getState().mainThreadId).toBeDefined();
+      expect(aui.tools.getState()).toBeDefined();
+      expect(aui.dataRenderers.getState()).toBeDefined();
+      expect(aui.thread.getState().isRunning).toBe(false);
+    } finally {
+      handle.destroy();
+    }
+  });
+
+  it("mounts with the default transport when no options are given", () => {
+    const handle = createAssistantClient(AuiConfig({ threads: AISDKChat({}) }));
+    try {
+      handle.subscribe(() => {});
+      const aui = handle.getClient();
+
+      expect(aui.thread.getState().messages).toHaveLength(0);
+      expect(aui.thread.getState().isRunning).toBe(false);
+    } finally {
+      handle.destroy();
+    }
+  });
+
+  it("sends through the default AssistantChatTransport with the generated id", async () => {
+    const sse = [
+      { type: "start" },
+      { type: "text-start", id: "t1" },
+      { type: "text-delta", id: "t1", delta: "ok" },
+      { type: "text-end", id: "t1" },
+      { type: "finish" },
+    ]
+      .map((chunk) => `data: ${JSON.stringify(chunk)}\n\n`)
+      .join("");
+    const fetchMock = vi.fn(
+      async () =>
+        new Response(sse, {
+          headers: { "content-type": "text/event-stream" },
+        }),
+    );
+    vi.stubGlobal("fetch", fetchMock);
+
+    const handle = createAssistantClient(
+      AuiConfig({ threads: AISDKChat({ id: "test-thread-1" }) }),
+    );
+    try {
+      handle.subscribe(() => {});
+      const aui = handle.getClient();
+
+      flushTapSync(() => aui.composer.setText("hi"));
+      flushTapSync(() => aui.composer.send());
+
+      await vi.waitFor(() => {
+        expect(fetchMock).toHaveBeenCalledTimes(1);
+        const state = aui.thread.getState();
+        expect(state.isRunning).toBe(false);
+        expect(state.messages).toHaveLength(2);
+      });
+
+      const [url, init] = fetchMock.mock.calls[0]! as [
+        RequestInfo,
+        RequestInit,
+      ];
+      expect(String(url)).toContain("/api/chat");
+      const body = JSON.parse(init.body as string);
+      expect(body.id).toBe("test-thread-1");
+      expect(body.messages).toHaveLength(1);
+    } finally {
+      handle.destroy();
+    }
+  });
+});

package/src/runtime/AISDKChat.ts

@@ -0,0 +1,47 @@
+"use client";
+
+import { resource, useResource } from "@assistant-ui/tap";
+import { useState } from "react";
+import { generateId } from "ai";
+import type { UIMessage } from "@ai-sdk/react";
+import {
+  RuntimeAdapter,
+  runtimeAdapterTransformScopes,
+} from "@assistant-ui/core/store";
+import { attachTransformScopes } from "@assistant-ui/store/client";
+import { useChatThread, type ChatThreadOptions } from "./useChatThread";
+
+export type AISDKChatOptions<UI_MESSAGE extends UIMessage = UIMessage> =
+  ChatThreadOptions<UI_MESSAGE>;
+
+const useAISDKChat = <UI_MESSAGE extends UIMessage = UIMessage>(
+  options?: AISDKChatOptions<UI_MESSAGE>,
+) => {
+  const [id] = useState(() => options?.id ?? generateId());
+  // The transport resolves the request id from the thread list item, falling
+  // back to the runtime's main item, whose id here is the external store's
+  // placeholder constant. The single thread of this entry is the chat itself,
+  // so the handed-over item initializes to the chat id.
+  const [threadListItem] = useState(() => ({
+    initialize: async () => ({ remoteId: id, externalId: undefined }),
+  }));
+  const runtime = useChatThread(options, {
+    id,
+    isMainThread: true,
+    getThreadListItem: () => threadListItem,
+    stopOnClientDestroy: true,
+  });
+  return useResource(RuntimeAdapter(runtime));
+};
+
+/**
+ * `AuiConfig` entry that runs the AI SDK chat as the `threads` scope. Hosts the
+ * same orchestration as `useChatRuntime` inside the client's own resource tree,
+ * so it works with any `AssistantClient` host, React or not. Single thread; the
+ * multi-thread and assistant-cloud surface is {@link AISDKThreads}. The chat
+ * id is captured when the entry mounts, so a later `id` change in the options
+ * has no effect.
+ */
+export const AISDKChat = resource(useAISDKChat);
+
+attachTransformScopes(useAISDKChat, runtimeAdapterTransformScopes);

package/src/runtime/AISDKThreads.cloud.test.ts

@@ -0,0 +1,136 @@
+// @vitest-environment jsdom
+
+import { describe, expect, it, vi } from "vitest";
+import { flushTapSync } from "@assistant-ui/tap";
+import { AuiConfig, createAssistantClient } from "@assistant-ui/store/client";
+import type { AssistantCloud } from "assistant-cloud";
+import type { RemoteThreadListAdapter } from "@assistant-ui/core";
+import type { ThreadHistoryAdapter } from "@assistant-ui/core";
+
+const load = vi.hoisted(() => vi.fn(async () => ({ messages: [] })));
+const append = vi.hoisted(() => vi.fn(async () => {}));
+
+const mocks = vi.hoisted(() => {
+  const history: ThreadHistoryAdapter = {
+    load: async () => ({ messages: [] }),
+    append: async () => {},
+    withFormat: () => ({
+      load,
+      append,
+      delete: async () => {},
+      reportTelemetry: () => {},
+    }),
+  };
+  const adapter: RemoteThreadListAdapter = {
+    list: vi.fn(async () => ({
+      threads: [
+        { status: "regular" as const, remoteId: "t1", title: "One" },
+        { status: "regular" as const, remoteId: "t2", title: "Two" },
+      ],
+    })),
+    initialize: vi.fn(async (threadId: string) => ({
+      remoteId: threadId,
+      externalId: undefined,
+    })),
+    rename: vi.fn(async () => {}),
+    archive: vi.fn(async () => {}),
+    unarchive: vi.fn(async () => {}),
+    delete: vi.fn(async () => {}),
+    generateTitle: vi.fn(
+      async () =>
+        new ReadableStream({
+          start(controller) {
+            controller.close();
+          },
+        }) as never,
+    ),
+    fetch: vi.fn(async (id: string) => ({
+      status: "regular" as const,
+      remoteId: id,
+      externalId: undefined,
+      title: id,
+    })),
+    unstable_useAdapters: function useAdapters() {
+      return { history };
+    },
+  };
+  return { adapter };
+});
+
+vi.mock("@assistant-ui/core/react", async (importOriginal) => ({
+  ...(await importOriginal<typeof import("@assistant-ui/core/react")>()),
+  useCloudThreadListAdapter: () => mocks.adapter,
+}));
+
+import { AISDKThreads } from "./AISDKThreads";
+import { createCancellableTransport } from "./__tests__/controlled-transport";
+
+describe("AISDKThreads cloud", () => {
+  it("reloads history when switching a keyed cloud thread", async () => {
+    const handle = createAssistantClient(
+      AuiConfig({
+        threads: AISDKThreads({
+          cloud: {} as AssistantCloud,
+          threadId: "t1",
+        }),
+      }),
+    );
+    handle.subscribe(() => {});
+    const aui = handle.getClient();
+    await aui.threads.getLoadThreadsPromise();
+    await vi.waitFor(() => {
+      expect(handle.getClient().threads.getState().mainThreadId).toBe("t1");
+    });
+    await vi.waitFor(() => {
+      expect(load).toHaveBeenCalled();
+    });
+    const afterFirst = load.mock.calls.length;
+    flushTapSync(() => aui.threads.switchToThread("t2"));
+    await vi.waitFor(() => {
+      expect(handle.getClient().threads.getState().mainThreadId).toBe("t2");
+    });
+    await vi.waitFor(() => {
+      expect(load.mock.calls.length).toBeGreaterThan(afterFirst);
+    });
+    handle.destroy();
+  });
+
+  it("stops an in-flight cloud chat when switching away", async () => {
+    const chat = createCancellableTransport();
+    const handle = createAssistantClient(
+      AuiConfig({
+        threads: AISDKThreads({
+          cloud: {} as AssistantCloud,
+          threadId: "t1",
+          transport: () => chat.transport,
+        }),
+      }),
+    );
+    handle.subscribe(() => {});
+    const aui = handle.getClient();
+    try {
+      await aui.threads.getLoadThreadsPromise();
+      await vi.waitFor(() => {
+        expect(handle.getClient().threads.getState().mainThreadId).toBe("t1");
+      });
+      await vi.waitFor(() => expect(load).toHaveBeenCalled());
+      await vi.waitFor(() => {
+        expect(handle.getClient().thread.getState().isLoading).toBe(false);
+      });
+      flushTapSync(() => handle.getClient().composer.setText("stream me"));
+      flushTapSync(() => handle.getClient().composer.send());
+      await vi.waitFor(() => {
+        expect(handle.getClient().thread.getState().isRunning).toBe(true);
+      });
+      flushTapSync(() => handle.getClient().threads.switchToThread("t2"));
+      await vi.waitFor(() => {
+        expect(handle.getClient().threads.getState().mainThreadId).toBe("t2");
+      });
+      await vi.waitFor(() => {
+        expect(chat.getCancelCount()).toBe(1);
+      });
+    } finally {
+      handle.destroy();
+    }
+  });
+});