@ulearning/ai-sdk-adapter 0.0.0-stage → 1.0.0-next.0

package/src/runtime/AISDKThreads.ts

@@ -0,0 +1,194 @@
+"use client";
+
+import { resource, useResource } from "@ulearning/ai-tap";
+import { useEffect, useMemo, useState } from "@ulearning/ai-tap";
+import type { Chat } from "@ai-sdk/vue";
+import type { UIMessage } from "ai";
+import type { ChatTransport } from "ai";
+import {
+  InMemoryThreadList,
+  inMemoryThreadListTransformScopes,
+} from "@ulearning/ai-core/store";
+import { ThreadClient } from "@ulearning/ai-core/store/internal";
+import {
+  attachTransformScopes,
+  useAssistantClientRef,
+  useAssistantScopeEffect,
+} from "@ulearning/ai-store/client";
+import { AssistantChatTransport } from "../transport/AssistantChatTransport";
+import {
+  createChat,
+  splitChatThreadOptions,
+  useChatThread,
+  type ChatThreadOptions,
+} from "./useChatThread";
+import { MessageRepository } from "@ulearning/ai-core/internal";
+import { useResourceCleanup } from "./useResourceCleanup";
+
+export type AISDKThreadsOptions<UI_MESSAGE extends UIMessage = UIMessage> =
+  Omit<ChatThreadOptions<UI_MESSAGE>, "id" | "transport" | "messages"> & {
+    /**
+     * The transport threads send through. A factory is invoked once per
+     * thread so each thread owns its instance. A plain
+     * `AssistantChatTransport` instance is cloned per thread (its
+     * assistant-ui wiring is per thread); any other transport instance is
+     * shared as-is. Defaults to one `AssistantChatTransport` per thread.
+     */
+    transport?:
+      | ChatTransport<UI_MESSAGE>
+      | (() => ChatTransport<UI_MESSAGE>)
+      | undefined;
+  };
+
+type AISDKThreadChatOptions<UI_MESSAGE extends UIMessage = UIMessage> =
+  AISDKThreadsOptions<UI_MESSAGE>;
+
+type ChatOptionsRef<UI_MESSAGE extends UIMessage> = {
+  current: AISDKThreadChatOptions<UI_MESSAGE> | undefined;
+};
+
+type ChatEntry<UI_MESSAGE extends UIMessage> = {
+  chat: Chat<UI_MESSAGE>;
+  transport: ChatTransport<UI_MESSAGE>;
+  repository: MessageRepository;
+  optionsRef: ChatOptionsRef<UI_MESSAGE>;
+};
+
+const createChatEntry = <UI_MESSAGE extends UIMessage>(
+  threadId: string,
+  options: AISDKThreadChatOptions<UI_MESSAGE> | undefined,
+): ChatEntry<UI_MESSAGE> => {
+  const optionsRef: ChatOptionsRef<UI_MESSAGE> = { current: options };
+  const { chatInit } = splitChatThreadOptions(
+    options as ChatThreadOptions<UI_MESSAGE> | undefined,
+  );
+  const transport =
+    typeof options?.transport === "function"
+      ? options.transport()
+      : options?.transport === undefined
+        ? new AssistantChatTransport()
+        : options.transport instanceof AssistantChatTransport
+          ? options.transport.__internal_clone()
+          : options.transport;
+  return {
+    chat: createChat({ ...chatInit, id: threadId, transport }, optionsRef),
+    transport,
+    repository: new MessageRepository(),
+    optionsRef,
+  };
+};
+
+const getOrCreateChatEntry = <UI_MESSAGE extends UIMessage>(
+  threadId: string,
+  options: AISDKThreadChatOptions<UI_MESSAGE> | undefined,
+  chats: Map<string, ChatEntry<UI_MESSAGE>>,
+): ChatEntry<UI_MESSAGE> => {
+  const existing = chats.get(threadId);
+  if (existing) return existing;
+  const created = createChatEntry(threadId, options);
+  chats.set(threadId, created);
+  return created;
+};
+
+const useAISDKChatThread = <UI_MESSAGE extends UIMessage = UIMessage>({
+  threadId,
+  options,
+  chats,
+}: {
+  threadId: string;
+  options: AISDKThreadChatOptions<UI_MESSAGE> | undefined;
+  chats: Map<string, ChatEntry<UI_MESSAGE>>;
+}) => {
+  const { chat, transport, repository } = getOrCreateChatEntry(
+    threadId,
+    options,
+    chats,
+  );
+
+  const fallbackItem = useMemo(
+    () => ({
+      initialize: async () => ({
+        remoteId: threadId,
+        externalId: undefined,
+      }),
+    }),
+    [threadId],
+  );
+  const runtime = useChatThread(
+    { ...options, transport } as ChatThreadOptions<UI_MESSAGE>,
+    {
+      id: threadId,
+      isMainThread: true,
+      getThreadListItem: () => fallbackItem,
+      chat,
+      messageRepositoryInstance: repository,
+      stopOnClientDestroy: false,
+    },
+  );
+
+  const clientRef = useAssistantClientRef();
+  useAssistantScopeEffect(
+    "modelContext",
+    () =>
+      runtime.registerModelContextProvider(clientRef.current!.modelContext()),
+    [runtime],
+  );
+
+  return useResource(ThreadClient({ runtime: runtime.thread }));
+};
+
+const AISDKChatThread = resource(useAISDKChatThread);
+
+const useAISDKThreads = <UI_MESSAGE extends UIMessage = UIMessage>(
+  options?: AISDKThreadsOptions<UI_MESSAGE>,
+) => {
+  const threadOptions = options;
+  const [chats] = useState(() => new Map<string, ChatEntry<UI_MESSAGE>>());
+
+  useEffect(() => {
+    for (const { optionsRef } of chats.values()) {
+      optionsRef.current = threadOptions;
+    }
+  });
+
+  useResourceCleanup(true, () => {
+    for (const { chat } of chats.values()) {
+      void chat.stop().catch(() => {});
+    }
+  });
+
+  const thread = (id: string) => {
+    const element = AISDKChatThread({
+      threadId: id,
+      options: threadOptions,
+      chats,
+    });
+    return element;
+  };
+
+  return useResource(
+    InMemoryThreadList({
+      thread,
+      onDelete: (id) => {
+        void chats
+          .get(id)
+          ?.chat.stop()
+          .catch(() => {});
+        chats.delete(id);
+      },
+    }),
+  );
+};
+
+/**
+ * `AuiConfig` entry that runs one AI SDK chat per thread. Hosts the same
+ * per-thread orchestration as {@link AISDKChat} inside the client's own
+ * resource tree, so it works with any `AssistantClient` host, React or not.
+ * Threads live in memory for the client's lifetime and keep
+ * their history across switches; only the visible thread is mounted, and a
+ * switched-away chat keeps streaming into its stored state until it settles
+ * or the thread is deleted. Model context is registered on each mounted thread.
+ */
+export const AISDKThreads = resource(useAISDKThreads);
+
+attachTransformScopes(useAISDKThreads, inMemoryThreadListTransformScopes);

