From 8420e284455c59d73551e8531466ef8ac6de81b2 Mon Sep 17 00:00:00 2001 From: siolyn Date: Mon, 8 Jun 2026 23:47:09 +0800 Subject: [PATCH 1/2] feat: show remaining context headroom in TUI badge --- src/tui/chrome.ts | 18 +++++++++++++- src/utils/token-estimator.ts | 3 +++ test/context-badge.test.ts | 48 ++++++++++++++++++++++++++++++++++++ test/token-estimator.test.ts | 3 +++ 4 files changed, 71 insertions(+), 1 deletion(-) diff --git a/src/tui/chrome.ts b/src/tui/chrome.ts index 74693bf..bae6b57 100644 --- a/src/tui/chrome.ts +++ b/src/tui/chrome.ts @@ -114,6 +114,16 @@ function colorBadge( return `${color}[${label}]${RESET} ${BOLD}${value}${RESET}` } +function formatCompactTokenCount(tokens: number): string { + const value = Math.max(0, Math.floor(tokens)) + if (value < 1_000) return String(value) + if (value < 1_000_000) return `${Math.round(value / 1_000)}K` + + const millions = value / 1_000_000 + const formatted = millions.toFixed(1).replace(/\.0$/, '') + return `${formatted}M` +} + function joinSegmentsWithinWidth( segments: string[], separator: string, @@ -232,6 +242,7 @@ export function renderPanel( export function renderContextBadge(stats: { utilization: number warningLevel: 'normal' | 'warning' | 'critical' | 'blocked' + remainingTokens?: number accounting?: { providerUsageTokens: number estimatedTokens: number @@ -251,6 +262,10 @@ export function renderContextBadge(stats: { const filled = Math.round(utilization * 10) const bar = '\u2593'.repeat(filled) + '\u2591'.repeat(10 - filled) + const headroom = + stats.remainingTokens !== undefined + ? ` ${formatCompactTokenCount(stats.remainingTokens)} left` + : '' const sourceLabel = accounting?.source === 'provider_usage' ? 'usage' @@ -261,7 +276,7 @@ export function renderContextBadge(stats: { : '' const suffix = sourceLabel ? ` ${sourceLabel}` : '' - return colorBadge('ctx', `${percent}% ${bar}${suffix}`, color) + return colorBadge('ctx', `${percent}% ${bar}${headroom}${suffix}`, color) } export function renderBanner( @@ -279,6 +294,7 @@ export function renderBanner( contextStats?: { utilization: number warningLevel: 'normal' | 'warning' | 'critical' | 'blocked' + remainingTokens?: number accounting?: { providerUsageTokens: number estimatedTokens: number diff --git a/src/utils/token-estimator.ts b/src/utils/token-estimator.ts index d1ded6d..6fedf2e 100644 --- a/src/utils/token-estimator.ts +++ b/src/utils/token-estimator.ts @@ -26,6 +26,7 @@ export type ContextStats = { providerUsageTokens: number contextWindow: number effectiveInput: number + remainingTokens: number utilization: number warningLevel: 'normal' | 'warning' | 'critical' | 'blocked' accounting: TokenAccountingResult @@ -181,6 +182,7 @@ export function computeContextStats( const window = getModelContextWindow(model) const accounting = tokenCountWithEstimation(messages) const utilization = Math.min(1, accounting.totalTokens / window.effectiveInput) + const remainingTokens = Math.max(0, window.effectiveInput - accounting.totalTokens) let warningLevel: ContextStats['warningLevel'] if (utilization >= 0.95) { @@ -199,6 +201,7 @@ export function computeContextStats( providerUsageTokens: accounting.providerUsageTokens, contextWindow: window.contextWindow, effectiveInput: window.effectiveInput, + remainingTokens, utilization, warningLevel, accounting, diff --git a/test/context-badge.test.ts b/test/context-badge.test.ts index 6279d92..6f7dc03 100644 --- a/test/context-badge.test.ts +++ b/test/context-badge.test.ts @@ -51,6 +51,7 @@ describe('renderContextBadge', () => { const result = renderContextBadge({ utilization: 0.82, warningLevel: 'warning', + remainingTokens: 18_000, accounting: { providerUsageTokens: 70_000, estimatedTokens: 12_000, @@ -59,9 +60,56 @@ describe('renderContextBadge', () => { }) const plain = stripAnsi(result) assert.ok(plain.includes('82%')) + assert.ok(plain.includes('18K left')) assert.ok(plain.includes('usage+est')) }) + it('shows compact remaining context headroom', () => { + const result = renderContextBadge({ + utilization: 0.68, + warningLevel: 'warning', + remainingTokens: 59_200, + }) + const plain = stripAnsi(result) + assert.ok(plain.includes('59K left')) + }) + + it('shows small remaining context headroom without a suffix', () => { + const result = renderContextBadge({ + utilization: 0.94, + warningLevel: 'critical', + remainingTokens: 999, + }) + const plain = stripAnsi(result) + assert.ok(plain.includes('999 left')) + }) + + it('shows million-scale remaining context headroom compactly', () => { + const result = renderContextBadge({ + utilization: 0.12, + warningLevel: 'normal', + remainingTokens: 1_234_000, + }) + const plain = stripAnsi(result) + assert.ok(plain.includes('1.2M left')) + }) + + it('shows zero remaining headroom when context is blocked', () => { + const result = renderContextBadge({ + utilization: 1, + warningLevel: 'blocked', + remainingTokens: 0, + }) + const plain = stripAnsi(result) + assert.ok(plain.includes('0 left')) + }) + + it('keeps the previous badge format when headroom is unavailable', () => { + const result = renderContextBadge({ utilization: 0.5, warningLevel: 'warning' }) + const plain = stripAnsi(result) + assert.ok(!plain.includes('left')) + }) + it('uses correct block characters for utilization', () => { const result = renderContextBadge({ utilization: 0.5, warningLevel: 'warning' }) const plain = stripAnsi(result) diff --git a/test/token-estimator.test.ts b/test/token-estimator.test.ts index 31070c2..11e4e69 100644 --- a/test/token-estimator.test.ts +++ b/test/token-estimator.test.ts @@ -97,6 +97,7 @@ describe('computeContextStats', () => { assert.ok(stats.estimatedTokens > 0) assert.equal(stats.contextWindow, 200_000) assert.equal(stats.effectiveInput, 184_000) + assert.equal(stats.remainingTokens, stats.effectiveInput - stats.totalTokens) }) it('computes blocked warning level for large messages', () => { @@ -109,6 +110,7 @@ describe('computeContextStats', () => { `expected blocked or critical, got ${stats.warningLevel}`, ) assert.equal(stats.utilization, 1, 'utilization should be capped at 1') + assert.equal(stats.remainingTokens, 0) }) it('computes warning level for medium messages', () => { @@ -133,6 +135,7 @@ describe('computeContextStats', () => { ] const stats = computeContextStats(messages, 'deepseek-chat') assert.equal(stats.utilization, 1) + assert.equal(stats.remainingTokens, 0) }) }) From 39125cdf892c1f72ba915f9698dd0d3c126ecb1d Mon Sep 17 00:00:00 2001 From: siolyn Date: Mon, 8 Jun 2026 23:56:58 +0800 Subject: [PATCH 2/2] test: tighten context headroom badge coverage --- src/tui/chrome.ts | 28 ++++++++---------- test/context-badge.test.ts | 58 +++++++++++++++++++++++++++++--------- test/session.test.ts | 1 + test/snip-compact.test.ts | 1 + 4 files changed, 59 insertions(+), 29 deletions(-) diff --git a/src/tui/chrome.ts b/src/tui/chrome.ts index bae6b57..fe062c5 100644 --- a/src/tui/chrome.ts +++ b/src/tui/chrome.ts @@ -4,6 +4,7 @@ import process from 'node:process' import type { RuntimeConfig } from '../config.js' import type { SlashCommand } from '../cli-commands.js' import type { PermissionRequest } from '../permissions.js' +import type { ContextStats } from '../utils/token-estimator.js' const RESET = '\u001b[0m' const DIM = '\u001b[2m' @@ -117,8 +118,15 @@ function colorBadge( function formatCompactTokenCount(tokens: number): string { const value = Math.max(0, Math.floor(tokens)) if (value < 1_000) return String(value) - if (value < 1_000_000) return `${Math.round(value / 1_000)}K` + if (value < 1_000_000) { + const thousands = Math.round(value / 1_000) + return thousands >= 1_000 ? formatCompactMillionTokenCount(value) : `${thousands}K` + } + + return formatCompactMillionTokenCount(value) +} +function formatCompactMillionTokenCount(value: number): string { const millions = value / 1_000_000 const formatted = millions.toFixed(1).replace(/\.0$/, '') return `${formatted}M` @@ -242,7 +250,7 @@ export function renderPanel( export function renderContextBadge(stats: { utilization: number warningLevel: 'normal' | 'warning' | 'critical' | 'blocked' - remainingTokens?: number + remainingTokens: number accounting?: { providerUsageTokens: number estimatedTokens: number @@ -262,10 +270,7 @@ export function renderContextBadge(stats: { const filled = Math.round(utilization * 10) const bar = '\u2593'.repeat(filled) + '\u2591'.repeat(10 - filled) - const headroom = - stats.remainingTokens !== undefined - ? ` ${formatCompactTokenCount(stats.remainingTokens)} left` - : '' + const headroom = ` ${formatCompactTokenCount(stats.remainingTokens)} left` const sourceLabel = accounting?.source === 'provider_usage' ? 'usage' @@ -291,16 +296,7 @@ export function renderBanner( mcpConnectedCount: number mcpConnectingCount: number mcpErrorCount: number - contextStats?: { - utilization: number - warningLevel: 'normal' | 'warning' | 'critical' | 'blocked' - remainingTokens?: number - accounting?: { - providerUsageTokens: number - estimatedTokens: number - source: 'provider_usage' | 'provider_usage_plus_estimate' | 'estimate_only' - } - } | null + contextStats?: ContextStats | null }, ): string { const panelWidth = Math.max(60, process.stdout.columns ?? 100) diff --git a/test/context-badge.test.ts b/test/context-badge.test.ts index 6f7dc03..58123da 100644 --- a/test/context-badge.test.ts +++ b/test/context-badge.test.ts @@ -8,7 +8,11 @@ function stripAnsi(input: string): string { describe('renderContextBadge', () => { it('renders normal level badge', () => { - const result = renderContextBadge({ utilization: 0.23, warningLevel: 'normal' }) + const result = renderContextBadge({ + utilization: 0.23, + warningLevel: 'normal', + remainingTokens: 142_000, + }) const plain = stripAnsi(result) assert.ok(plain.includes('ctx'), 'should contain ctx label') assert.ok(plain.includes('23%'), 'should show 23%') @@ -17,32 +21,52 @@ describe('renderContextBadge', () => { }) it('renders warning level badge', () => { - const result = renderContextBadge({ utilization: 0.68, warningLevel: 'warning' }) + const result = renderContextBadge({ + utilization: 0.68, + warningLevel: 'warning', + remainingTokens: 59_000, + }) const plain = stripAnsi(result) assert.ok(plain.includes('68%')) assert.ok(plain.includes('ctx')) }) it('renders critical level badge', () => { - const result = renderContextBadge({ utilization: 0.89, warningLevel: 'critical' }) + const result = renderContextBadge({ + utilization: 0.89, + warningLevel: 'critical', + remainingTokens: 20_000, + }) const plain = stripAnsi(result) assert.ok(plain.includes('89%')) }) it('renders blocked level badge', () => { - const result = renderContextBadge({ utilization: 0.96, warningLevel: 'blocked' }) + const result = renderContextBadge({ + utilization: 0.96, + warningLevel: 'blocked', + remainingTokens: 0, + }) const plain = stripAnsi(result) assert.ok(plain.includes('96%')) }) it('renders 0% utilization correctly', () => { - const result = renderContextBadge({ utilization: 0, warningLevel: 'normal' }) + const result = renderContextBadge({ + utilization: 0, + warningLevel: 'normal', + remainingTokens: 184_000, + }) const plain = stripAnsi(result) assert.ok(plain.includes('0%')) }) it('renders 100% utilization correctly', () => { - const result = renderContextBadge({ utilization: 1, warningLevel: 'blocked' }) + const result = renderContextBadge({ + utilization: 1, + warningLevel: 'blocked', + remainingTokens: 0, + }) const plain = stripAnsi(result) assert.ok(plain.includes('100%')) }) @@ -94,6 +118,16 @@ describe('renderContextBadge', () => { assert.ok(plain.includes('1.2M left')) }) + it('promotes rounded 1000K headroom to million-scale display', () => { + const result = renderContextBadge({ + utilization: 0.12, + warningLevel: 'normal', + remainingTokens: 999_500, + }) + const plain = stripAnsi(result) + assert.ok(plain.includes('1M left')) + }) + it('shows zero remaining headroom when context is blocked', () => { const result = renderContextBadge({ utilization: 1, @@ -104,14 +138,12 @@ describe('renderContextBadge', () => { assert.ok(plain.includes('0 left')) }) - it('keeps the previous badge format when headroom is unavailable', () => { - const result = renderContextBadge({ utilization: 0.5, warningLevel: 'warning' }) - const plain = stripAnsi(result) - assert.ok(!plain.includes('left')) - }) - it('uses correct block characters for utilization', () => { - const result = renderContextBadge({ utilization: 0.5, warningLevel: 'warning' }) + const result = renderContextBadge({ + utilization: 0.5, + warningLevel: 'warning', + remainingTokens: 92_000, + }) const plain = stripAnsi(result) // 50% → 5 filled blocks out of 10 const filledCount = (plain.match(/\u2593/g) || []).length diff --git a/test/session.test.ts b/test/session.test.ts index 627c8ea..1c86681 100644 --- a/test/session.test.ts +++ b/test/session.test.ts @@ -55,6 +55,7 @@ function contextStats(messages: ChatMessage[], effectiveInput = 20_000): Context providerUsageTokens: accounting.providerUsageTokens, contextWindow: effectiveInput, effectiveInput, + remainingTokens: Math.max(0, effectiveInput - accounting.totalTokens), utilization, warningLevel: utilization >= 0.95 diff --git a/test/snip-compact.test.ts b/test/snip-compact.test.ts index b243785..aa5f178 100644 --- a/test/snip-compact.test.ts +++ b/test/snip-compact.test.ts @@ -24,6 +24,7 @@ function contextStats(messages: ChatMessage[], effectiveInput = 20_000): Context providerUsageTokens: 0, contextWindow: effectiveInput, effectiveInput, + remainingTokens: Math.max(0, effectiveInput - totalTokens), utilization, warningLevel: utilization >= 0.95