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

package/src/transport/AssistantChatTransport.ts

@@ -0,0 +1,236 @@
+import type {
+  AssistantRuntime,
+  ThreadListItemRuntime,
+} from "@ulearning/ai-core";
+import {
+  DefaultChatTransport,
+  type HttpChatTransportInitOptions,
+  type UIMessage,
+} from "ai";
+import { toToolsJSONSchema } from "@ulearning/ai-stream";
+import { SSEEventDecoder } from "@ulearning/ai-stream/utils";
+import {
+  RESUMABLE_STREAM_ID_HEADER,
+  type AssistantChatResumableOptions,
+} from "./resumable";
+
+export type InitializableThreadListItem = Pick<
+  ThreadListItemRuntime,
+  "initialize"
+>;
+
+const FINISH_BUFFER_LIMIT = 4096;
+const FINISH_BUFFER_TAIL = 1024;
+const RESUMABLE_THREAD_ID_HEADER = "x-assistant-ui-resumable-thread-id";
+const RESUMABLE_RECONNECT_ID_HEADER = "x-assistant-ui-resumable-reconnect-id";
+
+// 101/204/205/304 are null-body statuses per the fetch spec: `new Response(body, { status })`
+// throws for them, and WebKit returns a non-null empty body, so the body check alone does not guard it.
+const NULL_BODY_STATUSES = new Set([101, 204, 205, 304]);
+
+export type AssistantChatTransportInitOptions<UI_MESSAGE extends UIMessage> =
+  HttpChatTransportInitOptions<UI_MESSAGE> & {
+    resumable?: AssistantChatResumableOptions;
+  };
+
+export class AssistantChatTransport<
+  UI_MESSAGE extends UIMessage,
+> extends DefaultChatTransport<UI_MESSAGE> {
+  private runtime: AssistantRuntime | undefined;
+  private getThreadListItem:
+    | (() => InitializableThreadListItem | undefined)
+    | undefined;
+  private readonly resumable: AssistantChatResumableOptions | undefined;
+
+  private readonly __internal_initOptions:
+    | AssistantChatTransportInitOptions<UI_MESSAGE>
+    | undefined;
+
+  /** Constructs an unwired copy with the same init options. */
+  public __internal_clone(): AssistantChatTransport<UI_MESSAGE> {
+    const Constructor = this.constructor as new (
+      initOptions?: AssistantChatTransportInitOptions<UI_MESSAGE>,
+    ) => AssistantChatTransport<UI_MESSAGE>;
+    return new Constructor(this.__internal_initOptions);
+  }
+
+  constructor(initOptions?: AssistantChatTransportInitOptions<UI_MESSAGE>) {
+    const { resumable, ...rest } = initOptions ?? {};
+    const userFetch = rest.fetch;
+    const userPrepareReconnect = rest.prepareReconnectToStreamRequest;
+
+    super({
+      ...rest,
+      ...(resumable && {
+        fetch: wrapFetchWithResumable(resumable, userFetch),
+        prepareReconnectToStreamRequest: wrapPrepareReconnect(
+          resumable,
+          userPrepareReconnect,
+        ),
+      }),
+      prepareSendMessagesRequest: async (options) => {
+        const threadId = options.id;
+        const context = this.runtime?.thread.getModelContext();
+        const threadListItem =
+          this.getThreadListItem?.() ?? this.runtime?.threads.mainItem;
+        const id = (await threadListItem?.initialize())?.remoteId ?? options.id;
+
+        const optionsEx = {
+          ...options,
+          id,
+          body: {
+            callSettings: context?.callSettings,
+            system: context?.system,
+            config: context?.config,
+            tools: toToolsJSONSchema(context?.tools ?? {}),
+            ...options?.body,
+          },
+        };
+        const preparedRequest =
+          await rest.prepareSendMessagesRequest?.(optionsEx);
+        const headers = resumable
+          ? new Headers(preparedRequest?.headers ?? options.headers)
+          : undefined;
+        headers?.set(RESUMABLE_THREAD_ID_HEADER, threadId);
+
+        return {
+          ...preparedRequest,
+          ...(headers && { headers }),
+          body: preparedRequest?.body ?? {
+            ...optionsEx.body,
+            id,
+            messages: options.messages,
+            trigger: options.trigger,
+            messageId: options.messageId,
+            metadata: options.requestMetadata,
+          },
+        };
+      },
+    });
+
+    this.resumable = resumable;
+    this.__internal_initOptions = initOptions;
+  }
+
+  setRuntime(runtime: AssistantRuntime) {
+    this.runtime = runtime;
+  }
+
+  getResumableAdapter(): AssistantChatResumableOptions | undefined {
+    return this.resumable;
+  }
+
+  __internal_setGetThreadListItem(
+    getter: () => InitializableThreadListItem | undefined,
+  ) {
+    this.getThreadListItem = getter;
+  }
+}
+
+function wrapFetchWithResumable(
+  resumable: AssistantChatResumableOptions,
+  userFetch: HttpChatTransportInitOptions<UIMessage>["fetch"],
+): NonNullable<HttpChatTransportInitOptions<UIMessage>["fetch"]> {
+  const baseFetch: typeof globalThis.fetch = userFetch
+    ? (input, init) => userFetch(input as RequestInfo | URL, init)
+    : globalThis.fetch.bind(globalThis);
+
+  return async (input, init) => {
+    const headers = new Headers(init?.headers);
+    const threadId = headers.get(RESUMABLE_THREAD_ID_HEADER) ?? undefined;
+    const reconnectingStreamId = headers.get(RESUMABLE_RECONNECT_ID_HEADER);
+    const checkpointAtRequest =
+      reconnectingStreamId ?? resumable.storage.getStreamId(threadId);
+    headers.delete(RESUMABLE_THREAD_ID_HEADER);
+    headers.delete(RESUMABLE_RECONNECT_ID_HEADER);
+    const res = await baseFetch(input, { ...init, headers });
+    const id = res.headers.get(RESUMABLE_STREAM_ID_HEADER);
+    const ownsCheckpoint =
+      !reconnectingStreamId ||
+      resumable.storage.getStreamId(threadId) === reconnectingStreamId;
+    if (id && ownsCheckpoint) resumable.storage.setStreamId(id, threadId);
+    if (
+      (res.status === 204 || res.status === 404) &&
+      reconnectingStreamId &&
+      resumable.storage.getStreamId(threadId) === reconnectingStreamId
+    ) {
+      resumable.storage.clear(threadId);
+    }
+    if (!res.body || NULL_BODY_STATUSES.has(res.status)) return res;
+
+    const detectFinish =
+      resumable.isFinishEvent ?? createDefaultFinishDetector();
+    const decoder = resumable.isFinishEvent ? new TextDecoder() : undefined;
+    let accumulator = "";
+    const tap = new TransformStream<Uint8Array, Uint8Array>({
+      transform(chunk, controller) {
+        controller.enqueue(chunk);
+        accumulator += decoder?.decode(chunk, { stream: true }) ?? "";
+        if (detectFinish(chunk, accumulator)) {
+          if (
+            ownsCheckpoint &&
+            resumable.storage.getStreamId(threadId) ===
+              (id ?? checkpointAtRequest)
+          ) {
+            resumable.storage.clear(threadId);
+          }
+          accumulator = "";
+        } else if (accumulator.length > FINISH_BUFFER_LIMIT) {
+          accumulator = accumulator.slice(-FINISH_BUFFER_TAIL);
+        }
+      },
+    });
+
+    return new Response(res.body.pipeThrough(tap), {
+      status: res.status,
+      statusText: res.statusText,
+      headers: res.headers,
+    });
+  };
+}
+
+function createDefaultFinishDetector() {
+  const decoder = new TextDecoder("utf-8", { ignoreBOM: true });
+  // Observing completion must not impose a smaller event limit than the AI SDK.
+  const events = new SSEEventDecoder({
+    maxLineLength: Number.MAX_SAFE_INTEGER,
+    maxEventLength: Number.MAX_SAFE_INTEGER,
+  });
+  return (chunk: Uint8Array) =>
+    events.push(decoder.decode(chunk, { stream: true })).some(({ data }) => {
+      try {
+        return JSON.parse(data)?.type === "finish";
+      } catch {
+        return false;
+      }
+    });
+}
+
+function wrapPrepareReconnect(
+  resumable: AssistantChatResumableOptions,
+  userPrepareReconnect: HttpChatTransportInitOptions<UIMessage>["prepareReconnectToStreamRequest"],
+): NonNullable<
+  HttpChatTransportInitOptions<UIMessage>["prepareReconnectToStreamRequest"]
+> {
+  return async (options) => {
+    const streamId = resumable.storage.getStreamId(options.id);
+    if (!streamId) {
+      throw new Error(
+        "AssistantChatTransport: no resumable stream id available; nothing to resume",
+      );
+    }
+    const api =
+      typeof resumable.resumeApi === "function"
+        ? resumable.resumeApi(streamId)
+        : resumable.resumeApi;
+    const userPrepared = await userPrepareReconnect?.({ ...options, api });
+    const headers = new Headers(userPrepared?.headers ?? options.headers);
+    headers.set(RESUMABLE_THREAD_ID_HEADER, options.id);
+    headers.set(RESUMABLE_RECONNECT_ID_HEADER, streamId);
+    return {
+      ...userPrepared,
+      headers,
+      api: userPrepared?.api ?? api,
+    };
+  };
+}

