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. */