@pi-archimedes/footer 2.7.3 → 2.9.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@pi-archimedes/footer",
-  "version": "2.7.3",
+  "version": "2.9.0",
   "repository": {
     "type": "git",
     "url": "https://github.com/danielcherubini/pi-archimedes.git"
@@ -19,15 +19,15 @@
     "./config": "./src/config.ts"
   },
   "dependencies": {
-    "@pi-archimedes/core": "2.7.3"
+    "@pi-archimedes/core": "2.9.0"
   },
   "peerDependencies": {
     "@earendil-works/pi-coding-agent": ">=0.1.0",
     "@earendil-works/pi-tui": ">=0.1.0"
   },
   "devDependencies": {
-    "@earendil-works/pi-coding-agent": "^0.85.1",
-    "@earendil-works/pi-tui": "^0.85.1",
+    "@earendil-works/pi-coding-agent": "^0.87.0",
+    "@earendil-works/pi-tui": "^0.87.0",
     "typescript": "^6.0.0"
   },
   "pi": {

package/src/index.ts

@@ -1,11 +1,13 @@
 /**
- * dir | branch [+status] | worktree | model | ◕thinking | ↑↓R W $cost ━ context%
+ * dir | branch [+status] | worktree | model | ◕thinking · ⚠ext-status | ↑↓R W $cost ━ context%
  *
  * Adaptive layout — wraps to additional lines instead of truncating:
  * - fits width → single line (system info · stats · context bar)
  * - doesn't fit → two lines (system info above, stats + bar below)
  * - left sections alone overflow → three or more lines
  * - below splitThreshold (default 150) → at least two lines, as configured
+ * - oversized extension statuses are word-wrapped into multiple chunks so no
+ *   content is ever truncated by clampLine
  */
 
 import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
@@ -15,7 +17,7 @@ import { loadFooterConfig } from "./config.js";
 import { CostAccumulator } from "./cost-accumulator.js";
 import { getGitStatus, isInsideLinkedWorktree } from "./utils/git.js";
 import { getContextWindowInfo, getTokenUsageStats, type TokenUsageStats } from "./utils/stats.js";
-import { formatContextBar, formatGitStatusIndicators, formatThinkingIndicator, formatTokenCount } from "./utils/format.js";
+import { formatContextBar, formatGitStatusIndicators, formatThinkingIndicator, formatTokenCount, wrapStatusToChunks } from "./utils/format.js";
 import { footerIcons } from "./utils/icons.js";
 import { packFooterLines, SEP_W, SEPARATOR } from "./utils/layout.js";
 
@@ -76,12 +78,18 @@ export function registerFooter(pi: ExtensionAPI): void {
 
             const branchIcon = inWorktree ? footerIcons.worktree : footerIcons.branch;
 
-            // System info sections: dir | branch [+status] | model | thinking
+            // Extension status texts (pi setStatus) — may contain ANSI; oversized statuses wrap
+            const extensionStatusChunks = [...footerData.getExtensionStatuses().values()]
+              .filter(Boolean)
+              .flatMap((s) => wrapStatusToChunks(s, width));
+
+            // System info sections: dir | branch [+status] | model | thinking | ext-statuses
             const leftSections = [
               colorize("syntaxFunction", " " + footerIcons.directory + currentDirectory),
               currentBranch ? colorize("success", branchIcon + " " + currentBranch + (gitStatusStr ? " " + gitStatusStr : "")) : "",
               colorize("syntaxType", footerIcons.model + " " + activeModel),
               thinkingIndicatorStr,
+              ...extensionStatusChunks,
             ].filter(Boolean);
 
             // Usage stats: ↑in ↓out RcacheRead WcacheWrite $cost contextWindow

package/src/utils/format.test.ts

@@ -1,10 +1,12 @@
 import { describe, it, expect } from "vitest";
 import { visibleWidth } from "@earendil-works/pi-tui";
+import { stripAnsi } from "@pi-archimedes/core/text";
 import {
   formatTokenCount,
   formatContextBar,
   formatGitStatusIndicators,
   formatThinkingIndicator,
+  wrapStatusToChunks,
 } from "./format.js";
 import { thinkingLevelColors, thinkingLevelIcons, type ColorFn } from "./icons.js";
 
@@ -161,3 +163,55 @@ describe("formatThinkingIndicator", () => {
     expect(formatThinkingIndicator("unknown", mockColor)).toBe("◑ unknown");
   });
 });