package/src/transport/resumable.test.ts

@@ -0,0 +1,262 @@
+// @vitest-environment jsdom
+
+import { afterEach, describe, expect, it, vi } from "vitest";
+import { createResumableSessionStorage } from "./resumable";
+
+const originalSessionStorageDescriptor = Object.getOwnPropertyDescriptor(
+  window,
+  "sessionStorage",
+);
+
+afterEach(() => {
+  vi.restoreAllMocks();
+  if (originalSessionStorageDescriptor) {
+    Object.defineProperty(
+      window,
+      "sessionStorage",
+      originalSessionStorageDescriptor,
+    );
+  }
+  window.sessionStorage.clear();
+});
+
+describe("createResumableSessionStorage", () => {
+  it("stores stream ids in sessionStorage", () => {
+    const storage = createResumableSessionStorage({ key: "test-stream-id" });
+
+    expect(storage.getStreamId()).toBeNull();
+
+    storage.setStreamId("stream-1");
+    expect(storage.getStreamId()).toBe("stream-1");
+
+    storage.clear();
+    expect(storage.getStreamId()).toBeNull();
+  });
+
+  it("notifies subscribers when the stored stream id changes", () => {
+    const storage = createResumableSessionStorage({ key: "test-stream-id" });
+    const listener = vi.fn();
+    const unsubscribe = storage.subscribe?.(listener);
+
+    storage.setStreamId("stream-1");
+    storage.clear();
+
+    expect(listener).toHaveBeenCalledTimes(2);
+
+    unsubscribe?.();
+    storage.setStreamId("stream-2");
+
+    expect(listener).toHaveBeenCalledTimes(2);
+  });
+
+  it("isolates the current stream owner and subscribers by thread", () => {
+    const storage = createResumableSessionStorage({ key: "test-stream-id" });
+    const listenerA = vi.fn();
+    const listenerB = vi.fn();
+    storage.subscribe?.(listenerA, "thread-a");
+    storage.subscribe?.(listenerB, "thread-b");
+
+    storage.setStreamId("stream-a", "thread-a");
+
+    expect(storage.getStreamId("thread-a")).toBe("stream-a");
+    expect(storage.getStreamId("thread-b")).toBeNull();
+    expect(listenerA).toHaveBeenCalledTimes(1);
+    expect(listenerB).not.toHaveBeenCalled();
+
+    storage.clear("thread-b");
+
+    expect(storage.getStreamId("thread-a")).toBe("stream-a");
+    expect(listenerA).toHaveBeenCalledTimes(1);
+    expect(listenerB).not.toHaveBeenCalled();
+
+    storage.clear("thread-a");
+
+    expect(storage.getStreamId("thread-a")).toBeNull();
+    expect(listenerA).toHaveBeenCalledTimes(2);
+    expect(listenerB).not.toHaveBeenCalled();
+  });
+
+  it("keeps the persisted key readable after the storage is recreated", () => {
+    const storage = createResumableSessionStorage({ key: "test-stream-id" });
+    storage.setStreamId("stream-a", "__LOCALID_before_reload");
+
+    const reloadedStorage = createResumableSessionStorage({
+      key: "test-stream-id",
+    });
+
+    expect(reloadedStorage.getStreamId("__LOCALID_after_reload")).toBe(
+      "stream-a",
+    );
+  });
+
+  it("preserves the current owner for unscoped storage updates", () => {
+    const storage = createResumableSessionStorage({ key: "test-stream-id" });
+    storage.setStreamId("stream-a", "thread-a");
+    storage.subscribe?.(vi.fn(), "thread-b");
+
+    storage.setStreamId("stream-b");
+
+    expect(storage.getStreamId("thread-a")).toBe("stream-b");
+    expect(storage.getStreamId("thread-b")).toBeNull();
+
+    storage.clear("thread-a");
+    expect(storage.getStreamId("thread-a")).toBeNull();
+  });
+
+  it("caches the last known value between writes", () => {
+    const storageMethods = {
+      getItem: vi.fn(() => null),
+      setItem: vi.fn(),
+      removeItem: vi.fn(),
+    } as unknown as Storage;
+    Object.defineProperty(window, "sessionStorage", {
+      configurable: true,
+      value: storageMethods,
+    });
+    const storage = createResumableSessionStorage({ key: "test-stream-id" });
+
+    expect(storage.getStreamId("thread-a")).toBeNull();
+    expect(storage.getStreamId("thread-a")).toBeNull();
+
+    expect(storageMethods.getItem).toHaveBeenCalledTimes(1);
+
+    storage.setStreamId("stream-a", "thread-a");
+
+    expect(storage.getStreamId("thread-a")).toBe("stream-a");
+    expect(storageMethods.getItem).toHaveBeenCalledTimes(1);
+  });
+
+  it("isolates subscriber errors", () => {
+    const storage = createResumableSessionStorage({ key: "test-stream-id" });
+    const error = new Error("listener failed");
+    const consoleError = vi
+      .spyOn(console, "error")
+      .mockImplementation(() => {});
+    const laterListener = vi.fn();
+
+    storage.subscribe?.(() => {
+      throw error;
+    });
+    storage.subscribe?.(laterListener);
+
+    expect(() => storage.setStreamId("stream-1")).not.toThrow();
+    expect(laterListener).toHaveBeenCalledOnce();
+    expect(consoleError).toHaveBeenCalledWith(
+      "[assistant-ui] resumable storage listener failed",
+      error,
+    );
+  });
+
+  it("degrades to null and no-op when sessionStorage access is blocked", () => {
+    Object.defineProperty(window, "sessionStorage", {
+      configurable: true,
+      get() {
+        throw new DOMException("blocked", "SecurityError");
+      },
+    });
+
+    const storage = createResumableSessionStorage();
+
+    expect(storage.getStreamId()).toBeNull();
+    expect(() => storage.setStreamId("stream-1")).not.toThrow();
+    expect(() => storage.clear()).not.toThrow();
+  });
+
+  it("degrades to null and no-op when sessionStorage methods throw", () => {
+    const throwingStorage = {
+      getItem: vi.fn(() => {
+        throw new DOMException("blocked", "SecurityError");
+      }),
+      setItem: vi.fn(() => {
+        throw new DOMException("blocked", "SecurityError");
+      }),
+      removeItem: vi.fn(() => {
+        throw new DOMException("blocked", "SecurityError");
+      }),
+    } as unknown as Storage;
+
+    Object.defineProperty(window, "sessionStorage", {
+      configurable: true,
+      value: throwingStorage,
+    });
+
+    const storage = createResumableSessionStorage();
+
+    expect(storage.getStreamId()).toBeNull();
+    expect(() => storage.setStreamId("stream-1")).not.toThrow();
+    expect(() => storage.clear()).not.toThrow();
+    expect(throwingStorage.getItem).toHaveBeenCalledTimes(1);
+    expect(throwingStorage.setItem).toHaveBeenCalledTimes(1);
+    expect(throwingStorage.removeItem).toHaveBeenCalledTimes(1);
+  });
+
+  it("scopes the stored id by a getter key, isolating instances", () => {
+    const storageA = createResumableSessionStorage({ key: () => "thread-a" });
+    const storageB = createResumableSessionStorage({ key: () => "thread-b" });
+
+    storageA.setStreamId("stream-a");
+    expect(storageA.getStreamId()).toBe("stream-a");
+    expect(storageB.getStreamId()).toBeNull();
+
+    storageB.setStreamId("stream-b");
+    expect(storageA.getStreamId()).toBe("stream-a");
+    expect(storageB.getStreamId()).toBe("stream-b");
+
+    storageA.clear();
+    expect(storageA.getStreamId()).toBeNull();
+    expect(storageB.getStreamId()).toBe("stream-b");
+  });
+
+  it("reads a getter key lazily on every access", () => {
+    let current = "first";
+    const storage = createResumableSessionStorage({ key: () => current });
+
+    storage.setStreamId("stream-1");
+    expect(storage.getStreamId()).toBe("stream-1");
+
+    current = "second";
+    expect(storage.getStreamId()).toBeNull();
+
+    storage.setStreamId("stream-2");
+    expect(storage.getStreamId()).toBe("stream-2");
+    expect(window.sessionStorage.getItem("first")).toBe("stream-1");
+    expect(window.sessionStorage.getItem("second")).toBe("stream-2");
+
+    storage.clear();
+    expect(window.sessionStorage.getItem("second")).toBeNull();
+    expect(window.sessionStorage.getItem("first")).toBe("stream-1");
+  });
+
+  it("disables storage access while the getter returns undefined", () => {
+    let current: string | undefined;
+    const storage = createResumableSessionStorage({ key: () => current });
+
+    storage.setStreamId("stream-1");
+    expect(storage.getStreamId()).toBeNull();
+    expect(window.sessionStorage.getItem("aui-resumable-stream-id")).toBeNull();
+
+    current = "thread-a";
+    storage.setStreamId("stream-1");
+    expect(storage.getStreamId()).toBe("stream-1");
+
+    current = undefined;
+    expect(storage.getStreamId()).toBeNull();
+    storage.clear();
+
+    current = "thread-a";
+    expect(storage.getStreamId()).toBe("stream-1");
+  });
+
+  it("treats a throwing key getter as no key", () => {
+    const storage = createResumableSessionStorage({
+      key: () => {
+        throw new Error("no thread context");
+      },
+    });
+
+    expect(() => storage.setStreamId("stream-1")).not.toThrow();
+    expect(storage.getStreamId()).toBeNull();
+    expect(() => storage.clear()).not.toThrow();
+    expect(window.sessionStorage.getItem("aui-resumable-stream-id")).toBeNull();
+  });
+});

