@assistant-ui/ai-sdk 0.0.3 → 0.0.5

package/src/usage.ts

@@ -1,4 +1,5 @@
 /// <reference types="@assistant-ui/core/react" />
+import { useMemo } from "react";
 import { useAuiState } from "@assistant-ui/store";
 
 export type ThreadTokenUsage = {
@@ -141,18 +142,12 @@ export function getThreadMessageTokenUsage(
   const topLevelUsage = normalizeUsage(metadata.usage);
   if (topLevelUsage) return withComputedTotal(topLevelUsage);
 
-  const legacyUsage = normalizeUsage(asRecord(metadata.custom)?.usage);
-  if (legacyUsage) return withComputedTotal(legacyUsage);
+  const customUsage = normalizeUsage(asRecord(metadata.custom)?.usage);
+  if (customUsage) return withComputedTotal(customUsage);
 
   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 {
@@ -168,7 +163,15 @@ function findLatestMessageWithUsage(
   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(): ThreadTokenUsage | undefined {
   const msg = useAuiState((s) => findLatestMessageWithUsage(s.thread.messages));
-  return getThreadMessageTokenUsage(msg);
+  return useMemo(() => getThreadMessageTokenUsage(msg), [msg]);
 }