diff --git a/packages/core/src/__tests__/agent-token-usage.test.ts b/packages/core/src/__tests__/agent-token-usage.test.ts new file mode 100644 index 000000000..6e651f8ff --- /dev/null +++ b/packages/core/src/__tests__/agent-token-usage.test.ts @@ -0,0 +1,72 @@ +import { afterEach, beforeEach, describe, expect, it } from "vitest"; +import { AgentStore } from "../agent-store.js"; +import { aggregateAgentTokenUsage } from "../agent-token-usage.js"; +import { createTaskStoreTestHarness } from "./store-test-helpers.js"; + +describe("aggregateAgentTokenUsage", () => { + const harness = createTaskStoreTestHarness(); + let agentStore: AgentStore; + + beforeEach(async () => { + await harness.beforeEach(); + agentStore = new AgentStore({ rootDir: harness.rootDir() }); + await agentStore.init(); + }); + + afterEach(async () => { + await harness.afterEach(); + }); + + it("returns null when agent does not exist", async () => { + const result = await aggregateAgentTokenUsage({ taskStore: harness.store(), agentStore, agentId: "missing" }); + expect(result).toBeNull(); + }); + + it("returns null for ephemeral agents", async () => { + const ephemeral = await agentStore.createAgent({ name: "temp", role: "executor", reportsTo: "FN-1", metadata: { type: "spawned" } }); + const result = await aggregateAgentTokenUsage({ taskStore: harness.store(), agentStore, agentId: ephemeral.id }); + expect(result).toBeNull(); + }); + + it("aggregates usage across windows", async () => { + const agent = await agentStore.createAgent({ name: "exec", role: "executor" }); + await harness.store().createTask({ + description: "recent", + assignedAgentId: agent.id, + tokenUsage: { + inputTokens: 100, + outputTokens: 10, + cachedTokens: 50, + cacheWriteTokens: 5, + totalTokens: 165, + firstUsedAt: "2026-05-13T09:00:00.000Z", + lastUsedAt: "2026-05-13T11:00:00.000Z", + }, + }); + await harness.store().createTask({ + description: "older", + assignedAgentId: agent.id, + tokenUsage: { + inputTokens: 40, + outputTokens: 4, + cachedTokens: 10, + cacheWriteTokens: 1, + totalTokens: 55, + firstUsedAt: "2026-05-05T09:00:00.000Z", + lastUsedAt: "2026-05-05T11:00:00.000Z", + }, + }); + + const result = await aggregateAgentTokenUsage({ + taskStore: harness.store(), + agentStore, + agentId: agent.id, + now: new Date("2026-05-13T12:00:00.000Z"), + }); + + expect(result).not.toBeNull(); + expect(result?.allTime).toMatchObject({ totalInputTokens: 140, totalCachedTokens: 60, totalCacheWriteTokens: 6, totalOutputTokens: 14, nTasks: 2 }); + expect(result?.last24h).toMatchObject({ totalInputTokens: 100, totalCachedTokens: 50, totalCacheWriteTokens: 5, totalOutputTokens: 10, nTasks: 1 }); + expect(result?.last7d).toMatchObject({ totalInputTokens: 100, totalCachedTokens: 50, totalCacheWriteTokens: 5, totalOutputTokens: 10, nTasks: 1 }); + }); +}); diff --git a/packages/core/src/agent-token-usage.ts b/packages/core/src/agent-token-usage.ts new file mode 100644 index 000000000..4496f7681 --- /dev/null +++ b/packages/core/src/agent-token-usage.ts @@ -0,0 +1,106 @@ +import type { AgentStore } from "./agent-store.js"; +import type { TaskStore } from "./store.js"; +import { isEphemeralAgent, type AgentRole } from "./types.js"; + +export interface AgentTokenUsageWindowSummary { + totalInputTokens: number; + totalCachedTokens: number; + totalCacheWriteTokens: number; + totalOutputTokens: number; + nTasks: number; + hitRatio: number; +} + +export interface AgentTokenUsageSummary { + agentId: string; + role: AgentRole; + last24h: AgentTokenUsageWindowSummary; + last7d: AgentTokenUsageWindowSummary; + allTime: AgentTokenUsageWindowSummary; +} + +export async function aggregateAgentTokenUsage({ + taskStore, + agentStore, + agentId, + now = new Date(), +}: { + taskStore: TaskStore; + agentStore: AgentStore; + agentId: string; + now?: Date; +}): Promise { + const agent = await agentStore.getAgent(agentId); + if (!agent || isEphemeralAgent(agent)) { + return null; + } + + const tasks = await taskStore.listTasks({ slim: true, includeArchived: true }); + const nowMs = now.getTime(); + const last24hMs = nowMs - (24 * 60 * 60 * 1000); + const last7dMs = nowMs - (7 * 24 * 60 * 60 * 1000); + + const allTime = createWindowSummary(); + const last24h = createWindowSummary(); + const last7d = createWindowSummary(); + + for (const task of tasks) { + if (!task.tokenUsage) continue; + const matchesAgent = task.assignedAgentId === agentId || task.sourceAgentId === agentId || task.checkedOutBy === agentId; + if (!matchesAgent) continue; + + const usage = task.tokenUsage; + applyTaskUsage(allTime, usage.inputTokens ?? 0, usage.cachedTokens ?? 0, usage.outputTokens ?? 0, usage.cacheWriteTokens ?? 0); + + const lastUsedAtMs = Date.parse(usage.lastUsedAt ?? ""); + if (!Number.isFinite(lastUsedAtMs)) continue; + + if (lastUsedAtMs >= last24hMs) { + applyTaskUsage(last24h, usage.inputTokens ?? 0, usage.cachedTokens ?? 0, usage.outputTokens ?? 0, usage.cacheWriteTokens ?? 0); + } + if (lastUsedAtMs >= last7dMs) { + applyTaskUsage(last7d, usage.inputTokens ?? 0, usage.cachedTokens ?? 0, usage.outputTokens ?? 0, usage.cacheWriteTokens ?? 0); + } + } + + return { + agentId, + role: agent.role as AgentRole, + last24h: finalizeWindowSummary(last24h), + last7d: finalizeWindowSummary(last7d), + allTime: finalizeWindowSummary(allTime), + }; +} + +function createWindowSummary(): AgentTokenUsageWindowSummary { + return { + totalInputTokens: 0, + totalCachedTokens: 0, + totalCacheWriteTokens: 0, + totalOutputTokens: 0, + nTasks: 0, + hitRatio: 0, + }; +} + +function applyTaskUsage( + summary: AgentTokenUsageWindowSummary, + inputTokens: number, + cachedTokens: number, + outputTokens: number, + cacheWriteTokens: number, +): void { + summary.totalInputTokens += inputTokens; + summary.totalCachedTokens += cachedTokens; + summary.totalCacheWriteTokens += cacheWriteTokens; + summary.totalOutputTokens += outputTokens; + summary.nTasks += 1; +} + +function finalizeWindowSummary(summary: AgentTokenUsageWindowSummary): AgentTokenUsageWindowSummary { + const denominator = summary.totalInputTokens + summary.totalCachedTokens; + return { + ...summary, + hitRatio: denominator > 0 ? summary.totalCachedTokens / denominator : 0, + }; +} diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index 4c4c98e59..046476818 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -87,6 +87,8 @@ export type { AgentProvisioningPolicyDecision, } from "./agent-provisioning-policy.js"; export { TaskStore } from "./store.js"; +export { aggregateAgentTokenUsage } from "./agent-token-usage.js"; +export type { AgentTokenUsageSummary, AgentTokenUsageWindowSummary } from "./agent-token-usage.js"; export { STALLED_REVIEW_REENQUEUE_THRESHOLD, STALLED_REVIEW_INVALID_TRANSITION_THRESHOLD,