From f02fe83e2848a930e07c2340a29b776f28dde6ec Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E7=A8=8B=E5=BA=8F=E5=91=98=E9=98=BF=E6=B1=9F=28Relakkes?=
 =?UTF-8?q?=29?= <relakkes@gmail.com>
Date: Tue, 5 May 2026 19:42:06 +0800
Subject: [PATCH] fix: keep context usage from dropping after responses

Context inspection was switching from the local context estimate to the latest
provider usage total once a response completed. That provider total omitted the
assistant output tokens, even though those tokens become part of the next turn's
context, so the desktop meter could fall from 11% to 10% after work finished.
The context total now includes output tokens and keeps the local estimate as a
lower bound.

Constraint: Provider usage and local context estimates use different token accounting paths
Rejected: Trust provider input tokens alone | it omits the latest assistant output from the next-turn context
Confidence: high
Scope-risk: narrow
Directive: Context usage totals must represent the next-turn conversation context, not billing-only input tokens
Tested: bun test src/utils/__tests__/context.test.ts
Tested: bun run check:server
---
 src/utils/__tests__/context.test.ts | 22 +++++++++++++++++++++
 src/utils/analyzeContext.ts         | 26 ++++++++++++-------------
 src/utils/context.ts                | 30 +++++++++++++++++++++++++++++
 3 files changed, 65 insertions(+), 13 deletions(-)
 create mode 100644 src/utils/__tests__/context.test.ts

diff --git a/src/utils/__tests__/context.test.ts b/src/utils/__tests__/context.test.ts
new file mode 100644
index 00000000..896de453
--- /dev/null
+++ b/src/utils/__tests__/context.test.ts
@@ -0,0 +1,22 @@
+import { describe, expect, test } from 'bun:test'
+import { calculateCurrentContextTokenTotal } from '../context.js'
+
+describe('calculateCurrentContextTokenTotal', () => {
+  test('includes assistant output tokens in the current context total', () => {
+    expect(calculateCurrentContextTokenTotal(26_000, {
+      input_tokens: 24_000,
+      cache_creation_input_tokens: 1_000,
+      cache_read_input_tokens: 1_000,
+      output_tokens: 3_000,
+    })).toBe(29_000)
+  })
+
+  test('keeps the local estimate as a lower bound when provider usage is smaller', () => {
+    expect(calculateCurrentContextTokenTotal(29_000, {
+      input_tokens: 24_000,
+      cache_creation_input_tokens: 1_000,
+      cache_read_input_tokens: 1_000,
+      output_tokens: 0,
+    })).toBe(29_000)
+  })
+})
diff --git a/src/utils/analyzeContext.ts b/src/utils/analyzeContext.ts
index 1631f7a4..0db8a4ca 100644
--- a/src/utils/analyzeContext.ts
+++ b/src/utils/analyzeContext.ts
@@ -48,7 +48,10 @@ import type {
 } from '../types/message.js'
 import { toolToAPISchema } from './api.js'
 import { filterInjectedMemoryFiles, getMemoryFiles } from './claudemd.js'
-import { getContextWindowForModel } from './context.js'
+import {
+  calculateCurrentContextTokenTotal,
+  getContextWindowForModel,
+} from './context.js'
 import { getCwd } from './cwd.js'
 import { logForDebugging } from './debug.js'
 import { isEnvTruthy } from './envUtils.js'
@@ -1192,20 +1195,17 @@ export async function analyzeContextUsage(
   // Total for display (everything except free space)
   const totalIncludingReserved = actualUsage
 
-  // Extract API usage from original messages (if provided) to match status line
-  // This uses the same source of truth as the status line for consistency
+  // Extract API usage from original messages (if provided) to anchor the
+  // estimate to the latest provider-reported request size.
   const apiUsage = getCurrentUsage(originalMessages ?? messages)
 
-  // When API usage is available, use it for total to match status line calculation
-  // Status line uses: input_tokens + cache_creation_input_tokens + cache_read_input_tokens
-  const totalFromAPI = apiUsage
-    ? apiUsage.input_tokens +
-      apiUsage.cache_creation_input_tokens +
-      apiUsage.cache_read_input_tokens
-    : null
-
-  // Use API total if available, otherwise fall back to estimated total
-  const finalTotalTokens = totalFromAPI ?? totalIncludingReserved
+  // Use the larger of the local estimate and the latest API-reported context.
+  // This keeps the context meter from dropping when a response completes and
+  // the output tokens become part of the next turn's context.
+  const finalTotalTokens = calculateCurrentContextTokenTotal(
+    totalIncludingReserved,
+    apiUsage,
+  )
 
   // Pre-calculate grid based on model context window and terminal width
   // For narrow screens (< 80 cols), use 5x5 for 200k models, 5x10 for 1M+ models
diff --git a/src/utils/context.ts b/src/utils/context.ts
index 0cce9b58..87fae9df 100644
--- a/src/utils/context.ts
+++ b/src/utils/context.ts
@@ -161,6 +161,36 @@ export function calculateContextPercentages(
   }
 }
 
+/**
+ * Calculate the current context size after the latest assistant response.
+ *
+ * API usage reports the prompt tokens used for the just-finished request plus
+ * that request's output tokens. The output becomes part of the next request's
+ * conversation context, so omitting it can make context usage appear to drop
+ * immediately after the model finishes responding. The local estimate is kept
+ * as a lower bound because it includes system/tool/message material that some
+ * provider usage payloads under-report.
+ */
+export function calculateCurrentContextTokenTotal(
+  estimatedTokens: number,
+  currentUsage: {
+    input_tokens: number
+    output_tokens?: number
+    cache_creation_input_tokens: number
+    cache_read_input_tokens: number
+  } | null,
+): number {
+  if (!currentUsage) return estimatedTokens
+
+  const totalFromAPI =
+    currentUsage.input_tokens +
+    currentUsage.cache_creation_input_tokens +
+    currentUsage.cache_read_input_tokens +
+    (currentUsage.output_tokens ?? 0)
+
+  return Math.max(estimatedTokens, totalFromAPI)
+}
+
 /**
  * Returns the model's default and upper limit for max output tokens.
  */