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

package/src/usage.test.ts

@@ -0,0 +1,204 @@
+import { describe, expect, it } from "vitest";
+import { getThreadMessageTokenUsage } from "./usage";
+
+function msg(metadata: unknown): { role: "assistant"; metadata: unknown } {
+  return {
+    role: "assistant",
+    metadata,
+  };
+}
+
+describe("getThreadMessageTokenUsage", () => {
+  it("reads usage from custom.usage", () => {
+    const usage = getThreadMessageTokenUsage(
+      msg({
+        custom: {
+          usage: { inputTokens: 4, outputTokens: 6 },
+        },
+      }),
+    );
+
+    expect(usage).toEqual({
+      totalTokens: 10,
+      inputTokens: 4,
+      outputTokens: 6,
+    });
+  });
+
+  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("getThreadMessageTokenUsage", () => {
+  it("returns token usage from an earlier assistant message", () => {
+    const messages = [
+      { role: "assistant", metadata: { usage: { totalTokens: 100 } } },
+      { role: "user", metadata: {} },
+      { role: "assistant", metadata: {} },
+    ];
+    const usage = getThreadMessageTokenUsage(messages[0]);
+
+    expect(usage).toEqual({ totalTokens: 100 });
+  });
+
+  it("returns token usage from the newest assistant message", () => {
+    const messages = [
+      { role: "assistant", metadata: { usage: { totalTokens: 100 } } },
+      {
+        role: "assistant",
+        metadata: { usage: { inputTokens: 40, outputTokens: 2 } },
+      },
+    ];
+    const usage = getThreadMessageTokenUsage(messages[1]);
+
+    expect(usage).toEqual({
+      totalTokens: 42,
+      inputTokens: 40,
+      outputTokens: 2,
+    });
+  });
+});

package/src/usage.ts

@@ -0,0 +1,179 @@
+/// <reference types="@ulearning/ai-core/vue" preserve="true" />
+import { computed, type ComputedRef } from "vue";
+import { useAuiState } from "@ulearning/ai-store";
+
+export type ThreadTokenUsage = {
+  totalTokens?: number;
+  inputTokens?: number;
+  outputTokens?: number;
+  reasoningTokens?: number;
+  cachedInputTokens?: number;
+};
+
+export interface TokenUsageExtractableMessage {
+  role?: string;
+  metadata?: unknown;
+}
+
+type UsageRecord = Record<string, unknown>;
+
+const USAGE_KEYS = [
+  "inputTokens",
+  "outputTokens",
+  "reasoningTokens",
+  "cachedInputTokens",
+  "totalTokens",
+] as const satisfies (keyof ThreadTokenUsage)[];
+
+function asRecord(value: unknown): UsageRecord | undefined {
+  if (!value || typeof value !== "object" || Array.isArray(value))
+    return undefined;
+  return value as UsageRecord;
+}
+
+function asPositiveTokenCount(value: unknown): number | undefined {
+  if (typeof value !== "number" || !Number.isFinite(value) || value < 0) {
+    return undefined;
+  }
+  return value;
+}
+
+function computeTotalTokens(usage: ThreadTokenUsage): number | undefined {
+  if (usage.totalTokens !== undefined) return usage.totalTokens;
+  if (usage.inputTokens !== undefined && usage.outputTokens !== undefined) {
+    return usage.inputTokens + usage.outputTokens;
+  }
+  return undefined;
+}
+
+function normalizeUsage(value: unknown): ThreadTokenUsage | undefined {
+  const record = asRecord(value);
+  if (!record) return undefined;
+
+  const result: ThreadTokenUsage = {};
+  let hasFields = false;
+  for (const key of USAGE_KEYS) {
+    const count = asPositiveTokenCount(record[key]);
+    if (count !== undefined) {
+      result[key] = count;
+      hasFields = true;
+    }
+  }
+  // AI SDK v7 moved these under token detail objects; v6 kept them top-level.
+  if (result.reasoningTokens === undefined) {
+    const count = asPositiveTokenCount(
+      asRecord(record.outputTokenDetails)?.reasoningTokens,
+    );
+    if (count !== undefined) {
+      result.reasoningTokens = count;
+      hasFields = true;
+    }
+  }
+  if (result.cachedInputTokens === undefined) {
+    const count = asPositiveTokenCount(
+      asRecord(record.inputTokenDetails)?.cacheReadTokens,
+    );
+    if (count !== undefined) {
+      result.cachedInputTokens = count;
+      hasFields = true;
+    }
+  }
+  return hasFields ? result : undefined;
+}
+
+function withComputedTotal(
+  usage: ThreadTokenUsage,
+): ThreadTokenUsage | undefined {
+  const totalTokens = computeTotalTokens(usage);
+  return { ...usage, ...(totalTokens !== undefined && { totalTokens }) };
+}
+
+function usageFromSteps(value: unknown): ThreadTokenUsage | undefined {
+  const steps = Array.isArray(value) ? value : [];
+
+  const sums: Record<string, number> = {};
+  const present: Record<string, boolean> = {};
+  let stepsWithUsage = 0;
+  let stepsWithComputableTotal = 0;
+
+  for (const step of steps) {
+    const usage = normalizeUsage(asRecord(step)?.usage);
+    if (!usage) continue;
+    stepsWithUsage++;
+
+    const stepTotal = computeTotalTokens(usage);
+    if (stepTotal !== undefined) {
+      sums.totalTokens = (sums.totalTokens ?? 0) + stepTotal;
+      stepsWithComputableTotal++;
+    }
+
+    for (const key of USAGE_KEYS) {
+      if (key === "totalTokens") continue;
+      if (usage[key] !== undefined) {
+        sums[key] = (sums[key] ?? 0) + usage[key];
+        present[key] = true;
+      }
+    }
+  }
+
+  if (stepsWithUsage === 0) return undefined;
+
+  const result: ThreadTokenUsage = {};
+  if (stepsWithComputableTotal === stepsWithUsage) {
+    result.totalTokens = sums.totalTokens!;
+  }
+  for (const key of USAGE_KEYS) {
+    if (key === "totalTokens") continue;
+    if (present[key]) {
+      result[key] = sums[key]!;
+    }
+  }
+  return result;
+}
+
+export function getThreadMessageTokenUsage(
+  message: TokenUsageExtractableMessage | undefined,
+): ThreadTokenUsage | undefined {
+  if (!message || message.role !== "assistant") return undefined;
+
+  const metadata = asRecord(message.metadata);
+  if (!metadata) return undefined;
+
+  const topLevelUsage = normalizeUsage(metadata.usage);
+  if (topLevelUsage) return withComputedTotal(topLevelUsage);
+
+  const customUsage = normalizeUsage(asRecord(metadata.custom)?.usage);
+  if (customUsage) return withComputedTotal(customUsage);
+
+  return usageFromSteps(metadata.steps);
+}
+
+function findLatestMessageWithUsage(
+  messages: readonly TokenUsageExtractableMessage[] | undefined,
+): TokenUsageExtractableMessage | undefined {
+  if (!messages) return undefined;
+
+  for (let idx = messages.length - 1; idx >= 0; idx -= 1) {
+    const message = messages[idx];
+    if (getThreadMessageTokenUsage(message)) {
+      return message;
+    }
+  }
+
+  return undefined;
+}
+
+/**
+ * Reads token usage from the newest assistant message that reports any.
+ *
+ * A route attaches usage through the AI SDK's `messageMetadata` option. Because
+ * a thread message carries a fixed metadata shape, the converter moves every
+ * other key the route returns into `metadata.custom`, which is where this hook
+ * looks.
+ */
+export function useThreadTokenUsage(): ComputedRef<
+  ThreadTokenUsage | undefined
+> {
+  const msg = useAuiState((s) => findLatestMessageWithUsage(s.thread.messages));
+  return computed(() => getThreadMessageTokenUsage(msg.value));
+}

package/src/utils/getVercelAIMessages.ts

@@ -0,0 +1,11 @@
+import {
+  getExternalStoreMessages,
+  type ThreadMessage,
+} from "@ulearning/ai-core";
+import type { UIMessage } from "ai";
+
+export const getVercelAIMessages = <UI_MESSAGE extends UIMessage = UIMessage>(
+  message: ThreadMessage,
+) => {
+  return getExternalStoreMessages(message) as UI_MESSAGE[];
+};

package/src/utils/sliceMessagesUntil.test.ts

@@ -0,0 +1,51 @@
+import { describe, expect, it } from "vitest";
+import type { UIMessage } from "ai";
+import { sliceMessagesUntil } from "./sliceMessagesUntil";
+
+const msg = (id: string, role: UIMessage["role"]): UIMessage =>
+  ({ id, role, parts: [] }) as any;
+
+describe("sliceMessagesUntil", () => {
+  it("returns an empty array when the message id is null", () => {
+    expect(
+      sliceMessagesUntil([msg("u1", "user"), msg("a1", "assistant")], null),
+    ).toEqual([]);
+  });
+
+  it("throws when the message id is not found", () => {
+    expect(() => sliceMessagesUntil([msg("u1", "user")], "missing")).toThrow(
+      'sliceMessagesUntil: Message "missing" not found in AI SDK messages. This is likely an internal bug in assistant-ui.',
+    );
+  });
+
+  it("slices up to and including the target when no assistant follows", () => {
+    const messages = [
+      msg("u1", "user"),
+      msg("a1", "assistant"),
+      msg("u2", "user"),
+    ];
+    expect(sliceMessagesUntil(messages, "u2").map((m) => m.id)).toEqual([
+      "u1",
+      "a1",
+      "u2",
+    ]);
+  });
+
+  it("stops at the target even when assistant messages follow it", () => {
+    const messages = [
+      msg("u1", "user"),
+      msg("a1", "assistant"),
+      msg("a2", "assistant"),
+      msg("u2", "user"),
+    ];
+    expect(sliceMessagesUntil(messages, "u1").map((m) => m.id)).toEqual(["u1"]);
+  });
+
+  it("returns the whole array when the target is the last message", () => {
+    const messages = [msg("u1", "user"), msg("a1", "assistant")];
+    expect(sliceMessagesUntil(messages, "a1").map((m) => m.id)).toEqual([
+      "u1",
+      "a1",
+    ]);
+  });
+});

package/src/utils/sliceMessagesUntil.ts

@@ -0,0 +1,16 @@
+import type { UIMessage } from "ai";
+
+export const sliceMessagesUntil = <UI_MESSAGE extends UIMessage = UIMessage>(
+  messages: UI_MESSAGE[],
+  messageId: string | null,
+) => {
+  if (messageId == null) return [];
+
+  const messageIdx = messages.findIndex((m) => m.id === messageId);
+  if (messageIdx === -1)
+    throw new Error(
+      `sliceMessagesUntil: Message "${messageId}" not found in AI SDK messages. This is likely an internal bug in assistant-ui.`,
+    );
+
+  return messages.slice(0, messageIdx + 1);
+};