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>
This commit is contained in:
gsxdsm
2026-04-27 04:51:29 -07:00
parent 30df38efda
commit d30edfda5d
8 changed files with 334 additions and 52 deletions

View File

@@ -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();
});
});