fix(chat): stop the context-usage percentage flickering during streaming

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 {

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%');
+  });
 });