@assistant-ui/ai-sdk 0.0.1

package/src/usage.ts

@@ -0,0 +1,174 @@
+/// <reference types="@assistant-ui/core/react" />
+import { useAuiState } from "@assistant-ui/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 legacyUsage = normalizeUsage(asRecord(metadata.custom)?.usage);
+  if (legacyUsage) return withComputedTotal(legacyUsage);
+
+  return usageFromSteps(metadata.steps);
+}
+
+export function getLatestThreadTokenUsage(
+  messages: readonly TokenUsageExtractableMessage[] | undefined,
+): ThreadTokenUsage | undefined {
+  return getThreadMessageTokenUsage(findLatestMessageWithUsage(messages));
+}
+
+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;
+}
+
+export function useThreadTokenUsage(): ThreadTokenUsage | undefined {
+  const msg = useAuiState((s) => findLatestMessageWithUsage(s.thread.messages));
+  return getThreadMessageTokenUsage(msg);
+}

package/src/utils/getVercelAIMessages.ts

@@ -0,0 +1,11 @@
+import {
+  getExternalStoreMessages,
+  type ThreadMessage,
+} from "@assistant-ui/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,55 @@
+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("extends the slice to include trailing assistant messages", () => {
+    const messages = [
+      msg("u1", "user"),
+      msg("a1", "assistant"),
+      msg("a2", "assistant"),
+      msg("u2", "user"),
+    ];
+    expect(sliceMessagesUntil(messages, "u1").map((m) => m.id)).toEqual([
+      "u1",
+      "a1",
+      "a2",
+    ]);
+  });
+
+  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,20 @@
+import type { UIMessage } from "ai";
+
+export const sliceMessagesUntil = <UI_MESSAGE extends UIMessage = UIMessage>(
+  messages: UI_MESSAGE[],
+  messageId: string | null,
+) => {
+  if (messageId == null) return [];
+
+  let 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.`,
+    );
+
+  while (messages[messageIdx + 1]?.role === "assistant") {
+    messageIdx++;
+  }
+
+  return messages.slice(0, messageIdx + 1);
+};