refactor(cost): single writer for cost_entries, finishRun reads SUM

Cost was written through three independent paths: worker-runner shipped
`finalUsage.costUsd` in run-finished -> finishRun wrote `runs.cost_usd`,
`tasks.cost_usd += costUsd`, `agents.spent_monthly_usd += costUsd`, and
inserted a per-run aggregate row in cost_entries. CostWriter
additionally inserted per-message rows in cost_entries from main. Any
CostWriter failure created silent drift between `SUM(cost_entries)`
(read by budget enforcement) and `runs.cost_usd` (read by the UI).

Make CostWriter the single writer, finishRun the single reader.

- Add `db.recordCostEntry(...)` as the first-class insert path for
  per-message cost rows. CostWriter calls it; tests call it directly.
- Remove the aggregate INSERT from finishRun. finishRun now reads
  `SELECT COALESCE(SUM(amount_usd), 0) FROM cost_entries WHERE run_id=?
  AND unattributed=0` and writes that as the authoritative run total to
  runs.cost_usd, tasks.cost_usd, and agents.spent_monthly_usd.
- Drop `costUsd` from `finishRun` input, the worker `run-finished`
  ChildMessage shape, the automation `RunFinishedEvent` type, and every
  caller + test fixture. Cost is owned by the per-message path.
- worker-runner stops extracting `usage.cost.total` — token counts stay
  for session accumulation, cost comes from pi-ai via CostWriter.
- run-finished-helpers now gates budget incident checks on
  `event.status === "succeeded"` instead of a cost field that no longer
  exists on the event.
- automation-run-completion fires `budget_threshold` on every succeeded
  run (post-run spent_monthly_usd already reflects the SUM path).
- Update budget tests to write a `recordCostEntry` before `finishRun`
  to simulate the per-message cost path; drop `costUsd` from cross-
  subsystem, edge-case-budget, database, and related fixtures.

One writer, one reader, one source of truth. No drift.
This commit is contained in:
ZacharyZhang-NY committed 2026-04-17 15:04:22 -04:00
1 parent 77f81a7c5a
commit 14dc9a4897
26 files changed
+102 -93

No files matched your search

+7 -5
View File
@@ -65,10 +65,9 @@ export class WorkerRunner {
post({ type: "run-status", runId: job.runId, status: "running", message: `Running agent ${job.agentName}.` });
let resolvedModelId: string | null = job.model;
let finalUsage: { inputTokens: number | null; outputTokens: number | null; costUsd: number | null } = {
let finalUsage: { inputTokens: number | null; outputTokens: number | null } = {
inputTokens: null,
outputTokens: null,
costUsd: null,
};
let errorMessage: string | null = null;
@@ -177,7 +176,6 @@ export class WorkerRunner {
finalUsage = {
inputTokens: typeof maybeAssistant.usage.input === "number" ? maybeAssistant.usage.input : finalUsage.inputTokens,
outputTokens: typeof maybeAssistant.usage.output === "number" ? maybeAssistant.usage.output : finalUsage.outputTokens,
costUsd: typeof maybeAssistant.usage.cost?.total === "number" ? maybeAssistant.usage.cost.total : finalUsage.costUsd,
};
}
}
@@ -227,8 +225,12 @@ export class WorkerRunner {
model: resolvedModelId,
sessionParams: job.sessionParams,
sessionDisplayId: null,
costUsd: finalUsage.costUsd,
unattributedCost: finalUsage.costUsd === null,
// Worker no longer owns cost attribution — CostWriter (main-side)
// writes per-message rows, finishRun sums them. This flag stays true
// only for run paths that never had a chance to emit usage (errors
// before the first message_end); per-message drift detection belongs
// elsewhere.
unattributedCost: finalUsage.inputTokens === null,
inputTokens: finalUsage.inputTokens,
outputTokens: finalUsage.outputTokens,
});