+
+// ── wrapStatusToChunks ──────────────────────────────────────────────────────
+
+describe("wrapStatusToChunks", () => {
+  it("short status fits → returned unchanged (including ANSI)", () => {
+    const status = "\u001b[33m⚠ token invalid\u001b[0m";
+    const chunks = wrapStatusToChunks(status, 80);
+    expect(chunks).toEqual([status]);
+  });
+
+  it("empty status → [status]", () => {
+    expect(wrapStatusToChunks("", 80)).toEqual([""]);
+    expect(wrapStatusToChunks("   ", 80)).toEqual(["   "]);
+  });
+
+  it("oversized status wraps: no content loss, every chunk ≤ width", () => {
+    const colored = "\u001b[33m⚠\u001b[0m";
+    // Build a ~40-word status with an ANSI-colored prefix
+    const words = Array.from({ length: 40 }, (_, i) => `word${i}`);
+    const status = colored + " " + words.join(" ");
+    const width = 40;
+    const chunks = wrapStatusToChunks(status, width);
+    // Every chunk must fit in width
+    expect(chunks.every((c) => visibleWidth(c) <= width)).toBe(true);
+    // No content loss — normalise whitespace in stripped versions
+    const joined = chunks.map(stripAnsi).join(" ").replace(/\s+/g, " ").trim();
+    const original = stripAnsi(status).replace(/\s+/g, " ").trim();
+    expect(joined).toBe(original);
+  });
+
+  it("oversized single word (URL) splits without loss", () => {
+    const url = "https://" + "a".repeat(192);
+    const width = 30;
+    const chunks = wrapStatusToChunks(url, width);
+    // Concat of stripped chunks equals stripped original
+    expect(chunks.map(stripAnsi).join("")).toBe(stripAnsi(url));
+    // Every chunk fits
+    expect(chunks.every((c) => visibleWidth(c) <= width)).toBe(true);
+  });
+
+  it("word-wider-than-width preserves ANSI — concatenated raw string equals input", () => {
+    const longWord = "\u001b[31m" + "Z".repeat(60) + "\u001b[0m";
+    const width = 20;
+    const chunks = wrapStatusToChunks(longWord, width);
+    // Some chunk must contain the opening escape
+    expect(chunks.some((c) => c.includes("\u001b[31m"))).toBe(true);
+    // Concatenation of raw chunks (no spaces) equals the input
+    expect(chunks.join("")).toBe(longWord);
+    // No chunk exceeds width
+    expect(chunks.every((c) => visibleWidth(c) <= width)).toBe(true);
+  });
+});

package/src/utils/format.ts

@@ -1,3 +1,4 @@
+import { visibleWidth } from "@earendil-works/pi-tui";
 import { footerIcons, gitDisplayIcons, gitStatusColors, thinkingLevelColors, thinkingLevelIcons, type ColorFn } from "./icons.js";
 
 // Token counts use binary units (1K = 1024)
