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:
1 parent
77f81a7c5a
commit
14dc9a4897
26 files changed
+102
-93
No files matched your search
@@ -48,7 +48,6 @@ export function finalizeRun(
|
||||
signal: event.signal,
|
||||
model: event.model,
|
||||
sessionDisplayId: event.sessionDisplayId,
|
||||
costUsd: event.costUsd,
|
||||
unattributedCost: event.unattributedCost,
|
||||
inputTokens: event.inputTokens,
|
||||
outputTokens: event.outputTokens,
|
||||
@@ -186,10 +185,10 @@ export function runPostCompletionAutomation(
|
||||
const { destroyed, MAX_RETRY_ATTEMPTS, RETRY_COOLDOWN_MS, pendingRetryTimers } = context;
|
||||
const { run, task } = rfCtx;
|
||||
|
||||
// Policy-based incident check after run completion (Phase 5). Reuse the
|
||||
// `task` already resolved in finalizeRun instead of double-fetching with
|
||||
// its own catch-and-swallow; if task is null the project-scope is null.
|
||||
if (budgetService && typeof event.costUsd === "number" && Number.isFinite(event.costUsd)) {
|
||||
// Policy-based incident check after run completion. Always runs once the
|
||||
// run is finalized — cost is read from the cost_entries SUM inside
|
||||
// finishRun/checkAndCreateIncidents, not from the worker event.
|
||||
if (budgetService && event.status === "succeeded") {
|
||||
budgetService.checkAndCreateIncidents(run.companyId, run.agentId, task?.projectId ?? null);
|
||||
}
|
||||
|
||||
|
||||
@@ -42,7 +42,6 @@ export function handleWorkerExit(
|
||||
signal: null,
|
||||
model: null,
|
||||
sessionDisplayId: null,
|
||||
costUsd: null,
|
||||
unattributedCost: true,
|
||||
});
|
||||
if (run.workspaceId) {
|
||||
@@ -128,7 +127,6 @@ export function handleWorkerExit(
|
||||
signal: null,
|
||||
model: null,
|
||||
sessionDisplayId: null,
|
||||
costUsd: null,
|
||||
unattributedCost: true,
|
||||
});
|
||||
if (candidate.workspaceId) {
|
||||
|
||||
@@ -88,7 +88,6 @@ function cancelRun(runId: string) {
|
||||
model: null,
|
||||
sessionParams: null,
|
||||
sessionDisplayId: null,
|
||||
costUsd: null,
|
||||
unattributedCost: true,
|
||||
inputTokens: null,
|
||||
outputTokens: null,
|
||||
|
||||
@@ -43,7 +43,6 @@ export type ChildMessage =
|
||||
model: string | null;
|
||||
sessionParams: Record<string, unknown> | null;
|
||||
sessionDisplayId: string | null;
|
||||
costUsd: number | null;
|
||||
unattributedCost: boolean;
|
||||
inputTokens: number | null;
|
||||
outputTokens: number | null;
|
||||
|
||||
@@ -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,
|
||||
});
|
||||
|
||||
@@ -32,7 +32,6 @@ export const workerMessageSchema = z.union([
|
||||
model: z.string().nullable(),
|
||||
sessionParams: z.record(z.string(), z.unknown()).nullable(),
|
||||
sessionDisplayId: z.string().nullable(),
|
||||
costUsd: z.number().nullable(),
|
||||
unattributedCost: z.boolean(),
|
||||
inputTokens: z.number().nullable(),
|
||||
outputTokens: z.number().nullable(),
|
||||
|
||||
Reference in new issue
Block a user