From 14dc9a48974cac095d201037673b10a393bfe365 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9B=B7=E7=94=B5=E8=8A=BD=E8=A1=A3?= Date: Fri, 17 Apr 2026 15:04:22 -0400 Subject: [PATCH] refactor(cost): single writer for cost_entries, finishRun reads SUM MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- src/main/agent-runtime/cost-writer.ts | 41 +++++------- .../automation/automation-run-completion.ts | 6 +- src/main/automation/types.ts | 1 - src/main/database-base.ts | 3 +- src/main/database/runs.ts | 67 +++++++++++++------ src/main/orchestration/run-dispatch.ts | 1 - src/main/orchestration/task-runner.ts | 1 - src/main/worker-service.ts | 1 - .../worker/handlers/run-finished-helpers.ts | 9 ++- src/main/worker/worker-crash-recovery.ts | 2 - src/main/worker/worker-entry.ts | 1 - src/main/worker/worker-helpers.ts | 1 - src/main/worker/worker-runner.ts | 12 ++-- src/main/worker/worker-types.ts | 1 - .../integration/agent-model-switching.spec.ts | 1 - .../autonomous-chain/crash-recovery.spec.ts | 1 - .../autonomous-chain/run-failure.spec.ts | 1 - tests/integration/cross-subsystem.spec.ts | 8 +-- tests/integration/database.spec.ts | 1 - .../document-run-provenance.spec.ts | 1 - tests/unit/atomic-checkout.test.ts | 1 - tests/unit/database.test.ts | 13 +++- tests/unit/edge-case-budget.test.ts | 14 ++-- tests/unit/edge-case-session.test.ts | 2 - tests/unit/session-management.test.ts | 3 - tests/unit/store.test.ts | 2 - 26 files changed, 102 insertions(+), 93 deletions(-) diff --git a/src/main/agent-runtime/cost-writer.ts b/src/main/agent-runtime/cost-writer.ts index 2b9fc0b4..5d107b51 100644 --- a/src/main/agent-runtime/cost-writer.ts +++ b/src/main/agent-runtime/cost-writer.ts @@ -18,18 +18,18 @@ interface ObserveInput { } /** - * Writes a `cost_entries` row for an assistant `message_end` and updates the - * cumulative `runs.cost_usd` for the active run. + * Writes a `cost_entries` row for each pi `message_end`. finishRun reads + * SUM(cost_entries.amount_usd) for the run-level aggregate, so this is + * the single source of truth for cost attribution. */ export class CostWriter { constructor(private readonly deps: CostWriterDeps) {} observeAssistantMessage(input: ObserveInput): void { // Fail closed on cost calculation — budget policies read `cost_entries` - // to enforce spending gates; an "unattributed" row would silently let - // an agent run unbounded. If pricing cannot be resolved, refuse to - // persist the row and let the caller see the error. The failure is - // visible (throws) instead of invisible (null amount, unattributed=1). + // to enforce spending gates; silently emitting an unattributed row + // would let an agent run unbounded. If pricing cannot be resolved, + // throw so the caller logs and the missing cost is visible. const definition = getConnectorDefinition(input.connectorId); if (!definition.provider) { throw new Error(`[cost-writer] connector ${input.connectorId} has no provider mapping`); @@ -45,25 +45,14 @@ export class CostWriter { ); } - this.deps.db.queryAll( - "insert into cost_entries (id, company_id, run_id, task_id, agent_id, connector_id, amount_usd, unattributed, created_at) values (?, ?, ?, ?, ?, ?, ?, ?, ?)", - cryptoRandomId(), - input.companyId, - input.runId, - input.taskId, - input.agentId, - input.connectorId, - cost.total, - 0, - new Date().toISOString(), - ); + this.deps.db.recordCostEntry({ + runId: input.runId, + companyId: input.companyId, + agentId: input.agentId, + taskId: input.taskId, + connectorId: input.connectorId, + amountUsd: cost.total, + unattributed: false, + }); } } - -function cryptoRandomId(): string { - // Avoid pulling node:crypto at module top-level so this file can be - // imported from non-Node contexts during type checks. - return globalThis.crypto?.randomUUID - ? globalThis.crypto.randomUUID() - : `${Date.now()}-${Math.random().toString(36).slice(2)}`; -} diff --git a/src/main/automation/automation-run-completion.ts b/src/main/automation/automation-run-completion.ts index af38c9ba..15466bc4 100644 --- a/src/main/automation/automation-run-completion.ts +++ b/src/main/automation/automation-run-completion.ts @@ -270,7 +270,6 @@ export function processRunCompletion( taskId: run.taskId, agentId: run.agentId, runId: run.id, - costUsd: event.costUsd ?? undefined, }; try { const agent = db.getAgent(run.agentId); @@ -306,7 +305,10 @@ export function processRunCompletion( dispatchAutomation("run_failed", automationContext); } - if (event.costUsd && event.costUsd > 0) { + // Budget threshold check runs after every succeeded run; finishRun + // has already summed cost_entries and updated spentMonthlyUsd, so we + // just read the post-run spend and fire if it crosses 80%. + if (isSuccess) { try { const agent = db.getAgent(run.agentId); if (agent.budgetMonthlyUsd > 0) { diff --git a/src/main/automation/types.ts b/src/main/automation/types.ts index 1af2243d..d4d1c3e4 100644 --- a/src/main/automation/types.ts +++ b/src/main/automation/types.ts @@ -72,5 +72,4 @@ export interface RunFinishedEvent { status: RunStatus; summary: string | null; errorMessage: string | null; - costUsd: number | null; } diff --git a/src/main/database-base.ts b/src/main/database-base.ts index 1ca72841..4f8fc7bb 100644 --- a/src/main/database-base.ts +++ b/src/main/database-base.ts @@ -302,7 +302,8 @@ export class AppDatabaseBase { readRunLog(runId: string) { return this.runs.readRunLog(runId); } readRunLogChunk(runId: string, offset = 0, limit = 65_536) { return this.runs.readRunLogChunk(runId, offset, limit); } markRunStatus(runId: string, status: RunStatus, message: string) { return this.runs.markRunStatus(runId, status, message); } - finishRun(input: { runId: string; status: RunStatus; summary: string | null; errorMessage: string | null; exitCode: number | null; signal: string | null; model: string | null; sessionDisplayId: string | null; costUsd: number | null; unattributedCost: boolean }) { return this.runs.finishRun(input); } + finishRun(input: { runId: string; status: RunStatus; summary: string | null; errorMessage: string | null; exitCode: number | null; signal: string | null; model: string | null; sessionDisplayId: string | null; unattributedCost: boolean }) { return this.runs.finishRun(input); } + recordCostEntry(input: { runId: string; companyId: string; agentId: string; taskId: string | null; connectorId: string; amountUsd: number | null; unattributed: boolean }) { return this.runs.recordCostEntry(input); } hasActiveRun(filter: { companyId: string; taskId?: string; agentId?: string }) { return this.runs.hasActiveRun(filter); } listQueuedRuns() { return this.runs.listQueuedRuns(); } getInterruptedRunsForRecovery() { return this.runs.getInterruptedRunsForRecovery(); } diff --git a/src/main/database/runs.ts b/src/main/database/runs.ts index 7a7424ae..54b45cbc 100644 --- a/src/main/database/runs.ts +++ b/src/main/database/runs.ts @@ -158,7 +158,6 @@ export function buildRunMethods( signal: string | null; model: string | null; sessionDisplayId: string | null; - costUsd: number | null; unattributedCost: boolean; }) { const run = getDb().prepare("select task_id as taskId, agent_id as agentId, company_id as companyId, connector_id as connectorId from runs where id = ?").get(input.runId) as @@ -166,13 +165,20 @@ export function buildRunMethods( | undefined; if (!run) return { budgetHardStopped: false }; const ts = nowIso(); + + // Cost is owned by CostWriter — per-message rows in `cost_entries`. + // finishRun reads the aggregate instead of writing another row, so + // there is exactly one source of truth and `cost_entries.run_id` + // cannot double-count against the previous finishRun INSERT. + const costRow = getDb() + .prepare("select coalesce(sum(amount_usd), 0) as total from cost_entries where run_id = ? and unattributed = 0") + .get(input.runId) as { total: number } | undefined; + const runCostUsd = costRow?.total ?? 0; + const agent = getDb().prepare( "select budget_monthly_usd as budgetMonthlyUsd, spent_monthly_usd as spentMonthlyUsd, status from agents where id = ?", ).get(run.agentId) as { budgetMonthlyUsd: number; spentMonthlyUsd: number; status: string } | undefined; - const nextSpend = - agent && typeof input.costUsd === "number" && Number.isFinite(input.costUsd) - ? agent.spentMonthlyUsd + input.costUsd - : agent?.spentMonthlyUsd ?? 0; + const nextSpend = (agent?.spentMonthlyUsd ?? 0) + runCostUsd; const budgetHardStopped = Boolean( agent && agent.budgetMonthlyUsd > 0 && @@ -187,31 +193,21 @@ export function buildRunMethods( ) .run( input.status, input.summary ?? "", input.errorMessage, input.exitCode, - input.signal, input.model, input.sessionDisplayId, input.costUsd, + input.signal, input.model, input.sessionDisplayId, runCostUsd, input.unattributedCost ? 1 : 0, ts, ts, input.runId, ); getDb().prepare("update tasks set active_run_id = null, updated_at = ? where id = ?").run(ts, run.taskId); getDb() .prepare("UPDATE tasks SET checkout_run_id = NULL, execution_locked_at = NULL, updated_at = ? WHERE id = ? AND checkout_run_id = ?") .run(ts, run.taskId, input.runId); - if (typeof input.costUsd === "number" && Number.isFinite(input.costUsd)) { + if (runCostUsd > 0) { getDb() .prepare("update tasks set cost_usd = cost_usd + ?, updated_at = ? where id = ?") - .run(input.costUsd, ts, run.taskId); + .run(runCostUsd, ts, run.taskId); getDb() .prepare("update agents set spent_monthly_usd = spent_monthly_usd + ?, updated_at = ? where id = ?") - .run(input.costUsd, ts, run.agentId); - getDb() - .prepare( - "update agents set status = case when budget_monthly_usd > 0 and (spent_monthly_usd + ?) >= budget_monthly_usd then 'paused' else status end, updated_at = ? where id = ?", - ) - .run(input.costUsd, ts, run.agentId); + .run(runCostUsd, ts, run.agentId); } - const taskProjectRow = getDb().prepare("SELECT project_id FROM tasks WHERE id = ?").get(run.taskId) as { project_id: string | null } | undefined; - const projectId = taskProjectRow?.project_id ?? null; - getDb() - .prepare("insert into cost_entries (id, company_id, run_id, task_id, agent_id, connector_id, amount_usd, unattributed, project_id, created_at) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)") - .run(randomUUID(), run.companyId, input.runId, run.taskId, run.agentId, run.connectorId, input.costUsd, input.unattributedCost ? 1 : 0, projectId, ts); if (budgetHardStopped) { getDb() .prepare("update agents set status = 'paused', updated_at = ? where id = ?") @@ -264,6 +260,39 @@ export function buildRunMethods( return Boolean(row?.present); }, + /** + * Single writer for `cost_entries`. CostWriter calls this per + * `message_end`; tests also call it to simulate per-turn cost. + * finishRun reads SUM(cost_entries) for the run-level aggregate, + * so this is the authoritative path for cost attribution. + */ + recordCostEntry(input: { + runId: string; + companyId: string; + agentId: string; + taskId: string | null; + connectorId: string; + amountUsd: number | null; + unattributed: boolean; + }): string { + const id = randomUUID(); + const ts = nowIso(); + const projectRow = input.taskId + ? getDb().prepare("SELECT project_id FROM tasks WHERE id = ?").get(input.taskId) as { project_id: string | null } | undefined + : undefined; + const projectId = projectRow?.project_id ?? null; + getDb() + .prepare( + "insert into cost_entries (id, company_id, run_id, task_id, agent_id, connector_id, amount_usd, unattributed, project_id, created_at) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + ) + .run( + id, input.companyId, input.runId, input.taskId, input.agentId, + input.connectorId, input.amountUsd, + input.unattributed ? 1 : 0, projectId, ts, + ); + return id; + }, + listQueuedRuns() { return readAll( `select id, company_id as companyId, task_id as taskId, agent_id as agentId, workspace_id as workspaceId, diff --git a/src/main/orchestration/run-dispatch.ts b/src/main/orchestration/run-dispatch.ts index 9a006094..72d32b72 100644 --- a/src/main/orchestration/run-dispatch.ts +++ b/src/main/orchestration/run-dispatch.ts @@ -85,7 +85,6 @@ export function createRunDispatch(deps: OrchestrationDependencies) { signal: null, model: null, sessionDisplayId: null, - costUsd: null, unattributedCost: true, }); return; diff --git a/src/main/orchestration/task-runner.ts b/src/main/orchestration/task-runner.ts index cf2eb14f..d8596a67 100644 --- a/src/main/orchestration/task-runner.ts +++ b/src/main/orchestration/task-runner.ts @@ -182,7 +182,6 @@ export function createTaskRunner( signal: null, model: null, sessionDisplayId: null, - costUsd: null, unattributedCost: true, }); db.addActivity({ diff --git a/src/main/worker-service.ts b/src/main/worker-service.ts index f798df42..d15155d9 100644 --- a/src/main/worker-service.ts +++ b/src/main/worker-service.ts @@ -177,7 +177,6 @@ export function createWorkerService({ signal: null, model: null, sessionDisplayId: null, - costUsd: null, unattributedCost: true, }); if (run.workspaceId) { diff --git a/src/main/worker/handlers/run-finished-helpers.ts b/src/main/worker/handlers/run-finished-helpers.ts index 71e70bd2..b862659b 100644 --- a/src/main/worker/handlers/run-finished-helpers.ts +++ b/src/main/worker/handlers/run-finished-helpers.ts @@ -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); } diff --git a/src/main/worker/worker-crash-recovery.ts b/src/main/worker/worker-crash-recovery.ts index ebed5018..6959d895 100644 --- a/src/main/worker/worker-crash-recovery.ts +++ b/src/main/worker/worker-crash-recovery.ts @@ -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) { diff --git a/src/main/worker/worker-entry.ts b/src/main/worker/worker-entry.ts index 7191f2a1..d234856b 100644 --- a/src/main/worker/worker-entry.ts +++ b/src/main/worker/worker-entry.ts @@ -88,7 +88,6 @@ function cancelRun(runId: string) { model: null, sessionParams: null, sessionDisplayId: null, - costUsd: null, unattributedCost: true, inputTokens: null, outputTokens: null, diff --git a/src/main/worker/worker-helpers.ts b/src/main/worker/worker-helpers.ts index 359ac41e..26ceb292 100644 --- a/src/main/worker/worker-helpers.ts +++ b/src/main/worker/worker-helpers.ts @@ -43,7 +43,6 @@ export type ChildMessage = model: string | null; sessionParams: Record | null; sessionDisplayId: string | null; - costUsd: number | null; unattributedCost: boolean; inputTokens: number | null; outputTokens: number | null; diff --git a/src/main/worker/worker-runner.ts b/src/main/worker/worker-runner.ts index fa165429..b71ed5b8 100644 --- a/src/main/worker/worker-runner.ts +++ b/src/main/worker/worker-runner.ts @@ -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, }); diff --git a/src/main/worker/worker-types.ts b/src/main/worker/worker-types.ts index b2311d96..334bc4d2 100644 --- a/src/main/worker/worker-types.ts +++ b/src/main/worker/worker-types.ts @@ -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(), diff --git a/tests/integration/agent-model-switching.spec.ts b/tests/integration/agent-model-switching.spec.ts index 7f8bc5ee..e5ed7ba6 100644 --- a/tests/integration/agent-model-switching.spec.ts +++ b/tests/integration/agent-model-switching.spec.ts @@ -135,7 +135,6 @@ describe("agent model switching", () => { signal: null, model: "gpt-5.4", sessionDisplayId: null, - costUsd: null, unattributedCost: true, }); diff --git a/tests/integration/autonomous-chain/crash-recovery.spec.ts b/tests/integration/autonomous-chain/crash-recovery.spec.ts index 2b702306..bafb1f92 100644 --- a/tests/integration/autonomous-chain/crash-recovery.spec.ts +++ b/tests/integration/autonomous-chain/crash-recovery.spec.ts @@ -81,7 +81,6 @@ describe("Step 7: Worker crash recovery (database layer)", () => { signal: null, model: null, sessionDisplayId: null, - costUsd: null, unattributedCost: true, }); diff --git a/tests/integration/autonomous-chain/run-failure.spec.ts b/tests/integration/autonomous-chain/run-failure.spec.ts index 7479943b..bd0a9bee 100644 --- a/tests/integration/autonomous-chain/run-failure.spec.ts +++ b/tests/integration/autonomous-chain/run-failure.spec.ts @@ -41,7 +41,6 @@ describe("Step 6: Run failure escalation (database layer)", () => { signal: null, model: "gpt-5-codex", sessionDisplayId: null, - costUsd: 0.5, unattributedCost: false, }); diff --git a/tests/integration/cross-subsystem.spec.ts b/tests/integration/cross-subsystem.spec.ts index 79f4c97a..1d486547 100644 --- a/tests/integration/cross-subsystem.spec.ts +++ b/tests/integration/cross-subsystem.spec.ts @@ -210,6 +210,7 @@ describe("Cross-Subsystem Pipeline (D-04)", () => { // 4. Execute run db.markRunStatus(runId, "running", "Starting"); + db.recordCostEntry({ runId, companyId: s.companyId, agentId: s.agentId, taskId, connectorId: "gemini-cli", amountUsd: 0.05, unattributed: false }); db.finishRun({ runId, status: "succeeded", @@ -219,7 +220,6 @@ describe("Cross-Subsystem Pipeline (D-04)", () => { signal: null, model: "gemini-2.5-pro", sessionDisplayId: null, - costUsd: 0.05, unattributedCost: false, }); @@ -339,6 +339,7 @@ describe("Cross-Subsystem Pipeline (D-04)", () => { db.checkoutTask(taskId, runId); db.markRunStatus(runId, "running", "Starting"); + db.recordCostEntry({ runId, companyId: s.companyId, agentId: s.agentId, taskId, connectorId: s.connectorId, amountUsd: 1.5, unattributed: false }); db.finishRun({ runId, status: "succeeded", @@ -348,7 +349,6 @@ describe("Cross-Subsystem Pipeline (D-04)", () => { signal: null, model: "gemini-2.5-pro", sessionDisplayId: null, - costUsd: 1.5, // $1.50 exceeds $1 budget unattributedCost: false, }); @@ -452,6 +452,7 @@ describe("Cross-Subsystem Pipeline (D-04)", () => { }); db.checkoutTask(taskId1, runId1); db.markRunStatus(runId1, "running", "Starting"); + db.recordCostEntry({ runId: runId1, companyId: s.companyId, agentId: s.agentId, taskId: taskId1, connectorId: s.connectorId, amountUsd: 3.0, unattributed: false }); db.finishRun({ runId: runId1, status: "succeeded", @@ -461,7 +462,6 @@ describe("Cross-Subsystem Pipeline (D-04)", () => { signal: null, model: "gemini-2.5-pro", sessionDisplayId: null, - costUsd: 3.0, unattributedCost: false, }); @@ -494,6 +494,7 @@ describe("Cross-Subsystem Pipeline (D-04)", () => { }); db.checkoutTask(taskId2, runId2); db.markRunStatus(runId2, "running", "Starting"); + db.recordCostEntry({ runId: runId2, companyId: s.companyId, agentId: s.agentId, taskId: taskId2, connectorId: s.connectorId, amountUsd: 3.0, unattributed: false }); db.finishRun({ runId: runId2, status: "succeeded", @@ -503,7 +504,6 @@ describe("Cross-Subsystem Pipeline (D-04)", () => { signal: null, model: "gemini-2.5-pro", sessionDisplayId: null, - costUsd: 3.0, unattributedCost: false, }); diff --git a/tests/integration/database.spec.ts b/tests/integration/database.spec.ts index 34ac0b92..6dc81921 100644 --- a/tests/integration/database.spec.ts +++ b/tests/integration/database.spec.ts @@ -135,7 +135,6 @@ describe("AppDatabase", () => { signal: null, model: "gpt-5-codex", sessionDisplayId: "session-1", - costUsd: 12.5, unattributedCost: false, }); diff --git a/tests/integration/document-run-provenance.spec.ts b/tests/integration/document-run-provenance.spec.ts index 2f4208a5..eb7b3ace 100644 --- a/tests/integration/document-run-provenance.spec.ts +++ b/tests/integration/document-run-provenance.spec.ts @@ -149,7 +149,6 @@ describe("document run provenance", () => { status: "succeeded", summary: "Updated the market report", errorMessage: null, - costUsd: null, }, () => false); const task = harness.db.getTask(harness.taskId); diff --git a/tests/unit/atomic-checkout.test.ts b/tests/unit/atomic-checkout.test.ts index 62fffd1d..3335ea02 100644 --- a/tests/unit/atomic-checkout.test.ts +++ b/tests/unit/atomic-checkout.test.ts @@ -198,7 +198,6 @@ describe("CAS checkout", () => { signal: null, model: null, sessionDisplayId: null, - costUsd: 0, unattributedCost: false, }); diff --git a/tests/unit/database.test.ts b/tests/unit/database.test.ts index 3cc6f9ab..cdee6806 100644 --- a/tests/unit/database.test.ts +++ b/tests/unit/database.test.ts @@ -279,6 +279,17 @@ describe("AppDatabase", () => { connectorId: "openai-codex", }); + // Simulate a per-message cost entry (written by CostWriter in prod). + db.recordCostEntry({ + runId: run.id, + companyId, + agentId, + taskId, + connectorId: "openai-codex", + amountUsd: 12, + unattributed: false, + }); + const outcome = db.finishRun({ runId: run.id, status: "succeeded", @@ -288,7 +299,6 @@ describe("AppDatabase", () => { signal: null, model: "gpt-5", sessionDisplayId: "session-budget", - costUsd: 12, unattributedCost: false, }); @@ -508,7 +518,6 @@ describe("AppDatabase", () => { signal: null, model: null, sessionDisplayId: null, - costUsd: 0, unattributedCost: true, }); diff --git a/tests/unit/edge-case-budget.test.ts b/tests/unit/edge-case-budget.test.ts index 4153aa5d..fe2e3e4e 100644 --- a/tests/unit/edge-case-budget.test.ts +++ b/tests/unit/edge-case-budget.test.ts @@ -159,6 +159,7 @@ describe("Edge Case: Budget Exhaustion Mid-Run (D-05.2)", () => { db.checkoutTask(taskId2, run0.id); db.markRunStatus(run0.id, "running", ""); // Complete Run-0 with $6 cost (exceeds $5 budget) + db.recordCostEntry({ runId: run0.id, companyId, agentId, taskId: taskId2, connectorId: "gemini-cli", amountUsd: 6.0, unattributed: false }); db.finishRun({ runId: run0.id, status: "succeeded", @@ -168,12 +169,12 @@ describe("Edge Case: Budget Exhaustion Mid-Run (D-05.2)", () => { signal: null, model: "gemini-pro", sessionDisplayId: null, - costUsd: 6.0, unattributedCost: false, }); // Now complete Run-1 (the in-flight run) -- it should succeed normally // Per D-04/D-05: in-flight runs are unaffected by budget exhaustion + db.recordCostEntry({ runId: run1.id, companyId, agentId, taskId, connectorId: "gemini-cli", amountUsd: 0.5, unattributed: false }); db.finishRun({ runId: run1.id, status: "succeeded", @@ -183,7 +184,6 @@ describe("Edge Case: Budget Exhaustion Mid-Run (D-05.2)", () => { signal: null, model: "gemini-pro", sessionDisplayId: null, - costUsd: 0.50, unattributedCost: false, }); @@ -226,6 +226,7 @@ describe("Edge Case: Budget Exhaustion Mid-Run (D-05.2)", () => { }); db.checkoutTask(taskId, run.id); db.markRunStatus(run.id, "running", ""); + db.recordCostEntry({ runId: run.id, companyId, agentId, taskId, connectorId: "gemini-cli", amountUsd: 6.0, unattributed: false }); db.finishRun({ runId: run.id, status: "succeeded", @@ -235,7 +236,6 @@ describe("Edge Case: Budget Exhaustion Mid-Run (D-05.2)", () => { signal: null, model: "gemini-pro", sessionDisplayId: null, - costUsd: 6.0, unattributedCost: false, }); @@ -275,6 +275,7 @@ describe("Edge Case: Budget Exhaustion Mid-Run (D-05.2)", () => { }); db.checkoutTask(taskId, run.id); db.markRunStatus(run.id, "running", ""); + db.recordCostEntry({ runId: run.id, companyId, agentId, taskId, connectorId: "gemini-cli", amountUsd: 6.0, unattributed: false }); db.finishRun({ runId: run.id, status: "succeeded", @@ -284,7 +285,6 @@ describe("Edge Case: Budget Exhaustion Mid-Run (D-05.2)", () => { signal: null, model: "gemini-pro", sessionDisplayId: null, - costUsd: 6.0, unattributedCost: false, }); @@ -333,6 +333,7 @@ describe("Edge Case: Budget Exhaustion Mid-Run (D-05.2)", () => { }); db.checkoutTask(taskId, run.id); db.markRunStatus(run.id, "running", ""); + db.recordCostEntry({ runId: run.id, companyId, agentId, taskId, connectorId: "gemini-cli", amountUsd: 1.5, unattributed: false }); db.finishRun({ runId: run.id, status: "succeeded", @@ -342,7 +343,6 @@ describe("Edge Case: Budget Exhaustion Mid-Run (D-05.2)", () => { signal: null, model: "gemini-pro", sessionDisplayId: null, - costUsd: 1.5, unattributedCost: false, }); @@ -386,6 +386,7 @@ describe("Edge Case: Budget Exhaustion Mid-Run (D-05.2)", () => { }); db.checkoutTask(taskId, run.id); db.markRunStatus(run.id, "running", ""); + db.recordCostEntry({ runId: run.id, companyId, agentId, taskId, connectorId: "gemini-cli", amountUsd: 4.0, unattributed: false }); db.finishRun({ runId: run.id, status: "succeeded", @@ -395,7 +396,6 @@ describe("Edge Case: Budget Exhaustion Mid-Run (D-05.2)", () => { signal: null, model: "gemini-pro", sessionDisplayId: null, - costUsd: 4.0, unattributedCost: false, }); @@ -435,6 +435,7 @@ describe("Edge Case: Budget Exhaustion Mid-Run (D-05.2)", () => { db.markRunStatus(run.id, "running", ""); // Complete with cost that exceeds budget -- cost entry is still recorded + db.recordCostEntry({ runId: run.id, companyId, agentId, taskId, connectorId: "gemini-cli", amountUsd: 5.0, unattributed: false }); db.finishRun({ runId: run.id, status: "succeeded", @@ -444,7 +445,6 @@ describe("Edge Case: Budget Exhaustion Mid-Run (D-05.2)", () => { signal: null, model: "gemini-pro", sessionDisplayId: null, - costUsd: 5.0, unattributedCost: false, }); diff --git a/tests/unit/edge-case-session.test.ts b/tests/unit/edge-case-session.test.ts index 4560c7f1..750f345e 100644 --- a/tests/unit/edge-case-session.test.ts +++ b/tests/unit/edge-case-session.test.ts @@ -202,7 +202,6 @@ describe("Edge Case: Session Rotation During Execution (D-05.3)", () => { signal: null, model: "gemini-pro", sessionDisplayId: null, - costUsd: 0.05, unattributedCost: false, }); @@ -348,7 +347,6 @@ describe("Edge Case: Session Rotation During Execution (D-05.3)", () => { signal: null, model: "gemini-pro", sessionDisplayId: null, - costUsd: 0.01, unattributedCost: false, }); diff --git a/tests/unit/session-management.test.ts b/tests/unit/session-management.test.ts index 2b661c59..7bda47d3 100644 --- a/tests/unit/session-management.test.ts +++ b/tests/unit/session-management.test.ts @@ -212,7 +212,6 @@ describe("task_sessions CRUD", () => { signal: null, model: "claude-4-opus", sessionDisplayId: null, - costUsd: 0.01, unattributedCost: false, }); @@ -226,7 +225,6 @@ describe("task_sessions CRUD", () => { signal: null, model: "claude-4-opus", sessionDisplayId: null, - costUsd: 0.02, unattributedCost: false, }); @@ -571,7 +569,6 @@ describe("session management service", () => { signal: null, model: "gemini-pro", sessionDisplayId: null, - costUsd: 0.05, unattributedCost: false, }); diff --git a/tests/unit/store.test.ts b/tests/unit/store.test.ts index 187c2048..647d3e84 100644 --- a/tests/unit/store.test.ts +++ b/tests/unit/store.test.ts @@ -270,7 +270,6 @@ const fakeTask = (id: string, title = "Task"): TaskRecord => ({ activeRunId: null, sessionJson: null, sessionDisplayId: null, - costUsd: 0, requiresUserReview: false, isHeartbeat: false, startedAt: null, @@ -294,7 +293,6 @@ const fakeRun = (id: string): RunRecord => ({ signal: null, model: null, sessionDisplayId: null, - costUsd: 0, unattributedCost: false, startedAt: null, finishedAt: null,