@@ -52,3 +53,126 @@ export function formatGitStatusIndicators(
 export function formatThinkingIndicator(thinkingLevel: string, colorize: ColorFn): string {
   return colorize(thinkingLevelColors[thinkingLevel] || "dim", `${thinkingLevelIcons[thinkingLevel] || "◑"} ${thinkingLevel}`);
 }
+
+/**
+ * Wrap an extension status into chunks that fit the terminal.
+ * - fits (visibleWidth ≤ width) → [status] (unchanged, including ANSI)
+ * - oversized → greedy word-wrap: whole words packed per chunk; a single word
+ *   wider than `width` is split at character level on the raw string (ANSI
+ *   sequences stay attached to the chars they prefix because we split on plain
+ *   whitespace only — never mid-escape).
+ * No content is ever dropped; every returned chunk has visibleWidth ≤ width.
+ */
+export function wrapStatusToChunks(status: string, width: number): string[] {
+  // Edge cases: empty/whitespace or non-positive width → return as-is
+  if (width <= 0 || !status.trim()) return [status];
+  // Already fits → return unchanged (preserves ANSI)
+  if (visibleWidth(status) <= width) return [status];
+
+  // Split on whitespace boundaries only — never inside ANSI escapes
+  // (ANSI escapes contain no plain whitespace, so splitting on /\s+/ is safe)
+  const words = status.split(/(\s+)/);
+  // words alternates: [word, sep, word, sep, ...]
+  // We reassemble non-sep tokens and carry sep tokens to preserve spacing context
+  // Simpler: collect [token, isSep] pairs
+  const tokens: Array<{ text: string; isSep: boolean }> = [];
+  for (const w of words) {
+    if (w === "") continue;
+    tokens.push({ text: w, isSep: /^\s+$/.test(w) });
+  }
+
+  const chunks: string[] = [];
+  let current = "";
+  let currentW = 0;
+
+  /**
+   * Split a single word (no whitespace) that is wider than `width` into
+   * sub-chunks of at most `width` visible columns, character by character.
+   * ANSI escapes are kept intact (we only split printable characters;
+   * the regex /\x1b\[[^m]*m/ never contains whitespace, so they were
+   * never split by the outer word-split, and here we advance char by char
+   * through the raw string, which keeps escape bytes together with their
+   * surrounding characters as long as we don't split inside an escape).
+   *
+   * Implementation: iterate character-by-character through the raw string,
+   * accumulating into a sub-chunk until adding the next char would exceed width.
+   */
+  function splitWideWord(word: string): string[] {
+    const result: string[] = [];
+    let sub = "";
+    let subW = 0;
+    let i = 0;
+    while (i < word.length) {
+      // Detect ANSI escape sequence starting at i (\x1b[...m)
+      // An escape contributes 0 visible width, so we append it whole.
+      if (word[i] === "\x1b") {
+        const end = word.indexOf("m", i + 1);
+        if (end !== -1) {
+          sub += word.slice(i, end + 1);
+          i = end + 1;
+          continue;
+        }
+      }
+      // Regular character (may be multi-byte Unicode)
+      const char = word[i]!;
+      const charW = visibleWidth(char);
+      if (subW + charW > width && sub !== "") {
+        result.push(sub);
+        sub = char;
+        subW = charW;
+      } else {
+        sub += char;
+        subW += charW;
+      }
+      i++;
+    }
+    if (sub !== "") result.push(sub);
+    return result.length > 0 ? result : [word];
+  }
+
+  function flushCurrent() {
+    if (current !== "") {
+      chunks.push(current);
+      current = "";
+      currentW = 0;
+    }
+  }
+
+  for (const token of tokens) {
+    if (token.isSep) {
+      // Separator — skip (we rebuild spacing implicitly via chunks)
+      continue;
+    }
+    const word = token.text;
+    const wordW = visibleWidth(word);
+
+    if (wordW > width) {
+      // Word is wider than the entire terminal — flush then char-split it
+      flushCurrent();
+      const subChunks = splitWideWord(word);
+      for (let j = 0; j < subChunks.length - 1; j++) {
+        chunks.push(subChunks[j]!);
+      }
+      // Last sub-chunk becomes the start of the next current
+      const last = subChunks[subChunks.length - 1]!;
+      current = last;
+      currentW = visibleWidth(last);
+    } else if (currentW === 0) {
+      // Starting a new chunk
+      current = word;
+      currentW = wordW;
+    } else if (currentW + 1 + wordW <= width) {
+      // Append with a single space
+      current += " " + word;
+      currentW += 1 + wordW;
+    } else {
+      // Doesn't fit — flush and start new chunk
+      flushCurrent();
+      current = word;
+      currentW = wordW;
+    }
+  }
+  flushCurrent();
+
+  return chunks.length > 0 ? chunks : [status];
+}

package/src/utils/render-smoke.test.ts

@@ -1,7 +1,7 @@
 import { describe, it, expect } from "vitest";
 import { stripAnsi } from "@pi-archimedes/core/text";
 import { packFooterLines, SEP_W, SEPARATOR } from "./layout.js";
-import { formatContextBar } from "./format.js";
+import { formatContextBar, wrapStatusToChunks } from "./format.js";
 import { visibleWidth } from "@earendil-works/pi-tui";
 
 // Minimal stand-ins: same visible chars as the real footer, no ANSI
@@ -14,13 +14,14 @@ const MODEL = "🧠 claude-sonnet-4-6";
 const THINKING = "◕ high"; // current format for the 'high' level
 const STATS = "↑141 ↓42k R8.7M W220k $4.25 116k/977k";
 
-function buildLines(width: number, pct: number, splitThreshold = 150) {
+function buildLines(width: number, pct: number, splitThreshold = 150, extraSections: string[] = []) {
   const leftSections = [
     DIR,
     BRANCH,
     WORKTREE,
     MODEL,
     THINKING,
+    ...extraSections,
   ].filter(Boolean);
   let groups = packFooterLines([...leftSections, STATS], width, SEP_W);
   if (groups.length < 2 && width < splitThreshold && leftSections.length > 0) {
@@ -84,4 +85,35 @@ describe("render glue simulation", () => {
     // No line exceeds 40 visible columns
     expect(lines.every((l) => visibleWidth(stripAnsi(l)) <= 40)).toBe(true);
   });
+
+  it("extension status chunk wraps without truncation on the user's terminal", () => {
+    const STATUS = "⚠ litellm token invalid — re-auth required";
+    const lines = buildLines(158, 12, 150, [STATUS]);
+    const stripped = lines.map((l) => stripAnsi(l));
+    const all = stripped.join(" ");
+    // (a) status text is fully present across joined lines (no clipping)
+    expect(all).toContain(STATUS);
+    // (b) no line exceeds 158 visible columns
+    expect(lines.every((l) => visibleWidth(stripAnsi(l)) <= 158)).toBe(true);
+  });
+
+  it("oversized extension status wraps through the production path without truncation", () => {
+    const LONG =
+      "⚠ litellm token invalid — re-auth required: " +
+      "gcloud auth application-default login (see https://support.google.com/a/answer/9368756 for details) — check Settings → Auth for more context " +
+      "and verify your credentials are up-to-date before retrying the long operation that failed";
+    // Verify it is >158 visible chars
+    expect(visibleWidth(LONG)).toBeGreaterThan(158);
+    const chunks = wrapStatusToChunks(LONG, 158);
+    const lines = buildLines(158, 12, 150, chunks);
+    const stripped = lines.map((l) => stripAnsi(l));
+    // (a) every word of the original status appears in the joined lines
+    const allText = stripped.join(" ");
+    const words = stripAnsi(LONG).split(/\s+/).filter(Boolean);
+    for (const word of words) {
+      expect(allText).toContain(word);
+    }
+    // (b) every line ≤ 158 visible cols
+    expect(lines.every((l) => visibleWidth(stripAnsi(l)) <= 158)).toBe(true);
+  });
 });

package/src/utils/stats.test.ts

@@ -21,9 +21,16 @@ interface UserMessage {
   content: string;
 }
 
+interface ToolResultMessage {
+  role: "toolResult";
+  usage?: MessageUsage | undefined;
+}
+
 interface SessionEntry {
-  type: "message";
-  message: AssistantMessage | UserMessage;
+  id: string;
+  type: "message" | "usage" | "compaction" | "branch_summary";
+  message?: AssistantMessage | UserMessage | ToolResultMessage | undefined;
+  usage?: MessageUsage | undefined;
 }
 
 interface MockContext {
@@ -34,20 +41,48 @@ interface MockContext {
   model?: { contextWindow?: number };
 }
 
-function makeAssistantEntry(usage: MessageUsage): SessionEntry {
+let idCounter = 1;
+
+function makeAssistantEntry(idOrUsage: string | MessageUsage, maybeUsage?: MessageUsage): SessionEntry {
+  const id = typeof idOrUsage === "string" ? idOrUsage : `asst-${idCounter++}`;
+  const usage = typeof idOrUsage === "string" ? maybeUsage! : idOrUsage;
   return {
+    id,
     type: "message",
     message: { role: "assistant", usage },
   };
 }
 
-function makeUserEntry(): SessionEntry {
+function makeUserEntry(id = `user-${idCounter++}`): SessionEntry {
   return {
+    id,
     type: "message",
     message: { role: "user", content: "hello" },
   };
 }
 
+function makeUsageEntry(id: string, usage: MessageUsage): SessionEntry {
+  return {
+    id,
+    type: "usage",
+    usage,
+  };
+}
+
+function makeToolResultEntry(id: string, usage?: MessageUsage): SessionEntry {
+  return {
+    id,
+    type: "message",
+    message: usage !== undefined ? { role: "toolResult", usage } : { role: "toolResult" },
+  };
+}
+
+function makeCompactionEntry(id: string, usage?: MessageUsage): SessionEntry {
+  return usage !== undefined
+    ? { id, type: "compaction", usage }
+    : { id, type: "compaction" };
+}
+
 function makeCtx(entries: SessionEntry[], contextUsage?: { contextWindow?: number; percent?: number }, modelContextWindow?: number): any {
   return {
     sessionManager: { getEntries: () => entries },
@@ -61,12 +96,14 @@ function makeCtx(entries: SessionEntry[], contextUsage?: { contextWindow?: numbe
 describe("getTokenUsageStats", () => {
   let getTokenUsageStats: typeof import("./stats.js").getTokenUsageStats;
   let invalidateStatsCache: typeof import("./stats.js").invalidateStatsCache;
+  let resetStatsState: typeof import("./stats.js").resetStatsState;
 
   async function loadModule() {
     vi.resetModules();
     const mod = await import("./stats.js");
     getTokenUsageStats = mod.getTokenUsageStats;
     invalidateStatsCache = mod.invalidateStatsCache;
+    resetStatsState = mod.resetStatsState;
   }
 
   it("empty entries returns zero stats", async () => {
@@ -124,6 +161,135 @@ describe("getTokenUsageStats", () => {
     expect(result.totalOutput).toBe(50);
   });
 
+  it("accumulates tokens and cost from UsageEntry (cache warming)", async () => {
+    await loadModule();
+    const ctx = makeCtx([
+      makeAssistantEntry("a1", { input: 100, output: 50, cacheRead: 10, cacheWrite: 5, cost: { total: 0.01 } }),
+      makeUsageEntry("u1", { input: 0, output: 0, cacheRead: 500, cacheWrite: 200, cost: { total: 0.005 } }),
+    ]);
+    const result = getTokenUsageStats(ctx);
+    expect(result).toEqual({
+      totalInput: 100,
+      totalOutput: 50,
+      totalCacheRead: 510,
+      totalCacheWrite: 205,
+      totalCost: 0.015,
+    });
+  });
+
+  it("accumulates usage from ToolResultMessage", async () => {
+    await loadModule();
+    const ctx = makeCtx([
+      makeAssistantEntry("a1", { input: 100, output: 50, cacheRead: 0, cacheWrite: 0, cost: { total: 0.01 } }),
+      makeToolResultEntry("t1", { input: 200, output: 80, cacheRead: 50, cacheWrite: 20, cost: { total: 0.02 } }),
+      makeToolResultEntry("t2"), // toolResult without usage ignored
+    ]);
+    const result = getTokenUsageStats(ctx);
+    expect(result).toEqual({
+      totalInput: 300,
+      totalOutput: 130,
+      totalCacheRead: 50,
+      totalCacheWrite: 20,
+      totalCost: 0.03,
+    });
+  });
+
+  it("accumulates usage from CompactionEntry and branch_summary", async () => {
+    await loadModule();
+    const ctx = makeCtx([
+      makeCompactionEntry("c1", { input: 1000, output: 200, cacheRead: 100, cacheWrite: 50, cost: { total: 0.04 } }),
+      {
+        id: "bs1",
+        type: "branch_summary",
+        usage: { input: 500, output: 100, cacheRead: 50, cacheWrite: 25, cost: { total: 0.02 } },
+      } as SessionEntry,
+    ]);
+    const result = getTokenUsageStats(ctx);
+    expect(result).toEqual({
+      totalInput: 1500,
+      totalOutput: 300,
+      totalCacheRead: 150,
+      totalCacheWrite: 75,
+      totalCost: 0.06,
+    });
+  });
+
+  it("detects branch switch: changing entry IDs at anchor position triggers clean re-scan", async () => {
+    await loadModule();
+    // Branch 1: entries [a1, a2]
+    const branch1 = [
+      makeAssistantEntry("a1", { input: 100, output: 50, cacheRead: 0, cacheWrite: 0, cost: { total: 0.01 } }),
+      makeAssistantEntry("a2", { input: 200, output: 100, cacheRead: 0, cacheWrite: 0, cost: { total: 0.02 } }),
+    ];
+    let ctx = makeCtx(branch1);
+    let result = getTokenUsageStats(ctx);
+    expect(result.totalInput).toBe(300);
+    expect(result.totalOutput).toBe(150);
+
+    // Branch switch: second entry is different id/usage (a2_alt instead of a2)
+    const branch2: SessionEntry[] = [
+      branch1[0]!,
+      makeAssistantEntry("a2_alt", { input: 50, output: 25, cacheRead: 0, cacheWrite: 0, cost: { total: 0.005 } }),
+    ];
+    ctx = makeCtx(branch2);
+    result = getTokenUsageStats(ctx);
+    // Should NOT be 350 (accumulated on old running total); must re-scan branch2 cleanly
+    expect(result).toEqual({
+      totalInput: 150,
+      totalOutput: 75,
+      totalCacheRead: 0,
+      totalCacheWrite: 0,
+      totalCost: 0.015,
+    });
+  });
+
+  it("detects branch switch: changing first entry ID triggers clean re-scan even if tail anchor matches", async () => {
+    await loadModule();
+    // Branch 1: entries [a1, a2]
+    const branch1 = [
+      makeAssistantEntry("a1", { input: 100, output: 50, cacheRead: 0, cacheWrite: 0, cost: { total: 0.01 } }),
+      makeAssistantEntry("a2", { input: 200, output: 100, cacheRead: 0, cacheWrite: 0, cost: { total: 0.02 } }),
+    ];
+    let ctx = makeCtx(branch1);
+    let result = getTokenUsageStats(ctx);
+    expect(result.totalInput).toBe(300);
+
+    // Branch switch where tail ID 'a2' happens to match, but head entry is 'a1_replaced'
+    const branch2: SessionEntry[] = [
+      makeAssistantEntry("a1_replaced", { input: 10, output: 5, cacheRead: 0, cacheWrite: 0, cost: { total: 0.001 } }),
+      branch1[1]!,
+    ];
+    ctx = makeCtx(branch2);
+    result = getTokenUsageStats(ctx);
+    expect(result).toEqual({
+      totalInput: 210,
+      totalOutput: 105,
+      totalCacheRead: 0,
+      totalCacheWrite: 0,
+      totalCost: 0.021,
+    });
+  });
+
+  it("updates total when tail entry usage mutates in-place with unchanged entry count", async () => {
+    await loadModule();
+    const tailEntry = makeAssistantEntry("a1", { input: 100, output: 50, cacheRead: 0, cacheWrite: 0, cost: { total: 0.01 } });
+    const entries = [tailEntry];
+    const ctx = makeCtx(entries);
+
+    const initial = getTokenUsageStats(ctx);
+    expect(initial.totalOutput).toBe(50);
+    expect(initial.totalCost).toBe(0.01);
+
+    // Direct in-place property mutation of existing usage object (aliasing test)
+    const usageObj = (tailEntry.message as AssistantMessage).usage;
+    usageObj.output = 120;
+    usageObj.cost = { total: 0.025 };
+
+    const updated = getTokenUsageStats(ctx);
+    expect(updated.totalOutput).toBe(120);
+    expect(updated.totalCost).toBe(0.025);
+  });
+
   it("cache returns same result within TTL", async () => {
     await loadModule();
     const ctx = makeCtx([
@@ -146,6 +312,19 @@ describe("getTokenUsageStats", () => {
     expect(first).not.toBe(second); // different reference (cache cleared)
   });
 
+  it("resetStatsState clears module state completely", async () => {
+    await loadModule();
+    const ctx = makeCtx([
+      makeAssistantEntry("a1", { input: 100, output: 50, cacheRead: 0, cacheWrite: 0, cost: { total: 0.01 } }),
+    ]);
+    const first = getTokenUsageStats(ctx);
+    expect(first.totalInput).toBe(100);
+    resetStatsState();
+    const second = getTokenUsageStats(ctx);
+    expect(second).toEqual(first);
+    expect(second).not.toBe(first);
+  });
+
   it("property: stats are monotonic (adding entries never decreases totals)", async () => {
     await loadModule();
     fc.assert(

package/src/utils/stats.ts

@@ -1,15 +1,11 @@
 import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
 
-interface MessageUsage {
-  input: number;
-  output: number;
-  cacheRead: number;
-  cacheWrite: number;
-  cost: { total: number };
-}
-
-interface AssistantMessage {
-  usage: MessageUsage;
+export interface MessageUsage {
+  input?: number;
+  output?: number;
+  cacheRead?: number;
+  cacheWrite?: number;
+  cost?: { total?: number };
 }
 
 export interface TokenUsageStats {
@@ -20,84 +16,183 @@ export interface TokenUsageStats {
   totalCost: number;
 }
 
-// Cache TTL: 500ms — stats don't change more frequently than message_end events
-const STATS_CACHE_TTL_MS = 500;
+export function extractEntryUsage(sessionEntry: any): MessageUsage | null {
+  if (sessionEntry?.type === "usage") {
+    return sessionEntry.usage ?? null;
+  }
+  if (sessionEntry?.type === "message" && sessionEntry.message?.role === "assistant") {
+    return sessionEntry.message.usage ?? null;
+  }
+  if (sessionEntry?.type === "message" && sessionEntry.message?.role === "toolResult" && sessionEntry.message.usage) {
+    return sessionEntry.message.usage;
+  }
+  if ((sessionEntry?.type === "compaction" || sessionEntry?.type === "branch_summary") && sessionEntry.usage) {
+    return sessionEntry.usage;
+  }
+  return null;
+}
 
-interface StatsCacheEntry {
-  value: TokenUsageStats;
-  timestamp: number;
-  entryCount: number;
+function snapshotUsage(usage: MessageUsage | null | undefined): MessageUsage | undefined {
+  if (!usage) return undefined;
+  return {
+    input: usage.input ?? 0,
+    output: usage.output ?? 0,
+    cacheRead: usage.cacheRead ?? 0,
+    cacheWrite: usage.cacheWrite ?? 0,
+    cost: { total: usage.cost?.total ?? 0 },
+  };
 }
 
-let statsCache: StatsCacheEntry | undefined;
+function usageEquals(a: MessageUsage | undefined, b: MessageUsage | undefined): boolean {
+  if (!a && !b) return true;
+  if (!a || !b) return false;
+  return (
+    (a.input ?? 0) === (b.input ?? 0) &&
+    (a.output ?? 0) === (b.output ?? 0) &&
+    (a.cacheRead ?? 0) === (b.cacheRead ?? 0) &&
+    (a.cacheWrite ?? 0) === (b.cacheWrite ?? 0) &&
+    (a.cost?.total ?? 0) === (b.cost?.total ?? 0)
+  );
+}
 
 // Running total from initial scan — avoids re-scanning old entries
 let runningTotal: TokenUsageStats | undefined;
 let runningTotalEntryCount = 0;
+let lastFirstEntryId: string | undefined;
+let lastAnchorEntryId: string | undefined;
+let lastTailUsage: MessageUsage | undefined;
+let statsCache: TokenUsageStats | undefined;
 
 export function getTokenUsageStats(ctx: ExtensionContext): TokenUsageStats {
   const entries = ctx.sessionManager.getEntries();
 
-  // Return cached result if entry count hasn't changed and cache is fresh
-  if (
-    statsCache &&
-    statsCache.entryCount === entries.length &&
-    Date.now() - statsCache.timestamp < STATS_CACHE_TTL_MS
-  ) {
-    return statsCache.value;
+  if (entries.length === 0) {
+    resetStatsState();
+    return {
+      totalInput: 0,
+      totalOutput: 0,
+      totalCacheRead: 0,
+      totalCacheWrite: 0,
+      totalCost: 0,
+    };
+  }
+
+  const anchorMatch =
+    runningTotal !== undefined &&
+    runningTotalEntryCount > 0 &&
+    entries[0]?.id === lastFirstEntryId &&
+    entries[runningTotalEntryCount - 1]?.id === lastAnchorEntryId;
+
+  if (anchorMatch && runningTotal) {
+    if (entries.length === runningTotalEntryCount) {
+      // Invariant: for unchanged entry count, only the tail entry's usage mutates during streaming.
+      const currentTailUsage = extractEntryUsage(entries[entries.length - 1]);
+      if (!usageEquals(currentTailUsage ?? undefined, lastTailUsage)) {
+        const pIn = lastTailUsage?.input ?? 0;
+        const pOut = lastTailUsage?.output ?? 0;
+        const pCr = lastTailUsage?.cacheRead ?? 0;
+        const pCw = lastTailUsage?.cacheWrite ?? 0;
+        const pCost = lastTailUsage?.cost?.total ?? 0;
+
+        const cIn = currentTailUsage?.input ?? 0;
+        const cOut = currentTailUsage?.output ?? 0;
+        const cCr = currentTailUsage?.cacheRead ?? 0;
+        const cCw = currentTailUsage?.cacheWrite ?? 0;
+        const cCost = currentTailUsage?.cost?.total ?? 0;
+
+        runningTotal = {
+          totalInput: runningTotal.totalInput + (cIn - pIn),
+          totalOutput: runningTotal.totalOutput + (cOut - pOut),
+          totalCacheRead: runningTotal.totalCacheRead + (cCr - pCr),
+          totalCacheWrite: runningTotal.totalCacheWrite + (cCw - pCw),
+          totalCost: runningTotal.totalCost + (cCost - pCost),
+        };
+        lastTailUsage = snapshotUsage(currentTailUsage);
+        statsCache = runningTotal;
+        return runningTotal;
+      }
+
+      if (statsCache !== undefined) {
+        return statsCache;
+      }
+      statsCache = { ...runningTotal };
+      return statsCache;
+    }
+
+    if (entries.length > runningTotalEntryCount) {
+      let totalInput = runningTotal.totalInput;
+      let totalOutput = runningTotal.totalOutput;
+      let totalCacheRead = runningTotal.totalCacheRead;
+      let totalCacheWrite = runningTotal.totalCacheWrite;
+      let totalCost = runningTotal.totalCost;
+
+      const prevTailUsage = extractEntryUsage(entries[runningTotalEntryCount - 1]);
+      if (!usageEquals(prevTailUsage ?? undefined, lastTailUsage)) {
+        const pIn = lastTailUsage?.input ?? 0;
+        const pOut = lastTailUsage?.output ?? 0;
+        const pCr = lastTailUsage?.cacheRead ?? 0;
+        const pCw = lastTailUsage?.cacheWrite ?? 0;
+        const pCost = lastTailUsage?.cost?.total ?? 0;
+
+        const cIn = prevTailUsage?.input ?? 0;
+        const cOut = prevTailUsage?.output ?? 0;
+        const cCr = prevTailUsage?.cacheRead ?? 0;
+        const cCw = prevTailUsage?.cacheWrite ?? 0;
+        const cCost = prevTailUsage?.cost?.total ?? 0;
+
+        totalInput += cIn - pIn;
+        totalOutput += cOut - pOut;
+        totalCacheRead += cCr - pCr;
+        totalCacheWrite += cCw - pCw;
+        totalCost += cCost - pCost;
+      }
+
+      for (let i = runningTotalEntryCount; i < entries.length; i++) {
+        const u = extractEntryUsage(entries[i]);
+        if (u) {
+          totalInput += u.input ?? 0;
+          totalOutput += u.output ?? 0;
+          totalCacheRead += u.cacheRead ?? 0;
+          totalCacheWrite += u.cacheWrite ?? 0;
+          totalCost += u.cost?.total ?? 0;
+        }
+      }
+
+      runningTotal = { totalInput, totalOutput, totalCacheRead, totalCacheWrite, totalCost };
+      runningTotalEntryCount = entries.length;
+      lastFirstEntryId = entries[0]?.id;
+      lastAnchorEntryId = entries[entries.length - 1]?.id;
+      lastTailUsage = snapshotUsage(extractEntryUsage(entries[entries.length - 1]));
+      statsCache = runningTotal;
+      return runningTotal;
+    }
   }
 
+  // Else (first run, entries truncated, or branch switch where anchor id mismatched):
   let totalInput = 0,
     totalOutput = 0,
     totalCacheRead = 0,
     totalCacheWrite = 0,
     totalCost = 0;
 
-  // If we have a running total, only scan new entries
-  const hasRunningTotal = runningTotal !== undefined && runningTotalEntryCount < entries.length;
-  const startIdx = hasRunningTotal ? runningTotalEntryCount : 0;
-
-  if (startIdx === 0) {
-    // Full scan — no running total yet
-    for (const sessionEntry of entries) {
-      if (sessionEntry?.type === "message" && sessionEntry.message?.role === "assistant") {
-        const assistantMessage = sessionEntry.message as AssistantMessage;
-        totalInput += assistantMessage.usage.input;
-        totalOutput += assistantMessage.usage.output;
-        totalCacheRead += assistantMessage.usage.cacheRead;
-        totalCacheWrite += assistantMessage.usage.cacheWrite;
-        totalCost += assistantMessage.usage.cost.total;
-      }
-    }
-    runningTotal = { totalInput, totalOutput, totalCacheRead, totalCacheWrite, totalCost };
-    runningTotalEntryCount = entries.length;
-  } else {
-    // Incremental: start from running total + scan new entries
-    const rt = runningTotal!; // Safe: hasRunningTotal guard above
-    totalInput = rt.totalInput;
-    totalOutput = rt.totalOutput;
-    totalCacheRead = rt.totalCacheRead;
-    totalCacheWrite = rt.totalCacheWrite;
-    totalCost = rt.totalCost;
-
-    for (let i = startIdx; i < entries.length; i++) {
-      const sessionEntry = entries[i];
-      if (sessionEntry?.type === "message" && sessionEntry.message?.role === "assistant") {
-        const assistantMessage = sessionEntry.message as AssistantMessage;
-        totalInput += assistantMessage.usage.input;
-        totalOutput += assistantMessage.usage.output;
-        totalCacheRead += assistantMessage.usage.cacheRead;
-        totalCacheWrite += assistantMessage.usage.cacheWrite;
-        totalCost += assistantMessage.usage.cost.total;
-      }
+  for (const sessionEntry of entries) {
+    const u = extractEntryUsage(sessionEntry);
+    if (u) {
+      totalInput += u.input ?? 0;
+      totalOutput += u.output ?? 0;
+      totalCacheRead += u.cacheRead ?? 0;
+      totalCacheWrite += u.cacheWrite ?? 0;
+      totalCost += u.cost?.total ?? 0;
     }
-    runningTotalEntryCount = entries.length;
   }
 
-  const result: TokenUsageStats = { totalInput, totalOutput, totalCacheRead, totalCacheWrite, totalCost };
-  runningTotal = result;
-  statsCache = { value: result, timestamp: Date.now(), entryCount: entries.length };
-  return result;
+  runningTotal = { totalInput, totalOutput, totalCacheRead, totalCacheWrite, totalCost };
+  runningTotalEntryCount = entries.length;
+  lastFirstEntryId = entries[0]?.id;
+  lastAnchorEntryId = entries[entries.length - 1]?.id;
+  lastTailUsage = snapshotUsage(extractEntryUsage(entries[entries.length - 1]));
+  statsCache = runningTotal;
+  return runningTotal;
 }
 
 /** Clear the stats cache — call when a new message arrives. */
@@ -105,6 +200,16 @@ export function invalidateStatsCache(): void {
   statsCache = undefined;
 }
 
+/** Reset all module-level stats state for clean isolation. */
+export function resetStatsState(): void {
+  runningTotal = undefined;
+  runningTotalEntryCount = 0;
+  lastFirstEntryId = undefined;
+  lastAnchorEntryId = undefined;
+  lastTailUsage = undefined;
+  statsCache = undefined;
+}
+
 export interface ContextWindowInfo {
   percent: string;
   percentValue: number;