package/src/transport/resumable.ts

@@ -0,0 +1,162 @@
+"use client";
+
+import { RESUMABLE_STREAM_ID_HEADER as RESUMABLE_STREAM_ID_HEADER_VALUE } from "@ulearning/ai-stream/resumable";
+
+/** Response header used by the [Resumable Streams](/docs/guides/resumable-streams) server and client wiring. */
+export const RESUMABLE_STREAM_ID_HEADER = RESUMABLE_STREAM_ID_HEADER_VALUE;
+
+const DEFAULT_STORAGE_KEY = "aui-resumable-stream-id";
+
+export type ResumableClientStorage = {
+  getStreamId(threadId?: string): string | null;
+  setStreamId(id: string, threadId?: string): void;
+  clear(threadId?: string): void;
+  /** Subscribes to stream id changes so automatic resume can react after mount. */
+  subscribe?(listener: () => void, threadId?: string): () => void;
+};
+
+const getSessionStorage = (): Storage | null => {
+  if (typeof window === "undefined") return null;
+  try {
+    return window.sessionStorage;
+  } catch {
+    return null;
+  }
+};
+
+type StorageSlot = {
+  value: string | null;
+  owner: string | undefined;
+};
+
+/** `sessionStorage`-backed storage for the pending resumable stream id. See the [Resumable Streams](/docs/guides/resumable-streams) guide for end-to-end wiring. */
+export function createResumableSessionStorage(options?: {
+  /**
+   * Storage key for the pending stream id. A static string namespaces per route
+   * or chat surface. A getter is read lazily on every access, so the key can be
+   * derived from the active thread's identity; while the getter returns
+   * `undefined`, reads report no pending stream and writes are dropped, so a
+   * thread whose identity is not known yet never touches another thread's key.
+   *
+   * Under a remote thread list with more than one thread, scope the key per
+   * thread and create one storage instance per thread runtime rather than a
+   * single shared one. A shared key is written and cleared by whichever thread
+   * acts last, so one conversation's stream can resume inside another.
+   */
+  key?: string | (() => string | undefined);
+}): ResumableClientStorage {
+  const keyOption = options?.key;
+  const resolveKey = (): string | undefined => {
+    if (typeof keyOption !== "function")
+      return keyOption ?? DEFAULT_STORAGE_KEY;
+    try {
+      return keyOption();
+    } catch {
+      return undefined;
+    }
+  };
+  const slots = new Map<string, StorageSlot>();
+  const listeners = new Set<{
+    listener: () => void;
+    threadId: string | undefined;
+  }>();
+  const readSlot = (key: string): StorageSlot => {
+    let slot = slots.get(key);
+    if (slot) return slot;
+    const storage = getSessionStorage();
+    let value: string | null = null;
+    if (storage) {
+      try {
+        value = storage.getItem(key);
+      } catch {
+        value = null;
+      }
+    }
+    slot = { value, owner: undefined };
+    slots.set(key, slot);
+    return slot;
+  };
+  const notify = (threadId?: string) => {
+    for (const subscription of listeners) {
+      if (threadId && subscription.threadId !== threadId) continue;
+      try {
+        subscription.listener();
+      } catch (error) {
+        console.error(
+          "[assistant-ui] resumable storage listener failed",
+          error,
+        );
+      }
+    }
+  };
+
+  return {
+    getStreamId(threadId) {
+      const key = resolveKey();
+      if (!key) return null;
+      const slot = readSlot(key);
+      if (slot.value === null) return null;
+      if (slot.owner && threadId && slot.owner !== threadId) return null;
+      return slot.value;
+    },
+    setStreamId(id, threadId) {
+      const key = resolveKey();
+      const storage = getSessionStorage();
+      if (!key || !storage) return;
+      try {
+        storage.setItem(key, id);
+      } catch {
+        // Ignore blocked or unavailable sessionStorage.
+        return;
+      }
+      const slot = readSlot(key);
+      slot.value = id;
+      if (threadId) {
+        slot.owner = threadId;
+      } else if (!slot.owner) {
+        slot.owner = Array.from(listeners).find(
+          (subscription) => subscription.threadId !== undefined,
+        )?.threadId;
+      }
+      notify(slot.owner);
+    },
+    clear(threadId) {
+      const key = resolveKey();
+      if (!key) return;
+      const slot = readSlot(key);
+      if (slot.owner && threadId && slot.owner !== threadId) return;
+      const storage = getSessionStorage();
+      if (!storage) return;
+      try {
+        storage.removeItem(key);
+      } catch {
+        // Ignore blocked or unavailable sessionStorage.
+        return;
+      }
+      slot.value = null;
+      slot.owner = undefined;
+      notify(threadId);
+    },
+    subscribe(listener, threadId) {
+      const key = resolveKey();
+      if (key && threadId) {
+        const slot = readSlot(key);
+        if (slot.value !== null) slot.owner ??= threadId;
+      }
+      const subscription = { listener, threadId };
+      listeners.add(subscription);
+      return () => listeners.delete(subscription);
+    },
+  };
+}
+
+export type AssistantChatResumableOptions = {
+  storage: ResumableClientStorage;
+  resumeApi: string | ((streamId: string) => string);
+  /**
+   * Called for each chunk with the accumulated decoded text. Return true when
+   * the stream has finished. The default detects complete SSE events whose
+   * top-level JSON type is `finish`. Cancellation does not invoke this callback.
+   */
+  isFinishEvent?: (chunk: Uint8Array, accumulator: string) => boolean;
+};