feat: live task token usage and stats-tab fixes

Capture per-session token usage from pi-coding-agent's getSessionStats()
after each promptWithFallback in the executor and merger paths, so
task.tokenUsage populates live during runs and reflects final totals on
done tasks. Previously the executor never read session usage and only
the heartbeat path bumped agent token totals, leaving task.tokenUsage
undefined even after completion.

Stats panel and done-card timing also now reflect live state: the modal
overlays the SSE-updated task prop on top of the one-shot fullDetail
snapshot, in-progress workflow steps contribute live elapsed to the
Workflow runtime metric, and the done card uses Timed duration (matching
the stats tab) with workflow runtime as fallback. Time indicator labels
coarsened to <1m / Nm / Nh / Nd.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

packages/engine/src/__tests__/session-token-usage.test.ts

@@ -0,0 +1,125 @@
+import { describe, it, expect, vi, beforeEach } from "vitest";
+import type { Task, TaskStore } from "@fusion/core";
+import { accumulateSessionTokenUsage } from "../session-token-usage.js";
+
+interface MockSessionStats {
+  tokens?: { input?: number; output?: number; cacheRead?: number; cacheWrite?: number };
+}
+
+function createSession(stats: MockSessionStats | undefined) {
+  return { getSessionStats: vi.fn(() => stats) } as unknown as Parameters<typeof accumulateSessionTokenUsage>[2];
+}
+
+function createStore(initial: Task["tokenUsage"]): TaskStore & { _task: Task; updateTask: ReturnType<typeof vi.fn> } {
+  const task = { id: "FN-1", tokenUsage: initial } as Task;
+  const updateTask = vi.fn(async (_id: string, updates: Partial<Task>) => {
+    if (updates.tokenUsage !== undefined) task.tokenUsage = updates.tokenUsage as Task["tokenUsage"];
+    return task;
+  });
+  const store = {
+    _task: task,
+    getTask: vi.fn(async () => task),
+    updateTask,
+  } as unknown as TaskStore & { _task: Task; updateTask: ReturnType<typeof vi.fn> };
+  return store;
+}
+
+describe("accumulateSessionTokenUsage", () => {
+  beforeEach(() => {
+    vi.useRealTimers();
+  });
+
+  it("writes initial token usage when task has none", async () => {
+    const store = createStore(undefined);
+    const session = createSession({ tokens: { input: 100, output: 30, cacheRead: 5, cacheWrite: 2 } });
+
+    await accumulateSessionTokenUsage(store, "FN-1", session);
+
+    expect(store.updateTask).toHaveBeenCalledTimes(1);
+    const call = store.updateTask.mock.calls[0]![1] as { tokenUsage: Task["tokenUsage"] };
+    expect(call.tokenUsage).toMatchObject({
+      inputTokens: 102, // input + cacheWrite
+      outputTokens: 30,
+      cachedTokens: 5,
+      totalTokens: 137,
+    });
+    expect(typeof call.tokenUsage!.firstUsedAt).toBe("string");
+    expect(typeof call.tokenUsage!.lastUsedAt).toBe("string");
+  });
+
+  it("accumulates only the delta on subsequent calls for the same session", async () => {
+    const store = createStore(undefined);
+    const session = createSession({ tokens: { input: 100, output: 30, cacheRead: 0, cacheWrite: 0 } });
+
+    await accumulateSessionTokenUsage(store, "FN-1", session);
+
+    // Second call: session has progressed.
+    (session as unknown as { getSessionStats: () => MockSessionStats }).getSessionStats = () => ({
+      tokens: { input: 250, output: 80, cacheRead: 0, cacheWrite: 0 },
+    });
+    await accumulateSessionTokenUsage(store, "FN-1", session);
+
+    expect(store.updateTask).toHaveBeenCalledTimes(2);
+    const second = store.updateTask.mock.calls[1]![1] as { tokenUsage: Task["tokenUsage"] };
+    expect(second.tokenUsage).toMatchObject({
+      inputTokens: 250,
+      outputTokens: 80,
+      cachedTokens: 0,
+      totalTokens: 330,
+    });
+  });
+
+  it("preserves firstUsedAt across updates", async () => {
+    const store = createStore(undefined);
+    const session = createSession({ tokens: { input: 10, output: 5, cacheRead: 0, cacheWrite: 0 } });
+
+    await accumulateSessionTokenUsage(store, "FN-1", session);
+    const first = store.updateTask.mock.calls[0]![1] as { tokenUsage: Task["tokenUsage"] };
+    const initialFirstUsed = first.tokenUsage!.firstUsedAt;
+
+    await new Promise((resolve) => setTimeout(resolve, 5));
+    (session as unknown as { getSessionStats: () => MockSessionStats }).getSessionStats = () => ({
+      tokens: { input: 20, output: 5, cacheRead: 0, cacheWrite: 0 },
+    });
+    await accumulateSessionTokenUsage(store, "FN-1", session);
+
+    const second = store.updateTask.mock.calls[1]![1] as { tokenUsage: Task["tokenUsage"] };
+    expect(second.tokenUsage!.firstUsedAt).toBe(initialFirstUsed);
+    expect(second.tokenUsage!.lastUsedAt >= initialFirstUsed).toBe(true);
+  });
+
+  it("does nothing when session has no getSessionStats", async () => {
+    const store = createStore(undefined);
+    const session = {} as Parameters<typeof accumulateSessionTokenUsage>[2];
+
+    await accumulateSessionTokenUsage(store, "FN-1", session);
+
+    expect(store.updateTask).not.toHaveBeenCalled();
+  });
+
+  it("does nothing when delta is zero", async () => {
+    const store = createStore({
+      inputTokens: 50,
+      outputTokens: 20,
+      cachedTokens: 0,
+      totalTokens: 70,
+      firstUsedAt: "2024-01-01T00:00:00.000Z",
+      lastUsedAt: "2024-01-01T00:00:00.000Z",
+    });
+    const session = createSession({ tokens: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 } });
+
+    await accumulateSessionTokenUsage(store, "FN-1", session);
+    // Calling again with the same zero stats should not produce another update.
+    await accumulateSessionTokenUsage(store, "FN-1", session);
+
+    expect(store.updateTask).not.toHaveBeenCalled();
+  });
+
+  it("swallows store errors instead of throwing", async () => {
+    const store = createStore(undefined);
+    (store.updateTask as ReturnType<typeof vi.fn>).mockRejectedValueOnce(new Error("db down"));
+    const session = createSession({ tokens: { input: 1, output: 1, cacheRead: 0, cacheWrite: 0 } });
+
+    await expect(accumulateSessionTokenUsage(store, "FN-1", session)).resolves.toBeUndefined();
+  });
+});

