From 78b14f5d05d60f33463fed0071123c2f15305b32 Mon Sep 17 00:00:00 2001 From: setkyar Date: Thu, 17 Sep 2026 16:43:52 +0700 Subject: [PATCH] fix(chat): stop the context-usage percentage flickering during streaming MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The indicator picked its context snapshot from the last assistant message carrying a `usage` object. A half-written streaming entry lands with an empty usage object, so the percentage dropped to 0% (or the capsule hid) until the next reload replaced it. - collectContextUsage now skips assistant messages whose usage totals zero, staying on the last message that actually measured the context. - updateContextUsage takes the snapshot it last rendered and returns the new one: a transient zero reuses the previous value, so once a value is on screen the capsule never flickers to 0% or hides. An empty model label likewise reuses the previous limit instead of the 128k default. - The controller holds the first paint until /api/models settles (and paints anyway if that fetch fails) so the denominator does not jump from the heuristic fallback to the registry value. Known limitation: the model label is not pinned the same way, so the first paint can still adjust once when worker status reports the model. Non-zero regressions are left alone on purpose — compaction legitimately drops the context. Fixes #110 --- .../components/session/chat/context-usage.js | 69 +++++++--- .../session/chat/context-usage.test.js | 119 ++++++++++++++++++ 2 files changed, 170 insertions(+), 18 deletions(-) diff --git a/web/src/components/session/chat/context-usage.js b/web/src/components/session/chat/context-usage.js index 12f2a3c5..5e493b58 100644 --- a/web/src/components/session/chat/context-usage.js +++ b/web/src/components/session/chat/context-usage.js @@ -79,13 +79,18 @@ export function collectContextUsage(entries = []) { const entry = entries[i]; if (entry?.type !== 'message' || !entry.message) continue; const msg = entry.message; - if (msg.role === 'assistant' && msg.usage) { - contextTokens = - msg.usage.totalTokens || - (msg.usage.input || 0) + - (msg.usage.output || 0) + - (msg.usage.cacheRead || 0) + - (msg.usage.cacheWrite || 0); + if (msg.role !== 'assistant' || !msg.usage) continue; + const total = + msg.usage.totalTokens || + (msg.usage.input || 0) + + (msg.usage.output || 0) + + (msg.usage.cacheRead || 0) + + (msg.usage.cacheWrite || 0); + // A streaming assistant message can land with an empty/zero usage object + // before the provider reports its totals. Skipping it keeps the indicator + // on the last message that actually measured the context. + if (total > 0) { + contextTokens = total; break; } } @@ -118,25 +123,35 @@ function formatLimit(n) { return n.toLocaleString(); } +// `previous` is the snapshot this function last rendered (the controller keeps +// it). Live updates can arrive with a half-written assistant entry or before the +// worker has reported a model, so anything missing falls back to the snapshot +// instead of flickering to 0% or hiding the indicator. Returns the new snapshot. export function updateContextUsage({ documentImpl = document, entries = [], knownModelLabel = '', contextWindows = {}, positionPopover = () => {}, + previous = null, } = {}) { const el = documentImpl.getElementById('pi-chat-context-usage'); - if (!el) return; + if (!el) return previous; const usage = collectContextUsage(entries); - if (usage.contextTokens <= 0 && usage.totalIOTokens <= 0) { + const contextTokens = + usage.contextTokens > 0 ? usage.contextTokens : previous?.contextTokens || 0; + if (contextTokens <= 0 && usage.totalIOTokens <= 0) { el.style.display = 'none'; - return; + return previous; } const { modelName, providerName } = splitModelLabel(knownModelLabel); - const limit = getModelContextLimit(modelName, providerName, contextWindows); - const percent = Math.min(100, Math.max(0, Math.round((usage.contextTokens / limit) * 100))); + const limit = + !modelName && previous?.limit + ? previous.limit + : getModelContextLimit(modelName, providerName, contextWindows); + const percent = Math.min(100, Math.max(0, Math.round((contextTokens / limit) * 100))); const fillPath = el.querySelector('.pi-context-fill'); const textSpan = el.querySelector('.pi-context-text'); @@ -147,7 +162,7 @@ export function updateContextUsage({ const formatNumber = (num) => num.toLocaleString(); el.setAttribute( 'title', - `Click for details (${formatNumber(usage.contextTokens)} / ${formatNumber(limit)} tokens used in context)`, + `Click for details (${formatNumber(contextTokens)} / ${formatNumber(limit)} tokens used in context)`, ); el.classList.remove('warning', 'danger'); @@ -171,7 +186,7 @@ export function updateContextUsage({ if (valOutput) valOutput.textContent = formatTokensDetail(usage.outputTokens); if (valTotal) valTotal.textContent = formatTokensDetail(usage.totalIOTokens); - if (usedSpan) usedSpan.textContent = formatTokensDetail(usage.contextTokens); + if (usedSpan) usedSpan.textContent = formatTokensDetail(contextTokens); if (limitSpan) limitSpan.textContent = formatLimit(limit); if (popoverBar) popoverBar.style.width = `${percent}%`; @@ -186,6 +201,8 @@ export function updateContextUsage({ } el.style.display = 'inline-flex'; + + return { contextTokens, limit }; } export function createContextUsageController({ @@ -196,17 +213,30 @@ export function createContextUsageController({ positionPopover = () => {}, } = {}) { let contextWindows = {}; + let modelsPending = false; + let lastShown = null; - const update = () => - updateContextUsage({ + const render = () => { + lastShown = updateContextUsage({ documentImpl, entries, knownModelLabel: getKnownModelLabel(), contextWindows, positionPopover, + previous: lastShown, }); + }; + + // Hold the first paint until /api/models settles: drawing against the + // heuristic fallback limit and then swapping in the registry value makes the + // percentage visibly jump. + const update = () => { + if (modelsPending) return; + render(); + }; if (chatApi && typeof chatApi.listModels === 'function') { + modelsPending = true; chatApi .listModels() .then((res) => { @@ -215,9 +245,12 @@ export function createContextUsageController({ }) .then((data) => { contextWindows = buildContextWindows(data.models || []); - update(); }) - .catch(() => {}); + .catch(() => {}) + .then(() => { + modelsPending = false; + render(); + }); } return { diff --git a/web/src/components/session/chat/context-usage.test.js b/web/src/components/session/chat/context-usage.test.js index 4580bbe0..6ad0484c 100644 --- a/web/src/components/session/chat/context-usage.test.js +++ b/web/src/components/session/chat/context-usage.test.js @@ -73,6 +73,22 @@ describe('context usage helpers', () => { expect(usage.totalIOTokens).toBe(4300); expect(usage.contextTokens).toBe(1800); }); + + it('ignores a streaming assistant entry that has no measured usage yet', () => { + const usage = collectContextUsage([ + { + type: 'message', + message: { role: 'assistant', usage: { input: 5000, output: 200 } }, + }, + { type: 'message', message: { role: 'user', content: 'follow-up' } }, + { + type: 'message', + message: { role: 'assistant', usage: { input: 0, output: 0, cacheRead: 0 } }, + }, + ]); + + expect(usage.contextTokens).toBe(5200); + }); }); describe('updateContextUsage', () => { @@ -148,4 +164,107 @@ describe('updateContextUsage', () => { expect(document.querySelector('.pi-popover-limit').textContent).toBe('1.2M'); }); + + it('holds the last value while streaming entries arrive without usage', () => { + renderDom(); + const entries = [ + { + type: 'message', + message: { role: 'assistant', usage: { input: 60000, output: 4000 } }, + }, + ]; + const controller = createContextUsageController({ + documentImpl: document, + entries, + getKnownModelLabel: () => 'gpt-4o @ openai', + }); + + controller.update(); + const el = document.getElementById('pi-chat-context-usage'); + expect(el.querySelector('.pi-context-text').textContent).toBe('50%'); + + // A half-written assistant entry (and then a tool result) must not drop the + // indicator to 0% or hide it. + entries.push({ type: 'message', message: { role: 'assistant', usage: {} } }); + controller.update(); + expect(el.querySelector('.pi-context-text').textContent).toBe('50%'); + expect(el.style.display).toBe('inline-flex'); + + entries.push({ type: 'message', message: { role: 'toolResult', content: 'ok' } }); + controller.update(); + expect(el.querySelector('.pi-context-text').textContent).toBe('50%'); + expect(el.style.display).toBe('inline-flex'); + + entries[1].message.usage = { input: 70000, output: 6000 }; + controller.update(); + expect(el.querySelector('.pi-context-text').textContent).toBe('59%'); + }); + + it('keeps the last limit when the model label is momentarily unknown', () => { + renderDom(); + const entries = [ + { type: 'message', message: { role: 'assistant', usage: { totalTokens: 100000 } } }, + ]; + let label = 'claude-sonnet-4 @ anthropic'; + const controller = createContextUsageController({ + documentImpl: document, + entries, + getKnownModelLabel: () => label, + }); + + controller.update(); + const text = document.querySelector('.pi-context-text'); + expect(text.textContent).toBe('50%'); + + label = ''; + controller.update(); + expect(text.textContent).toBe('50%'); + }); + + it('waits for the model registry before the first paint', async () => { + renderDom(); + let resolveModels; + const controller = createContextUsageController({ + documentImpl: document, + entries: [ + { type: 'message', message: { role: 'assistant', usage: { totalTokens: 100000 } } }, + ], + getKnownModelLabel: () => 'DEEPSEEK-V4-PRO @ DEEPSEEK', + chatApi: { + listModels: () => new Promise((resolve) => (resolveModels = resolve)), + }, + }); + + controller.update(); + const el = document.getElementById('pi-chat-context-usage'); + expect(el.style.display).toBe('none'); + + resolveModels({ + ok: true, + json: () => + Promise.resolve({ + models: [{ id: 'DEEPSEEK-V4-PRO', provider: 'DEEPSEEK', contextWindow: 1000000 }], + }), + }); + await new Promise((resolve) => setTimeout(resolve, 0)); + + expect(el.style.display).toBe('inline-flex'); + expect(el.querySelector('.pi-context-text').textContent).toBe('10%'); + }); + + it('still paints when the model registry fetch fails', async () => { + renderDom(); + createContextUsageController({ + documentImpl: document, + entries: [{ type: 'message', message: { role: 'assistant', usage: { totalTokens: 64000 } } }], + getKnownModelLabel: () => 'gpt-4o @ openai', + chatApi: { listModels: () => Promise.reject(new Error('offline')) }, + }); + + await new Promise((resolve) => setTimeout(resolve, 0)); + + const el = document.getElementById('pi-chat-context-usage'); + expect(el.style.display).toBe('inline-flex'); + expect(el.querySelector('.pi-context-text').textContent).toBe('50%'); + }); });