@assistant-ui/ai-sdk 0.0.1

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,161 @@
+"use client";
+
+import { RESUMABLE_STREAM_ID_HEADER as RESUMABLE_STREAM_ID_HEADER_VALUE } from "assistant-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);
+  /**
+   * Defaults to scanning for the AI SDK UIMessageStream `finish` marker.
+   * Cancellation never invokes this callback, only natural completion does.
+   */
+  isFinishEvent?: (chunk: Uint8Array, accumulator: string) => boolean;
+};

package/src/usage.test.ts

@@ -0,0 +1,186 @@
+import { describe, expect, it } from "vitest";
+import { getLatestThreadTokenUsage, getThreadMessageTokenUsage } from "./usage";
+
+function msg(metadata: unknown): { role: "assistant"; metadata: unknown } {
+  return {
+    role: "assistant",
+    metadata,
+  };
+}
+
+describe("getThreadMessageTokenUsage", () => {
+  it("does not double-count reasoning/cached in fallback totalTokens", () => {
+    const usage = getThreadMessageTokenUsage(
+      msg({
+        usage: {
+          inputTokens: 4,
+          outputTokens: 6,
+          reasoningTokens: 9,
+          cachedInputTokens: 3,
+        },
+      }),
+    );
+
+    // totalTokens = input + output only; reasoning/cached are detail fields
+    expect(usage).toEqual({
+      totalTokens: 10,
+      inputTokens: 4,
+      outputTokens: 6,
+      reasoningTokens: 9,
+      cachedInputTokens: 3,
+    });
+  });
+
+  it("reads v7 token detail objects for reasoning and cached tokens", () => {
+    const usage = getThreadMessageTokenUsage(
+      msg({
+        usage: {
+          inputTokens: 4,
+          outputTokens: 6,
+          inputTokenDetails: { cacheReadTokens: 3, cacheWriteTokens: 1 },
+          outputTokenDetails: { textTokens: 5, reasoningTokens: 9 },
+        },
+      }),
+    );
+
+    expect(usage).toEqual({
+      totalTokens: 10,
+      inputTokens: 4,
+      outputTokens: 6,
+      reasoningTokens: 9,
+      cachedInputTokens: 3,
+    });
+  });
+
+  it("does not fabricate zero splits when only totalTokens is present", () => {
+    const usage = getThreadMessageTokenUsage(
+      msg({ usage: { totalTokens: 12 } }),
+    );
+
+    expect(usage).toEqual({ totalTokens: 12 });
+    expect(usage).not.toHaveProperty("inputTokens");
+    expect(usage).not.toHaveProperty("outputTokens");
+  });
+
+  it("retains partial usage when only inputTokens is present", () => {
+    const usage = getThreadMessageTokenUsage(
+      msg({ usage: { inputTokens: 10 } }),
+    );
+
+    expect(usage).toEqual({ inputTokens: 10 });
+    expect(usage).not.toHaveProperty("totalTokens");
+  });
+
+  it("retains partial usage when only outputTokens is present", () => {
+    const usage = getThreadMessageTokenUsage(
+      msg({ usage: { outputTokens: 4 } }),
+    );
+
+    expect(usage).toEqual({ outputTokens: 4 });
+    expect(usage).not.toHaveProperty("totalTokens");
+  });
+
+  it("retains detail-only usage when only reasoning/cached tokens are present", () => {
+    const usage = getThreadMessageTokenUsage(
+      msg({ usage: { reasoningTokens: 7, cachedInputTokens: 2 } }),
+    );
+
+    expect(usage).toEqual({ reasoningTokens: 7, cachedInputTokens: 2 });
+    expect(usage).not.toHaveProperty("totalTokens");
+  });
+
+  it("aggregates multi-step usage without inflating totals", () => {
+    const usage = getThreadMessageTokenUsage(
+      msg({
+        steps: [
+          { usage: { inputTokens: 3, outputTokens: 2, reasoningTokens: 11 } },
+          { usage: { inputTokens: 4, outputTokens: 1, reasoningTokens: 13 } },
+        ],
+      }),
+    );
+
+    expect(usage).toEqual({
+      totalTokens: 10,
+      inputTokens: 7,
+      outputTokens: 3,
+      reasoningTokens: 24,
+    });
+  });
+
+  it("omits totalTokens if any step lacks a computable total", () => {
+    const usage = getThreadMessageTokenUsage(
+      msg({
+        steps: [
+          { usage: { inputTokens: 5, outputTokens: 5 } }, // total = 10
+          { usage: { reasoningTokens: 10 } }, // total = undefined
+        ],
+      }),
+    );
+
+    // Sums known partials but omits the invalid total
+    expect(usage).toEqual({
+      inputTokens: 5,
+      outputTokens: 5,
+      reasoningTokens: 10,
+    });
+    expect(usage).not.toHaveProperty("totalTokens");
+  });
+
+  it("aggregates totalTokens if all steps have a computable total", () => {
+    const usage = getThreadMessageTokenUsage(
+      msg({
+        steps: [
+          { usage: { inputTokens: 5, outputTokens: 5 } }, // implicit total = 10
+          { usage: { totalTokens: 15 } }, // explicit total = 15
+        ],
+      }),
+    );
+
+    expect(usage).toEqual({
+      totalTokens: 25,
+      inputTokens: 5,
+      outputTokens: 5,
+    });
+  });
+
+  it("omits totalTokens when one step has only total and another has only input", () => {
+    const usage = getThreadMessageTokenUsage(
+      msg({
+        steps: [{ usage: { totalTokens: 10 } }, { usage: { inputTokens: 3 } }],
+      }),
+    );
+
+    expect(usage).toEqual({
+      inputTokens: 3,
+    });
+    expect(usage).not.toHaveProperty("totalTokens");
+  });
+});
+
+describe("getLatestThreadTokenUsage", () => {
+  it("falls back to the latest assistant message with usage", () => {
+    const usage = getLatestThreadTokenUsage([
+      { role: "assistant", metadata: { usage: { totalTokens: 100 } } },
+      { role: "user", metadata: {} },
+      { role: "assistant", metadata: {} },
+    ]);
+
+    expect(usage).toEqual({ totalTokens: 100 });
+  });
+
+  it("prefers the newest assistant message when it has usage", () => {
+    const usage = getLatestThreadTokenUsage([
+      { role: "assistant", metadata: { usage: { totalTokens: 100 } } },
+      {
+        role: "assistant",
+        metadata: { usage: { inputTokens: 40, outputTokens: 2 } },
+      },
+    ]);
+
+    expect(usage).toEqual({
+      totalTokens: 42,
+      inputTokens: 40,
+      outputTokens: 2,
+    });
+  });
+});