fix: use backend thread token usage for header total (#2800)

* fix: use backend thread token usage for header total * Refactor thread token usage fetch
2026-05-24 08:55:59 +00:00 · 2026-05-09 19:40:32 +08:00
parent 881ff71252
commit 417416087b
16 changed files with 540 additions and 35 deletions
@@ -65,6 +65,40 @@ export function accumulateUsage(messages: Message[]): TokenUsage | null {
  return hasUsage ? cumulative : null;
 }

+function hasNonZeroUsage(
+  usage: TokenUsage | null | undefined,
+): usage is TokenUsage {
+  return (
+    usage !== null &&
+    usage !== undefined &&
+    (usage.inputTokens > 0 || usage.outputTokens > 0 || usage.totalTokens > 0)
+  );
+}
+
+function addUsage(base: TokenUsage, delta: TokenUsage): TokenUsage {
+  return {
+    inputTokens: base.inputTokens + delta.inputTokens,
+    outputTokens: base.outputTokens + delta.outputTokens,
+    totalTokens: base.totalTokens + delta.totalTokens,
+  };
+}
+
+export function selectHeaderTokenUsage({
+  backendUsage,
+  messages,
+  pendingMessages = [],
+}: {
+  backendUsage?: TokenUsage | null;
+  messages: Message[];
+  pendingMessages?: Message[];
+}): TokenUsage | null {
+  if (hasNonZeroUsage(backendUsage)) {
+    const pendingUsage = accumulateUsage(pendingMessages);
+    return pendingUsage ? addUsage(backendUsage, pendingUsage) : backendUsage;
+  }
+  return accumulateUsage(messages);
+}
+
 /**
 * Format a token count for display: 1234 -> "1,234", 12345 -> "12.3K"
 */