packages/engine/src/executor.ts

@@ -11,6 +11,7 @@ import { findWorktreeUser } from "./merger.js";
 import { generateWorktreeName, slugify } from "./worktree-names.js";
 import { Type, type Static } from "@mariozechner/pi-ai";
 import { describeModel, promptWithFallback, compactSessionContext } from "./pi.js";
+import { accumulateSessionTokenUsage } from "./session-token-usage.js";
 import { createResolvedAgentSession, extractRuntimeHint } from "./agent-session-helpers.js";
 import { buildSessionSkillContext } from "./session-skill-context.js";
 import { reviewStep, type ReviewVerdict } from "./reviewer.js";
@@ -2015,6 +2016,7 @@ export class TaskExecutor {
           // session.prompt() resolves normally even when retries are exhausted —
           // the error is stored on session.state.error instead of being thrown.
           checkSessionError(session);
+          await accumulateSessionTokenUsage(this.store, task.id, session);
 
           // Check if proactive context compaction is needed based on token cap setting.
           // This runs after the main prompt completes to avoid interrupting active work.
@@ -2072,6 +2074,7 @@ export class TaskExecutor {
 
             await promptWithFallback(session, resumePrompt);
             checkSessionError(session);
+            await accumulateSessionTokenUsage(this.store, task.id, session);
           }
 
           // If dependency was added during execution, discard worktree and move to triage
@@ -2233,6 +2236,7 @@ export class TaskExecutor {
               stuckDetector?.recordActivity(task.id);
               await promptWithFallback(retrySession, retryPrompt);
               checkSessionError(retrySession);
+              await accumulateSessionTokenUsage(this.store, task.id, retrySession);
 
               if (!taskDone) {
                 const implicitCheck = await this.store.getTask(task.id);
@@ -2439,6 +2443,7 @@ export class TaskExecutor {
 
               await promptWithFallback(activeEntry.session, reducedPrompt);
               checkSessionError(activeEntry.session);
+              await accumulateSessionTokenUsage(this.store, task.id, activeEntry.session);
 
               // Reduced-prompt retry succeeded — return to let the finally block clean up
               // without marking the task as failed.
@@ -3881,6 +3886,7 @@ and show an appropriate message to the user.\`
       );
 
       checkSessionError(session);
+      await accumulateSessionTokenUsage(this.store, task.id, session);
       session.dispose();
       await agentLogger.flush();
 

packages/engine/src/merger.ts

@@ -128,6 +128,7 @@ import { join } from "node:path";
 import { getTaskMergeBlocker, type TaskStore, type MergeResult, type MergeDetails, type WorkflowStep, type WorkflowStepResult, type Settings, type AgentPromptsConfig } from "@fusion/core";
 import { resolveAgentPrompt } from "@fusion/core";
 import { describeModel, promptWithFallback } from "./pi.js";
+import { accumulateSessionTokenUsage } from "./session-token-usage.js";
 import { createResolvedAgentSession, extractRuntimeHint } from "./agent-session-helpers.js";
 import { buildSessionSkillContext } from "./session-skill-context.js";
 import type { WorktreePool } from "./worktree-pool.js";
@@ -976,6 +977,7 @@ ${failureContext.output.slice(0, VERIFICATION_LOG_MAX_CHARS)}
         },
         signal: options.signal,
       });
+      await accumulateSessionTokenUsage(store, taskId, session);
 
       // Re-run deterministic verification command after the fix attempt.
       await store.logEntry(
@@ -1843,6 +1845,7 @@ You are assisting with a paused \`git pull --rebase\`.
       },
       signal: options?.signal,
     });
+    await accumulateSessionTokenUsage(store, taskId, session);
   } finally {
     session.dispose();
   }
@@ -3634,6 +3637,7 @@ async function runAiAgentForCommit(params: AiAgentParams): Promise<{ success: bo
 
     throw err;
   } finally {
+    await accumulateSessionTokenUsage(store, taskId, session);
     await agentLogger.flush();
     session.dispose();
   }
@@ -4022,6 +4026,7 @@ If issues are found that need attention, describe them clearly.`;
     );
 
     checkSessionError(session);
+    await accumulateSessionTokenUsage(store, taskId, session);
     session.dispose();
     await agentLogger.flush();
 

packages/engine/src/session-token-usage.ts

@@ -0,0 +1,92 @@
+import type { TaskStore } from "@fusion/core";
+import type { AgentSession } from "@mariozechner/pi-coding-agent";
+import { createLogger } from "./logger.js";
+
+const log = createLogger("session-token-usage");
+
+interface SessionBaseline {
+  input: number;
+  output: number;
+  cached: number;
+}
+
+// Per-session cumulative-token baselines so repeated calls only persist deltas.
+// The session object is keyed weakly so disposed sessions get garbage-collected.
+const sessionBaselines = new WeakMap<AgentSession, SessionBaseline>();
+
+interface SessionStatsLike {
+  tokens?: {
+    input?: number;
+    output?: number;
+    cacheRead?: number;
+    cacheWrite?: number;
+  };
+}
+
+function readSessionStats(session: AgentSession): SessionStatsLike | undefined {
+  const accessor = (session as unknown as { getSessionStats?: () => SessionStatsLike }).getSessionStats;
+  if (typeof accessor !== "function") return undefined;
+  try {
+    return accessor.call(session);
+  } catch {
+    return undefined;
+  }
+}
+
+/**
+ * Capture the session's cumulative token usage and accumulate any *new* deltas
+ * onto `task.tokenUsage`. Safe to call repeatedly on the same session — each
+ * call only persists what's been added since the previous call (per-session
+ * baseline tracking). Failures are logged and swallowed so token bookkeeping
+ * never blocks the task pipeline.
+ */
+export async function accumulateSessionTokenUsage(
+  store: TaskStore,
+  taskId: string,
+  session: AgentSession,
+): Promise<void> {
+  try {
+    const stats = readSessionStats(session);
+    const tokens = stats?.tokens;
+    if (!tokens) return;
+
+    // Treat cache-write tokens as input (they're billed as input on first write
+    // and read back at a discount on subsequent turns).
+    const currentInput = (tokens.input ?? 0) + (tokens.cacheWrite ?? 0);
+    const currentOutput = tokens.output ?? 0;
+    const currentCached = tokens.cacheRead ?? 0;
+
+    const baseline = sessionBaselines.get(session) ?? { input: 0, output: 0, cached: 0 };
+    const inputDelta = Math.max(0, currentInput - baseline.input);
+    const outputDelta = Math.max(0, currentOutput - baseline.output);
+    const cachedDelta = Math.max(0, currentCached - baseline.cached);
+
+    sessionBaselines.set(session, {
+      input: currentInput,
+      output: currentOutput,
+      cached: currentCached,
+    });
+
+    if (inputDelta === 0 && outputDelta === 0 && cachedDelta === 0) return;
+
+    const task = await store.getTask(taskId);
+    const now = new Date().toISOString();
+    const newInput = (task.tokenUsage?.inputTokens ?? 0) + inputDelta;
+    const newOutput = (task.tokenUsage?.outputTokens ?? 0) + outputDelta;
+    const newCached = (task.tokenUsage?.cachedTokens ?? 0) + cachedDelta;
+
+    await store.updateTask(taskId, {
+      tokenUsage: {
+        inputTokens: newInput,
+        outputTokens: newOutput,
+        cachedTokens: newCached,
+        totalTokens: newInput + newOutput + newCached,
+        firstUsedAt: task.tokenUsage?.firstUsedAt ?? now,
+        lastUsedAt: now,
+      },
+    });
+  } catch (err) {
+    const message = err instanceof Error ? err.message : String(err);
+    log.warn(`${taskId}: session token usage accumulate failed: ${message}`);
+  }
+}