@pi-archimedes/footer 2.8.0 → 2.9.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@pi-archimedes/footer",
-  "version": "2.8.0",
+  "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.8.0"
+    "@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/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;