diff --git a/src/agents/pi-embedded-runner/run/attempt.ts b/src/agents/pi-embedded-runner/run/attempt.ts index 3457fdf0161..3aedfcb84a9 100644 --- a/src/agents/pi-embedded-runner/run/attempt.ts +++ b/src/agents/pi-embedded-runner/run/attempt.ts @@ -135,6 +135,7 @@ import { import { pruneProcessedHistoryImages } from "./history-image-prune.js"; import { detectAndLoadPromptImages } from "./images.js"; import type { EmbeddedRunAttemptParams, EmbeddedRunAttemptResult } from "./types.js"; +import { recordTokenUsage } from "../../usage-log.js"; type PromptBuildHookRunner = { hasHooks: (hookName: "before_prompt_build" | "before_agent_start") => boolean; @@ -2758,6 +2759,19 @@ export async function runEmbeddedAttempt( }); } + recordTokenUsage({ + workspaceDir: params.workspaceDir, + runId: params.runId, + sessionId: params.sessionId, + sessionKey: params.sessionKey, + provider: params.provider, + model: params.modelId, + label: "llm_output", + usage: getUsageTotals(), + }).catch((err) => { + log.warn(`token usage log failed: ${String(err)}`); + }); + return { aborted, timedOut, diff --git a/src/agents/usage-log.ts b/src/agents/usage-log.ts new file mode 100644 index 00000000000..b3f63595793 --- /dev/null +++ b/src/agents/usage-log.ts @@ -0,0 +1,83 @@ +import fs from "fs/promises"; +import path from "path"; + +export type TokenUsageRecord = { + id: string; + taskId?: string; + label: string; + tokensUsed: number; + tokenLimit?: number; + inputTokens?: number; + outputTokens?: number; + cacheReadTokens?: number; + cacheWriteTokens?: number; + model?: string; + provider?: string; + runId?: string; + sessionId?: string; + sessionKey?: string; + createdAt: string; +}; + +function makeId() { + return `usage_${Date.now().toString(36)}_${Math.random().toString(36).slice(2, 8)}`; +} + +async function readJsonArray(file: string): Promise { + try { + const raw = await fs.readFile(file, "utf-8"); + const parsed = JSON.parse(raw); + return Array.isArray(parsed) ? (parsed as TokenUsageRecord[]) : []; + } catch { + return []; + } +} + +export async function recordTokenUsage(params: { + workspaceDir: string; + runId?: string; + sessionId?: string; + sessionKey?: string; + provider?: string; + model?: string; + label: string; + usage?: { + input?: number; + output?: number; + cacheRead?: number; + cacheWrite?: number; + total?: number; + }; +}) { + const usage = params.usage; + if (!usage) return; + const total = + usage.total ?? + (usage.input ?? 0) + (usage.output ?? 0) + (usage.cacheRead ?? 0) + (usage.cacheWrite ?? 0); + if (!total || total <= 0) return; + + const memoryDir = path.join(params.workspaceDir, "memory"); + const file = path.join(memoryDir, "token-usage.json"); + await fs.mkdir(memoryDir, { recursive: true }); + + const entry: TokenUsageRecord = { + id: makeId(), + taskId: params.runId, + label: params.label, + tokensUsed: Math.trunc(total), + ...(usage.input != null && usage.input > 0 && { inputTokens: Math.trunc(usage.input) }), + ...(usage.output != null && usage.output > 0 && { outputTokens: Math.trunc(usage.output) }), + ...(usage.cacheRead != null && usage.cacheRead > 0 && { cacheReadTokens: Math.trunc(usage.cacheRead) }), + ...(usage.cacheWrite != null && usage.cacheWrite > 0 && { cacheWriteTokens: Math.trunc(usage.cacheWrite) }), + model: params.model, + provider: params.provider, + runId: params.runId, + sessionId: params.sessionId, + sessionKey: params.sessionKey, + createdAt: new Date().toISOString(), + }; + + const records = await readJsonArray(file); + records.push(entry); + await fs.writeFile(file, JSON.stringify(records, null, 2)); +}