package/src/runtime/DynamicChatTransport.test.ts

@@ -0,0 +1,228 @@
+import type { AssistantRuntime } from "@ulearning/ai-core";
+import type { UIMessage } from "@ai-sdk/react";
+import type { ChatTransport, UIMessageChunk } from "ai";
+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("forwards replay behavior through the dynamic transport and thread proxy", () => {
+    const createTransport = (
+      resumeStreamIsReplay: boolean,
+    ): ChatTransport<UIMessage> => ({
+      resumeStreamIsReplay,
+      sendMessages: async () => new ReadableStream<UIMessageChunk>(),
+      reconnectToStream: async () => new ReadableStream<UIMessageChunk>(),
+    });
+    const dynamicTransport = new DynamicChatTransport(createTransport(false));
+    const owner = {};
+    const proxy = dynamicTransport.createThreadProxy(owner, () => ({
+      runtime: createRuntime("system"),
+      getThreadListItem: () => undefined,
+    }));
+
+    expect(dynamicTransport.resumeStreamIsReplay).toBe(false);
+    expect(proxy.resumeStreamIsReplay).toBe(false);
+
+    dynamicTransport.setTransport(createTransport(true));
+
+    expect(dynamicTransport.resumeStreamIsReplay).toBe(true);
+    expect(proxy.resumeStreamIsReplay).toBe(true);
+  });
+
+  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();
+  });
+});

package/src/runtime/DynamicChatTransport.ts

