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:
125
packages/engine/src/__tests__/session-token-usage.test.ts
Normal file
125
packages/engine/src/__tests__/session-token-usage.test.ts
Normal 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();
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user