Merge pull request #2914 from pocharlies/fix/context-usage-server-total

fix(ui): stop the context meter from counting every internal round-trip

packages/ui/src/components/chat/work-status/contextUsage.test.ts

@@ -14,8 +14,8 @@ describe('computeContextUsage', () => {
   });
 
   test('reports the latest turn rather than a sum across turns', () => {
-    // Each assistant turn reports the whole window it saw, so adding them up
-    // would report several times the real fill.
+    // A turn's tokens describe that turn's window, so adding turns up would
+    // report several times the real fill.
     const usage = computeContextUsage(
       [
         assistant({ input: 400, output: 0, reasoning: 0 }, 'old'),
@@ -61,4 +61,24 @@ describe('computeContextUsage', () => {
     const usage = computeContextUsage([assistant({ input: 10 })], 100);
     expect(usage?.totalTokens).toBe(10);
   });
+
+  test('prefers the server-reported total over summing round-trip fields', () => {
+    // Real payload from opencode 1.18.18: ~14 tool-call round-trips accumulated
+    // cache.read to 3.29M while the 1M window really held 232,872. Summing
+    // rendered 330.6%; the reported total renders the real 23.3%.
+    const usage = computeContextUsage(
+      [assistant({ total: 232_872, input: 0, output: 14_523, reasoning: 0, cache: { read: 3_291_956, write: 0 } })],
+      1_000_000,
+    );
+    expect(usage?.totalTokens).toBe(232_872);
+    expect(usage?.percent.toFixed(4)).toBe('23.2872');
+  });
+
+  test('selects a message whose only signal is the reported total', () => {
+    const usage = computeContextUsage(
+      [assistant({ total: 5_000, input: 0, output: 0, reasoning: 0 })],
+      100_000,
+    );
+    expect(usage?.totalTokens).toBe(5_000);
+  });
 });

packages/ui/src/components/chat/work-status/contextUsage.ts

@@ -13,7 +13,11 @@
  * global read to race with.
  */
 
+import { contextTokensFromBreakdown } from '@/stores/utils/tokenUtils';
+
 type MessageTokens = {
+  /** Server-reported window of the turn's final round-trip; absent on older servers. */
+  total?: number;
   input?: number;
   output?: number;
   reasoning?: number;
@@ -37,18 +41,12 @@ type WorkStatusContextUsage = {
 /** The store's own fallback when a model exposes no context limit. */
 export const DEFAULT_CONTEXT_LIMIT = 200_000;
 
-const sumTokens = (tokens: MessageTokens): number => (
-  (tokens.input ?? 0)
-  + (tokens.output ?? 0)
-  + (tokens.reasoning ?? 0)
-  + (tokens.cache?.read ?? 0)
-  + (tokens.cache?.write ?? 0)
-);
-
 /**
  * Usage from the newest assistant message that reported a non-zero token count.
- * Each assistant turn reports the whole window it saw, so the latest one is the
- * current fill — not a sum across turns.
+ * The latest turn describes the current fill — not a sum across turns. Within
+ * a turn, the server-reported `total` is the final round-trip's window;
+ * summing the breakdown fields instead overstates multi-step turns, whose
+ * input/cache fields accumulate across round-trips.
  */
 export const computeContextUsage = (
   messages: readonly MessageLike[],
@@ -60,7 +58,7 @@ export const computeContextUsage = (
     const message = messages[index];
     if (message?.role !== 'assistant' || !message.tokens) continue;
 
-    const totalTokens = sumTokens(message.tokens);
+    const totalTokens = contextTokensFromBreakdown(message.tokens);
     if (totalTokens <= 0) continue;
 
     const limit = contextLimit > 0 ? contextLimit : DEFAULT_CONTEXT_LIMIT;