@@ -0,0 +1,281 @@
+import type { AssistantRuntime } from "@ulearning/ai-core";
+import type { UIMessage } from "ai";
+import type { ChatTransport } from "ai";
+import {
+  AssistantChatTransport,
+  type InitializableThreadListItem,
+} from "../transport/AssistantChatTransport";
+import type { ResumableClientStorage } from "../transport/resumable";
+import { getResumableAdapter } from "./getResumableAdapter";
+
+const resumedStreamIdsByStorage = new WeakMap<
+  ResumableClientStorage,
+  Set<string>
+>();
+
+export 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;
+};
+
+type ResumableStorageSubscription = {
+  listener: () => void;
+  threadId: string | undefined;
+  unsubscribe: (() => void) | undefined;
+};
+
+class DynamicResumableStorage implements ResumableClientStorage {
+  private readonly subscriptions = new Set<ResumableStorageSubscription>();
+  private storage: ResumableClientStorage | undefined;
+  private resumedStreamIds: Set<string>;
+  private hasPendingNotification = false;
+
+  constructor(storage: ResumableClientStorage | undefined) {
+    this.storage = storage;
+    this.resumedStreamIds = getResumedStreamIds(storage);
+  }
+
+  public setStorage(storage: ResumableClientStorage | undefined) {
+    if (this.storage === storage) return;
+    this.storage = storage;
+    const nextResumedStreamIds = getResumedStreamIds(storage);
+    for (const streamId of this.resumedStreamIds) {
+      nextResumedStreamIds.add(streamId);
+    }
+    this.resumedStreamIds = nextResumedStreamIds;
+    this.hasPendingNotification = true;
+  }
+
+  public flushChange() {
+    if (!this.hasPendingNotification) return;
+    this.hasPendingNotification = false;
+    for (const subscription of this.subscriptions) {
+      subscription.unsubscribe?.();
+      subscription.unsubscribe = this.storage?.subscribe?.(
+        subscription.listener,
+        subscription.threadId,
+      );
+      subscription.listener();
+    }
+  }
+
+  public getResumedStreamIds() {
+    return this.resumedStreamIds;
+  }
+
+  public getStreamId(threadId?: string) {
+    return this.storage?.getStreamId(threadId) ?? null;
+  }
+
+  public setStreamId(id: string, threadId?: string) {
+    this.storage?.setStreamId(id, threadId);
+  }
+
+  public clear(threadId?: string) {
+    this.storage?.clear(threadId);
+  }
+
+  public subscribe(listener: () => void, threadId?: string) {
+    const subscription: ResumableStorageSubscription = {
+      listener,
+      threadId,
+      unsubscribe: this.storage?.subscribe?.(listener, threadId),
+    };
+    this.subscriptions.add(subscription);
+    return () => {
+      this.subscriptions.delete(subscription);
+      subscription.unsubscribe?.();
+    };
+  }
+}
+
+type ThreadTransportContext<UI_MESSAGE extends UIMessage> = {
+  owner: object;
+  sourceTransport?: ChatTransport<UI_MESSAGE> | undefined;
+  transport?: ChatTransport<UI_MESSAGE> | undefined;
+  runtime?: AssistantRuntime | undefined;
+  getThreadListItem?:
+    | (() => InitializableThreadListItem | undefined)
+    | undefined;
+};
+
+type ThreadTransportBinding = {
+  runtime: AssistantRuntime;
+  getThreadListItem: () => InitializableThreadListItem | undefined;
+};
+
+export class DynamicChatTransport<
+  UI_MESSAGE extends UIMessage,
+> implements ChatTransport<UI_MESSAGE> {
+  private readonly threadContexts = new Map<
+    string,
+    ThreadTransportContext<UI_MESSAGE>
+  >();
+  private transport: ChatTransport<UI_MESSAGE>;
+  private readonly resumableStorage: DynamicResumableStorage;
+
+  constructor(transport: ChatTransport<UI_MESSAGE>) {
+    this.transport = transport;
+    this.resumableStorage = new DynamicResumableStorage(
+      getResumableAdapter(transport)?.storage,
+    );
+  }
+
+  public readonly sendMessages: ChatTransport<UI_MESSAGE>["sendMessages"] = (
+    options,
+  ) => this.getTransport(options.chatId).sendMessages(options);
+
+  public readonly reconnectToStream: ChatTransport<UI_MESSAGE>["reconnectToStream"] =
+    (options) => this.getTransport(options.chatId).reconnectToStream(options);
+
+  public get resumeStreamIsReplay() {
+    return this.transport.resumeStreamIsReplay ?? false;
+  }
+
+  public readonly getCurrentTransport = (chatId: string) =>
+    this.getThreadTransport(this.threadContexts.get(chatId)) ?? this.transport;
+
+  public readonly getCurrentResumableStorage = () => this.resumableStorage;
+
+  public readonly getResumedStreamIds = () =>
+    this.resumableStorage.getResumedStreamIds();
+
+  public createThreadProxy(
+    owner: object,
+    getBinding: () => ThreadTransportBinding,
+  ): ChatTransport<UI_MESSAGE> {
+    const dynamicTransport = this;
+    return {
+      get resumeStreamIsReplay() {
+        return dynamicTransport.resumeStreamIsReplay;
+      },
+      sendMessages: (options) =>
+        this.getOrCreateBoundTransport(
+          options.chatId,
+          owner,
+          getBinding,
+        ).sendMessages(options),
+      reconnectToStream: (options) =>
+        this.getOrCreateBoundTransport(
+          options.chatId,
+          owner,
+          getBinding,
+        ).reconnectToStream(options),
+    };
+  }
+
+  public setTransport(transport: ChatTransport<UI_MESSAGE>) {
+    if (this.transport === transport) return;
+    this.transport = transport;
+    this.resumableStorage.setStorage(getResumableAdapter(transport)?.storage);
+  }
+
+  public flushTransportChange() {
+    this.resumableStorage.flushChange();
+  }
+
+  public registerThread(chatId: string, owner: object) {
+    const existing = this.threadContexts.get(chatId);
+    if (existing?.owner === owner) return;
+    this.threadContexts.set(chatId, {
+      owner,
+    });
+  }
+
+  public setThreadContext(
+    chatId: string,
+    owner: object,
+    runtime: AssistantRuntime,
+    getThreadListItem: () => InitializableThreadListItem | undefined,
+  ) {
+    this.registerThread(chatId, owner);
+    const context = this.threadContexts.get(chatId)!;
+    context.runtime = runtime;
+    context.getThreadListItem = getThreadListItem;
+    if (context.transport === undefined) {
+      this.getThreadTransport(context);
+      return;
+    }
+    if (
+      context.sourceTransport === this.transport &&
+      context.transport !== undefined
+    ) {
+      this.wireTransport(context, context.transport);
+    }
+  }
+
+  public unregisterThread(chatId: string, owner: object) {
+    if (this.threadContexts.get(chatId)?.owner === owner) {
+      this.threadContexts.delete(chatId);
+    }
+  }
+
+  private getTransport(chatId: string) {
+    const context = this.threadContexts.get(chatId);
+    if (!context) {
+      throw new Error(
+        `DynamicChatTransport has no registered context for chat "${chatId}"`,
+      );
+    }
+    return this.getThreadTransport(context)!;
+  }
+
+  private getBoundTransport(
+    chatId: string,
+    owner: object,
+    binding: ThreadTransportBinding,
+  ) {
+    this.registerThread(chatId, owner);
+    const context = this.threadContexts.get(chatId)!;
+    context.runtime = binding.runtime;
+    context.getThreadListItem = binding.getThreadListItem;
+    return this.getThreadTransport(context)!;
+  }
+
+  private getOrCreateBoundTransport(
+    chatId: string,
+    owner: object,
+    getBinding: () => ThreadTransportBinding,
+  ) {
+    const context = this.threadContexts.get(chatId);
+    return context?.owner === owner
+      ? this.getThreadTransport(context)!
+      : this.getBoundTransport(chatId, owner, getBinding());
+  }
+
+  private getThreadTransport(
+    context: ThreadTransportContext<UI_MESSAGE> | undefined,
+  ) {
+    if (!context) return undefined;
+    if (context.sourceTransport !== this.transport) {
+      context.transport = this.createThreadTransport(this.transport);
+      context.sourceTransport = this.transport;
+    }
+    this.wireTransport(context, context.transport!);
+    return context.transport!;
+  }
+
+  private createThreadTransport(transport: ChatTransport<UI_MESSAGE>) {
+    return transport instanceof AssistantChatTransport
+      ? transport.__internal_clone()
+      : transport;
+  }
+
+  private wireTransport(
+    context: ThreadTransportContext<UI_MESSAGE>,
+    transport: ChatTransport<UI_MESSAGE>,
+  ) {
+    if (!(transport instanceof AssistantChatTransport)) return;
+    if (context.runtime) transport.setRuntime(context.runtime);
+    if (context.getThreadListItem) {
+      transport.__internal_setGetThreadListItem(context.getThreadListItem);
+    }
+  }
+}