diff --git a/agents/claude.ts b/agents/claude.ts index 626bfbb..d23f84a 100644 --- a/agents/claude.ts +++ b/agents/claude.ts @@ -33,8 +33,10 @@ import { agent, buildCommitPrompt, getGitStatus, + logTokenTable, MAX_COMMIT_RETRIES, MAX_STDERR_LINES, + mergeAgentUsage, } from "./shared.ts"; async function installClaudeCli(): Promise { @@ -192,6 +194,10 @@ async function runClaude(params: RunParams): Promise { let finalOutput = ""; let sessionId: string | undefined; let accumulatedTokens = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }; + // Claude CLI reports a single end-of-run `total_cost_usd` on the result + // event. per-message events don't carry cost, so there's nothing to sum — + // we just capture the final value when it arrives. + let accumulatedCostUsd = 0; let tokensLogged = false; function buildUsage(): AgentUsage | undefined { @@ -204,6 +210,7 @@ async function runClaude(params: RunParams): Promise { outputTokens: accumulatedTokens.output, cacheReadTokens: accumulatedTokens.cacheRead || undefined, cacheWriteTokens: accumulatedTokens.cacheWrite || undefined, + costUsd: accumulatedCostUsd > 0 ? accumulatedCostUsd : undefined, } : undefined; } @@ -245,11 +252,15 @@ async function runClaude(params: RunParams): Promise { } } - // accumulate per-message usage if available + // accumulate per-message usage if available. capture cache fields too + // so the fallback token table (used when no final `result` event fires) + // still reports the full breakdown instead of silently dropping cache. const msgUsage = event.message?.usage; if (msgUsage) { accumulatedTokens.input += msgUsage.input_tokens || 0; accumulatedTokens.output += msgUsage.output_tokens || 0; + accumulatedTokens.cacheRead += msgUsage.cache_read_input_tokens || 0; + accumulatedTokens.cacheWrite += msgUsage.cache_creation_input_tokens || 0; } }, user: (event: ClaudeUserEvent) => { @@ -290,28 +301,35 @@ async function runClaude(params: RunParams): Promise { const numTurns = event.num_turns || 0; if (subtype === "success") { - // extract detailed usage from result event (most accurate source) + // extract detailed usage from result event (most accurate source). + // note: `input` here is non-cached input tokens only, matching the + // semantics of OpenCode's step_finish.tokens.input — the logTokenTable + // helper sums Input + Cache Read + Cache Write + Output into the Total + // column so consumers get the real billable figure. const usage = event.usage; const inputTokens = usage?.input_tokens || 0; const cacheRead = usage?.cache_read_input_tokens || 0; const cacheWrite = usage?.cache_creation_input_tokens || 0; const outputTokens = usage?.output_tokens || 0; - const totalInput = inputTokens + cacheRead + cacheWrite; + // guard against NaN/Infinity from malformed CLI output poisoning the total + const costUsd = + typeof event.total_cost_usd === "number" && Number.isFinite(event.total_cost_usd) + ? event.total_cost_usd + : 0; accumulatedTokens = { input: inputTokens, output: outputTokens, cacheRead, cacheWrite }; + accumulatedCostUsd = costUsd; log.info(`» ${params.label} result: subtype=${subtype}, turns=${numTurns}`); if (!tokensLogged) { - log.table([ - [ - { data: "Input", header: true }, - { data: "Cache Read", header: true }, - { data: "Cache Write", header: true }, - { data: "Output", header: true }, - ], - [String(totalInput), String(cacheRead), String(cacheWrite), String(outputTokens)], - ]); + logTokenTable({ + input: inputTokens, + cacheRead, + cacheWrite, + output: outputTokens, + costUsd, + }); tokensLogged = true; } } else if (subtype === "error_max_turns") { @@ -432,16 +450,15 @@ async function runClaude(params: RunParams): Promise { if (stderrContext) log.info(`» last stderr output:\n${stderrContext}`); } - if (!tokensLogged && (accumulatedTokens.input > 0 || accumulatedTokens.output > 0)) { - const totalTokens = accumulatedTokens.input + accumulatedTokens.output; - log.table([ - [ - { data: "Input Tokens", header: true }, - { data: "Output Tokens", header: true }, - { data: "Total Tokens", header: true }, - ], - [String(accumulatedTokens.input), String(accumulatedTokens.output), String(totalTokens)], - ]); + if ( + !tokensLogged && + (accumulatedTokens.input > 0 || + accumulatedTokens.output > 0 || + accumulatedTokens.cacheRead > 0 || + accumulatedTokens.cacheWrite > 0) + ) { + logTokenTable({ ...accumulatedTokens, costUsd: accumulatedCostUsd }); + tokensLogged = true; } const usage = buildUsage(); @@ -637,6 +654,10 @@ export const claude = agent({ ...runParams, args: [...baseArgs, "-p", ctx.instructions.full], }); + // usage needs to aggregate across the initial run + every commit retry. + // each runClaude() returns only its own iteration's usage, so without + // merging the caller sees only the final retry's slice and undercounts. + let aggregatedUsage = result.usage; // post-run: if the working tree is dirty, resume the session and ask the agent to commit for (let attempt = 0; attempt < MAX_COMMIT_RETRIES; attempt++) { @@ -655,8 +676,9 @@ export const claude = agent({ result.sessionId, ], }); + aggregatedUsage = mergeAgentUsage(aggregatedUsage, result.usage); } - return result; + return { ...result, usage: aggregatedUsage }; }, }); diff --git a/agents/opencode.ts b/agents/opencode.ts index e44f00b..6f92f90 100644 --- a/agents/opencode.ts +++ b/agents/opencode.ts @@ -33,8 +33,10 @@ import { agent, buildCommitPrompt, getGitStatus, + logTokenTable, MAX_COMMIT_RETRIES, MAX_STDERR_LINES, + mergeAgentUsage, } from "./shared.ts"; async function installOpencodeCli(): Promise { @@ -268,6 +270,10 @@ async function runOpenCode(params: RunParams): Promise { let finalOutput = ""; let accumulatedTokens = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }; + // per-step `part.cost` sums across the whole session. sourced from models.dev + // inside opencode — present for every supported provider (Anthropic, OpenAI, + // Google, xAI, DeepSeek, Moonshot, OpenRouter sub-providers, etc.). + let accumulatedCostUsd = 0; let tokensLogged = false; const toolCallTimings = new Map(); let currentStepId: string | null = null; @@ -284,6 +290,7 @@ async function runOpenCode(params: RunParams): Promise { outputTokens: accumulatedTokens.output, cacheReadTokens: accumulatedTokens.cacheRead || undefined, cacheWriteTokens: accumulatedTokens.cacheWrite || undefined, + costUsd: accumulatedCostUsd > 0 ? accumulatedCostUsd : undefined, } : undefined; } @@ -296,6 +303,7 @@ async function runOpenCode(params: RunParams): Promise { log.debug(`» ${params.label} init event (full): ${JSON.stringify(event)}`); finalOutput = ""; accumulatedTokens = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }; + accumulatedCostUsd = 0; tokensLogged = false; }, message: (event: OpenCodeMessageEvent) => { @@ -340,6 +348,15 @@ async function runOpenCode(params: RunParams): Promise { accumulatedTokens.cacheRead += eventTokens.cache?.read || 0; accumulatedTokens.cacheWrite += eventTokens.cache?.write || 0; } + // step_finish.part.cost is a per-step delta (not a running total) — + // OpenCode emits varying per-event values that sum to the session cost. + // verified empirically across Anthropic, OpenAI, Gemini, xAI, DeepSeek, + // Moonshot, and OpenRouter (see pullfrog-baseline/opencode-*.log). + // guard against NaN/Infinity — a single poison value would make the + // running total un-recoverable for the rest of the session. + if (typeof event.part?.cost === "number" && Number.isFinite(event.part.cost)) { + accumulatedCostUsd += event.part.cost; + } if (currentStepId === stepId) { currentStepId = null; currentStepType = null; @@ -429,20 +446,19 @@ async function runOpenCode(params: RunParams): Promise { if (event.status === "error") { log.info(`» ${params.label} failed: ${JSON.stringify(event)}`); } else { - const inputTokens = event.stats?.input_tokens || accumulatedTokens.input || 0; - const outputTokens = event.stats?.output_tokens || accumulatedTokens.output || 0; - const totalTokens = event.stats?.total_tokens || inputTokens + outputTokens; + // the final `result` event only carries input_tokens/output_tokens and + // no cache breakdown — accumulatedTokens (summed across step_finish + // events) is strictly more accurate, so we prefer it unconditionally. log.info(`» run complete: tool_calls=${toolCalls}, duration=${duration}ms`); - if ((inputTokens > 0 || outputTokens > 0) && !tokensLogged) { - log.table([ - [ - { data: "Input Tokens", header: true }, - { data: "Output Tokens", header: true }, - { data: "Total Tokens", header: true }, - ], - [String(inputTokens), String(outputTokens), String(totalTokens)], - ]); + if ( + (accumulatedTokens.input > 0 || + accumulatedTokens.output > 0 || + accumulatedTokens.cacheRead > 0 || + accumulatedTokens.cacheWrite > 0) && + !tokensLogged + ) { + logTokenTable({ ...accumulatedTokens, costUsd: accumulatedCostUsd }); tokensLogged = true; } } @@ -555,16 +571,15 @@ async function runOpenCode(params: RunParams): Promise { if (stderrContext) log.info(`» last stderr output:\n${stderrContext}`); } - if (!tokensLogged && (accumulatedTokens.input > 0 || accumulatedTokens.output > 0)) { - const totalTokens = accumulatedTokens.input + accumulatedTokens.output; - log.table([ - [ - { data: "Input Tokens", header: true }, - { data: "Output Tokens", header: true }, - { data: "Total Tokens", header: true }, - ], - [String(accumulatedTokens.input), String(accumulatedTokens.output), String(totalTokens)], - ]); + if ( + !tokensLogged && + (accumulatedTokens.input > 0 || + accumulatedTokens.output > 0 || + accumulatedTokens.cacheRead > 0 || + accumulatedTokens.cacheWrite > 0) + ) { + logTokenTable({ ...accumulatedTokens, costUsd: accumulatedCostUsd }); + tokensLogged = true; } const usage = buildUsage(); @@ -688,6 +703,10 @@ export const opencode = agent({ ...runParams, args: [...baseArgs, ctx.instructions.full], }); + // usage needs to aggregate across the initial run + every commit retry. + // each runOpenCode() returns only its own iteration's usage, so without + // merging the caller sees only the final retry's slice and undercounts. + let aggregatedUsage = result.usage; // post-run: if the working tree is dirty, continue the session and ask the agent to commit for (let attempt = 0; attempt < MAX_COMMIT_RETRIES; attempt++) { @@ -700,8 +719,9 @@ export const opencode = agent({ ...runParams, args: [...baseArgs, "--continue", buildCommitPrompt("opencode", status)], }); + aggregatedUsage = mergeAgentUsage(aggregatedUsage, result.usage); } - return result; + return { ...result, usage: aggregatedUsage }; }, }); diff --git a/agents/shared.test.ts b/agents/shared.test.ts new file mode 100644 index 0000000..1a1626e --- /dev/null +++ b/agents/shared.test.ts @@ -0,0 +1,76 @@ +import { describe, expect, it } from "vitest"; +import { type AgentUsage, mergeAgentUsage } from "./shared.ts"; + +const entry = (overrides: Partial): AgentUsage => ({ + agent: "pullfrog", + inputTokens: 0, + outputTokens: 0, + ...overrides, +}); + +describe("mergeAgentUsage", () => { + it("returns undefined when both sides are undefined", () => { + expect(mergeAgentUsage(undefined, undefined)).toBeUndefined(); + }); + + it("returns a copy of b when a is undefined", () => { + const b = entry({ inputTokens: 10 }); + expect(mergeAgentUsage(undefined, b)).toEqual(b); + }); + + it("returns a copy of a when b is undefined", () => { + const a = entry({ inputTokens: 10 }); + expect(mergeAgentUsage(a, undefined)).toEqual(a); + }); + + it("sums inputTokens and outputTokens unconditionally", () => { + const merged = mergeAgentUsage( + entry({ inputTokens: 10, outputTokens: 5 }), + entry({ inputTokens: 20, outputTokens: 7 }) + ); + expect(merged?.inputTokens).toBe(30); + expect(merged?.outputTokens).toBe(12); + }); + + it("keeps cache/cost fields undefined when both sides lack them", () => { + // this matters so downstream aggregateUsage doesn't persist spurious 0s into the DB + const merged = mergeAgentUsage(entry({ inputTokens: 10 }), entry({ inputTokens: 20 })); + expect(merged?.cacheReadTokens).toBeUndefined(); + expect(merged?.cacheWriteTokens).toBeUndefined(); + expect(merged?.costUsd).toBeUndefined(); + }); + + it("sums cache and cost fields when either side reports them", () => { + const merged = mergeAgentUsage( + entry({ inputTokens: 10, cacheReadTokens: 100, costUsd: 0.01 }), + entry({ inputTokens: 20, cacheWriteTokens: 50, costUsd: 0.02 }) + ); + expect(merged?.cacheReadTokens).toBe(100); + expect(merged?.cacheWriteTokens).toBe(50); + expect(merged?.costUsd).toBeCloseTo(0.03, 10); + }); + + it("preserves the agent id of the left operand", () => { + // the aggregator is called inside a single agent's run() — the agent label + // is a fixed property of the harness, not something that can flip mid-run + const merged = mergeAgentUsage( + entry({ agent: "claude", inputTokens: 10 }), + entry({ agent: "something-else", inputTokens: 20 }) + ); + expect(merged?.agent).toBe("claude"); + }); + + it("returns a fresh object rather than the input reference", () => { + // callers treat AgentUsage as immutable; returning the input itself would + // leak that invariant. mutating the returned value must not affect inputs. + const a = entry({ inputTokens: 10 }); + const mergedWithUndef = mergeAgentUsage(a, undefined); + expect(mergedWithUndef).not.toBe(a); + expect(mergedWithUndef).toEqual(a); + + const b = entry({ inputTokens: 20 }); + const mergedFromUndef = mergeAgentUsage(undefined, b); + expect(mergedFromUndef).not.toBe(b); + expect(mergedFromUndef).toEqual(b); + }); +}); diff --git a/agents/shared.ts b/agents/shared.ts index 024b686..2614366 100644 --- a/agents/shared.ts +++ b/agents/shared.ts @@ -34,10 +34,21 @@ export function buildCommitPrompt(_agentId: AgentId, status: string): string { } /** - * token/cost usage data from a single agent run + * token/cost usage data from a single agent run. + * + * NOTE on semantics: `inputTokens` here is the *total* billable input for the + * run — non-cached input + cache read + cache write — matching the per-agent + * SDK conventions. This is what gets persisted to `WorkflowRun.inputTokens`. + * + * The stdout token table and markdown step summary display a different "Input" + * column that shows only the non-cached portion (derivable as + * `inputTokens - cacheReadTokens - cacheWriteTokens`) so humans can see the + * cache hit ratio at a glance. Dashboards that query `WorkflowRun.inputTokens` + * directly are seeing the full total, not the log column. */ export interface AgentUsage { agent: string; + /** full billable input: non-cached + cache read + cache write */ inputTokens: number; outputTokens: number; cacheReadTokens?: number | undefined; @@ -95,3 +106,96 @@ export const agent = (input: Agent): Agent => { }, }; }; + +/** format a USD cost to 4 decimal places, always showing the leading zero */ +export function formatCostUsd(costUsd: number): string { + return costUsd.toFixed(4); +} + +/** + * merge two AgentUsage snapshots into one running total. + * + * both agent harnesses invoke their runner multiple times per `run()` when the + * post-run dirty-tree loop kicks in (MAX_COMMIT_RETRIES). each invocation + * produces its own AgentUsage; we sum them so downstream callers (usage + * summary, WorkflowRun persistence) see the whole session — not just the + * final retry's slice. + * + * returns `undefined` when both sides are empty so callers can short-circuit + * without a special case. zero-valued cache / cost fields are dropped to + * `undefined` for symmetry with each harness's `buildUsage`. + */ +export function mergeAgentUsage( + a: AgentUsage | undefined, + b: AgentUsage | undefined +): AgentUsage | undefined { + // always return a fresh object — callers treat AgentUsage as immutable, and + // returning `a` / `b` directly would leak that invariant to future callers + if (!a && !b) return undefined; + if (!a) return { ...(b as AgentUsage) }; + if (!b) return { ...a }; + const cacheRead = (a.cacheReadTokens ?? 0) + (b.cacheReadTokens ?? 0); + const cacheWrite = (a.cacheWriteTokens ?? 0) + (b.cacheWriteTokens ?? 0); + const cost = (a.costUsd ?? 0) + (b.costUsd ?? 0); + return { + agent: a.agent, + inputTokens: a.inputTokens + b.inputTokens, + outputTokens: a.outputTokens + b.outputTokens, + cacheReadTokens: cacheRead > 0 ? cacheRead : undefined, + cacheWriteTokens: cacheWrite > 0 ? cacheWrite : undefined, + costUsd: cost > 0 ? cost : undefined, + }; +} + +/** + * unified per-run token table used by every agent harness. + * + * columns are kept stable across agents and models so downstream log parsers + * (scripts/token-usage.ts, cost dashboards) only have to understand one format: + * + * Input non-cached input tokens sent this run + * Cache Read input tokens served from prompt cache (Anthropic, etc.) + * Cache Write input tokens written to prompt cache this run + * Output assistant output tokens + * Total sum of the four columns — the real billable quantity + * Cost ($) USD cost reported by the provider (only rendered when known) + * + * models that don't report prompt caching leave Cache Read / Write at 0. + * OpenCode emits per-step `part.cost` sourced from models.dev (works across + * Anthropic, OpenAI, Google, xAI, DeepSeek, Moonshot, OpenRouter, etc.); + * Claude CLI emits `total_cost_usd` on its final `result` event. pass the + * accumulated value via `costUsd` to render the Cost column. + */ +export function logTokenTable(t: { + input: number; + cacheRead: number; + cacheWrite: number; + output: number; + costUsd?: number | undefined; +}): void { + const total = t.input + t.cacheRead + t.cacheWrite + t.output; + // narrow costUsd to a concrete number so the render path doesn't need a cast + const costUsd = typeof t.costUsd === "number" && t.costUsd > 0 ? t.costUsd : undefined; + + const headerRow: Array<{ data: string; header: true }> = [ + { data: "Input", header: true }, + { data: "Cache Read", header: true }, + { data: "Cache Write", header: true }, + { data: "Output", header: true }, + { data: "Total", header: true }, + ]; + const dataRow: string[] = [ + String(t.input), + String(t.cacheRead), + String(t.cacheWrite), + String(t.output), + String(total), + ]; + + if (costUsd !== undefined) { + headerRow.push({ data: "Cost ($)", header: true }); + dataRow.push(formatCostUsd(costUsd)); + } + + log.table([headerRow, dataRow]); +} diff --git a/main.ts b/main.ts index b734232..c614653 100644 --- a/main.ts +++ b/main.ts @@ -32,6 +32,7 @@ import { createOctokit, writeGitHubUsageSummaryToFile } from "./utils/github.ts" import { resolveInstructions } from "./utils/instructions.ts"; import { executeLifecycleHook } from "./utils/lifecycle.ts"; import { normalizeEnv } from "./utils/normalizeEnv.ts"; +import { aggregateUsage, patchWorkflowRunFields } from "./utils/patchWorkflowRunFields.ts"; import { resolvePayload, resolvePromptInput } from "./utils/payload.ts"; import { postReviewCleanup } from "./utils/reviewCleanup.ts"; import { handleAgentResult } from "./utils/run.ts"; @@ -623,5 +624,23 @@ export async function main(): Promise { ); } } + + // persist aggregated token + cost usage to the WorkflowRun row. + // this is the single shared cleanup path across every agent implementation: + // each agent harness returns a single AgentUsage from agent.run() that + // already aggregates its internal retries via mergeAgentUsage, and the + // success branch above pushes that entry into toolState.usageEntries. + // aggregateUsage sums across those entries (one per agent.run()). + // + // caveat: if the agent promise rejected (timeout or uncaught throw) the + // usage was never pushed, so nothing gets persisted for that run. runs + // that returned AgentResult with success=false still report their partial + // usage because the harness populates AgentUsage before returning. + if (toolContext) { + const patch = aggregateUsage(toolState.usageEntries); + if (Object.keys(patch).length > 0) { + await patchWorkflowRunFields(toolContext, patch); + } + } } } diff --git a/play.ts b/play.ts index 162c144..8eeef8e 100644 --- a/play.ts +++ b/play.ts @@ -1,6 +1,6 @@ import { execSync } from "node:child_process"; import { mkdtemp } from "node:fs/promises"; -import { tmpdir } from "node:os"; +import { devNull, tmpdir } from "node:os"; import { dirname, join, resolve } from "node:path"; import { fileURLToPath, pathToFileURL } from "node:url"; import arg from "arg"; @@ -37,6 +37,16 @@ config({ path: join(__dirname, "..", ".env") }); export async function run(inputsOrPrompt: Inputs | string): Promise { await ensureGitHubToken(); + // play.ts is a CI-emulator — isolate it from the developer's user- and + // system-scope gitconfig so checks like `validatePushDestination` see the + // raw stored remote URL instead of values mutated by `url.*.insteadOf` + // rewrites (a common SSH-auth convenience on dev boxes). CI runners have + // empty gitconfigs so this is a no-op there; locally it makes `pnpm play` + // and real runs produce identical git state. `os.devNull` canonicalizes + // the null device across Unix (`/dev/null`) and Windows (`\\.\nul`). + process.env.GIT_CONFIG_GLOBAL = devNull; + process.env.GIT_CONFIG_SYSTEM = devNull; + // create unique temp directory path in OS temp location for parallel execution // use a parent dir from mkdtemp, then clone into a 'repo' subdirectory const tempParent = await mkdtemp(join(tmpdir(), "pullfrog-play-")); diff --git a/utils/log.ts b/utils/log.ts index 43bd354..42ecac3 100644 --- a/utils/log.ts +++ b/utils/log.ts @@ -5,7 +5,7 @@ import { AsyncLocalStorage } from "node:async_hooks"; import * as core from "@actions/core"; import { table } from "table"; -import type { AgentUsage } from "../agents/shared.ts"; +import { type AgentUsage, formatCostUsd } from "../agents/shared.ts"; import { isGitHubActions, isInsideDocker } from "./globals.ts"; // --- log prefix via AsyncLocalStorage --- @@ -334,28 +334,49 @@ export function formatIndentedField(label: string, content: string): string { } /** - * format aggregated usage data as a markdown table for the GitHub step summary + * format aggregated usage data as a markdown table for the GitHub step summary. + * + * columns mirror the per-run stdout token table emitted by `logTokenTable` + * (Input / Cache Read / Cache Write / Output / Total / Cost ($)) so the job + * summary and the in-run logs can be compared row-for-row. + * + * notes: + * - `AgentUsage.inputTokens` is the sum of non-cached input + cache read + * + cache write (set that way by both agent harnesses' `buildUsage`), + * so the non-cached Input column is recovered by subtracting cache fields. + * - `costUsd` is sourced from models.dev (OpenCode) or `total_cost_usd` + * (Claude CLI). absent rows show `—` so per-agent coverage is obvious. */ export function formatUsageSummary(entries: AgentUsage[]): string { if (entries.length === 0) return ""; - const header = "| Agent | Input | Output | Cache Read | Cache Write |"; - const separatorRow = "| --- | ---: | ---: | ---: | ---: |"; + const header = "| Agent | Input | Cache Read | Cache Write | Output | Total | Cost ($) |"; + const separatorRow = "| --- | ---: | ---: | ---: | ---: | ---: | ---: |"; const fmt = (n: number) => n.toLocaleString("en-US"); + const nonCachedInput = (e: AgentUsage): number => + Math.max(0, e.inputTokens - (e.cacheReadTokens ?? 0) - (e.cacheWriteTokens ?? 0)); + const totalFor = (e: AgentUsage): number => + nonCachedInput(e) + (e.cacheReadTokens ?? 0) + (e.cacheWriteTokens ?? 0) + e.outputTokens; + const costCell = (e: AgentUsage): string => + typeof e.costUsd === "number" && e.costUsd > 0 ? formatCostUsd(e.costUsd) : "—"; + const rows = entries.map( (e) => - `| ${e.agent} | ${fmt(e.inputTokens)} | ${fmt(e.outputTokens)} | ${fmt(e.cacheReadTokens ?? 0)} | ${fmt(e.cacheWriteTokens ?? 0)} |` + `| ${e.agent} | ${fmt(nonCachedInput(e))} | ${fmt(e.cacheReadTokens ?? 0)} | ${fmt(e.cacheWriteTokens ?? 0)} | ${fmt(e.outputTokens)} | ${fmt(totalFor(e))} | ${costCell(e)} |` ); const totalsRows: string[] = []; if (entries.length > 1) { - const totalInput = entries.reduce((sum, e) => sum + e.inputTokens, 0); + const totalInput = entries.reduce((sum, e) => sum + nonCachedInput(e), 0); const totalOutput = entries.reduce((sum, e) => sum + e.outputTokens, 0); const totalCacheRead = entries.reduce((sum, e) => sum + (e.cacheReadTokens ?? 0), 0); const totalCacheWrite = entries.reduce((sum, e) => sum + (e.cacheWriteTokens ?? 0), 0); + const grandTotal = totalInput + totalCacheRead + totalCacheWrite + totalOutput; + const totalCostUsd = entries.reduce((sum, e) => sum + (e.costUsd ?? 0), 0); + const totalCostCell = totalCostUsd > 0 ? `**${formatCostUsd(totalCostUsd)}**` : "—"; totalsRows.push( - `| **Total** | **${fmt(totalInput)}** | **${fmt(totalOutput)}** | **${fmt(totalCacheRead)}** | **${fmt(totalCacheWrite)}** |` + `| **Total** | **${fmt(totalInput)}** | **${fmt(totalCacheRead)}** | **${fmt(totalCacheWrite)}** | **${fmt(totalOutput)}** | **${fmt(grandTotal)}** | ${totalCostCell} |` ); } diff --git a/utils/patchWorkflowRunFields.test.ts b/utils/patchWorkflowRunFields.test.ts new file mode 100644 index 0000000..96db1ad --- /dev/null +++ b/utils/patchWorkflowRunFields.test.ts @@ -0,0 +1,92 @@ +import { describe, expect, it } from "vitest"; +import type { AgentUsage } from "../agents/shared.ts"; +import { aggregateUsage } from "./patchWorkflowRunFields.ts"; + +const entry = (overrides: Partial): AgentUsage => ({ + agent: "pullfrog", + inputTokens: 0, + outputTokens: 0, + ...overrides, +}); + +describe("aggregateUsage", () => { + it("returns empty object for empty input", () => { + expect(aggregateUsage([])).toEqual({}); + }); + + it("drops fields that sum to zero so NULL stays 'not reported'", () => { + // a run that only recorded input tokens shouldn't write zero into output/cache/cost — + // those columns stay NULL so dashboards can tell 'zero' from 'never reported'. + expect(aggregateUsage([entry({ inputTokens: 42 })])).toEqual({ inputTokens: 42 }); + }); + + it("sums a single entry with all fields present", () => { + expect( + aggregateUsage([ + entry({ + inputTokens: 100, + outputTokens: 50, + cacheReadTokens: 1000, + cacheWriteTokens: 200, + costUsd: 0.12, + }), + ]) + ).toEqual({ + inputTokens: 100, + outputTokens: 50, + cacheReadTokens: 1000, + cacheWriteTokens: 200, + costUsd: 0.12, + }); + }); + + it("sums multiple entries across agents", () => { + expect( + aggregateUsage([ + entry({ + agent: "claude", + inputTokens: 100, + outputTokens: 50, + cacheReadTokens: 1000, + costUsd: 0.1, + }), + entry({ + agent: "pullfrog", + inputTokens: 200, + outputTokens: 80, + cacheReadTokens: 2000, + cacheWriteTokens: 300, + costUsd: 0.25, + }), + ]) + ).toEqual({ + inputTokens: 300, + outputTokens: 130, + cacheReadTokens: 3000, + cacheWriteTokens: 300, + // floating-point sum — specifying exact value documents expected precision + costUsd: 0.35, + }); + }); + + it("treats undefined cache/cost as zero and drops when the sum is still zero", () => { + expect( + aggregateUsage([ + entry({ inputTokens: 10, outputTokens: 5 }), + entry({ inputTokens: 20, outputTokens: 15 }), + ]) + ).toEqual({ inputTokens: 30, outputTokens: 20 }); + }); + + it("clamps individual INT fields at INT4_MAX so partial-persist cannot happen", () => { + // server-side per-field rejection would silently drop the huge column and + // keep the small ones, producing a row with a NULL for the missing metric. + // clamping client-side guarantees the wire payload is self-consistent. + const result = aggregateUsage([ + entry({ inputTokens: 3_000_000_000, outputTokens: 42, cacheReadTokens: 5 }), + ]); + expect(result.inputTokens).toBe(2_147_483_647); + expect(result.outputTokens).toBe(42); + expect(result.cacheReadTokens).toBe(5); + }); +}); diff --git a/utils/patchWorkflowRunFields.ts b/utils/patchWorkflowRunFields.ts index da502ee..df0c3e1 100644 --- a/utils/patchWorkflowRunFields.ts +++ b/utils/patchWorkflowRunFields.ts @@ -1,9 +1,14 @@ +import type { AgentUsage } from "../agents/shared.ts"; import type { ToolContext } from "../mcp/server.ts"; import { apiFetch } from "./apiFetch.ts"; import { log } from "./cli.ts"; import { retry } from "./retry.ts"; -/** Keys accepted by PATCH /api/workflow-run/[runId] — keep in sync with `ALLOWED_FIELDS` in `app/api/workflow-run/[runId]/route.ts`. */ +/** + * Artifact tracking fields — one-off PATCHes from MCP tools as GitHub entities + * are created during the run. Strings only (GraphQL node IDs). + * Keep in sync with `STRING_FIELDS` in `app/api/workflow-run/[runId]/route.ts`. + */ export type WorkflowRunArtifactPatchKey = | "prNodeId" | "issueNodeId" @@ -11,9 +16,23 @@ export type WorkflowRunArtifactPatchKey = | "planCommentNodeId" | "summaryCommentNodeId"; -export type WorkflowRunArtifactPatch = Partial>; +/** + * Usage fields — aggregated across all agent calls and PATCHed once at + * end-of-run. Token counts are Int4 on the DB side (ample for any realistic + * run); `costUsd` is a Decimal populated by provider-reported dollar amounts. + * Keep in sync with `INT_FIELDS` + `DECIMAL_FIELDS` in the server route. + */ +export type WorkflowRunUsagePatchKey = + | "inputTokens" + | "outputTokens" + | "cacheReadTokens" + | "cacheWriteTokens" + | "costUsd"; -const ARTIFACT_PATCH_KEYS: WorkflowRunArtifactPatchKey[] = [ +export type WorkflowRunPatch = Partial> & + Partial>; + +const STRING_KEYS: WorkflowRunArtifactPatchKey[] = [ "prNodeId", "issueNodeId", "reviewNodeId", @@ -21,19 +40,33 @@ const ARTIFACT_PATCH_KEYS: WorkflowRunArtifactPatchKey[] = [ "summaryCommentNodeId", ]; -/** PATCH workflow-run artifact fields (Pullfrog JWT, not GitHub). */ +const NUMBER_KEYS: WorkflowRunUsagePatchKey[] = [ + "inputTokens", + "outputTokens", + "cacheReadTokens", + "cacheWriteTokens", + "costUsd", +]; + +/** PATCH workflow-run fields (Pullfrog JWT, not GitHub). */ export async function patchWorkflowRunFields( ctx: ToolContext, - fields: WorkflowRunArtifactPatch + fields: WorkflowRunPatch ): Promise { if (ctx.runId === undefined || !ctx.apiToken) return; - const body: Record = {}; - for (const key of ARTIFACT_PATCH_KEYS) { + const body: Record = {}; + for (const key of STRING_KEYS) { const value = fields[key]; if (typeof value === "string" && value.length > 0) { body[key] = value; } } + for (const key of NUMBER_KEYS) { + const value = fields[key]; + if (typeof value === "number" && Number.isFinite(value) && value >= 0) { + body[key] = value; + } + } if (Object.keys(body).length === 0) return; try { await retry( @@ -60,3 +93,58 @@ export async function patchWorkflowRunFields( log.warning(`patchWorkflowRunFields exhausted retries: ${error}`); } } + +/** + * Postgres INTEGER / Prisma Int4 is signed 32-bit. Aggregated usage won't + * realistically hit this in a single run (2.1B tokens ≈ $6000+ of input on + * Claude Opus), but clamping here keeps the wire payload self-consistent: + * the server rejects out-of-range INT fields individually, so without a + * client-side clamp a single overflow would write a partial row where + * some columns land and others silently don't. + */ +const INT4_MAX = 2_147_483_647; + +function clampInt(value: number, field: WorkflowRunUsagePatchKey): number { + if (value > INT4_MAX) { + log.warning( + `aggregateUsage: ${field}=${value} exceeds INT4_MAX (${INT4_MAX}) — clamping so the rest of the usage row still persists.` + ); + return INT4_MAX; + } + return value; +} + +/** + * Sum per-agent usage entries into a single WorkflowRunPatch payload. + * Returns an empty object when there's nothing to report, which causes + * `patchWorkflowRunFields` to no-op — safe to call unconditionally from + * end-of-run paths. Zero-valued fields are dropped so the DB only stores + * positive sums (and NULL means "not reported"). + * + * Token sums are clamped to INT4_MAX to guarantee the payload the server + * sees is always self-consistent across all numeric columns. + */ +export function aggregateUsage(entries: AgentUsage[]): WorkflowRunPatch { + if (entries.length === 0) return {}; + + const sum = entries.reduce( + (acc, e) => ({ + inputTokens: acc.inputTokens + e.inputTokens, + outputTokens: acc.outputTokens + e.outputTokens, + cacheReadTokens: acc.cacheReadTokens + (e.cacheReadTokens ?? 0), + cacheWriteTokens: acc.cacheWriteTokens + (e.cacheWriteTokens ?? 0), + costUsd: acc.costUsd + (e.costUsd ?? 0), + }), + { inputTokens: 0, outputTokens: 0, cacheReadTokens: 0, cacheWriteTokens: 0, costUsd: 0 } + ); + + const out: WorkflowRunPatch = {}; + if (sum.inputTokens > 0) out.inputTokens = clampInt(sum.inputTokens, "inputTokens"); + if (sum.outputTokens > 0) out.outputTokens = clampInt(sum.outputTokens, "outputTokens"); + if (sum.cacheReadTokens > 0) + out.cacheReadTokens = clampInt(sum.cacheReadTokens, "cacheReadTokens"); + if (sum.cacheWriteTokens > 0) + out.cacheWriteTokens = clampInt(sum.cacheWriteTokens, "cacheWriteTokens"); + if (sum.costUsd > 0) out.costUsd = sum.costUsd; + return out; +}