ZacharyZhang-NY/AgentCompany
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 2 Wiki Activity

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

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/unit/atomic-checkout.test.ts
-1
View File
@@ -198,7 +198,6 @@ describe("CAS checkout", () => {
signal: null,
model: null,
sessionDisplayId: null,
costUsd: 0,
unattributedCost: false,
});
tests/unit/database.test.ts
+11 -2
View File
@@ -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,
});
tests/unit/edge-case-budget.test.ts
+7 -7
View File
@@ -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,
});
tests/unit/edge-case-session.test.ts
-2
View File
@@ -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,
});
tests/unit/session-management.test.ts
-3
View File
@@ -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,
});
tests/unit/store.test.ts
-2
View File
@@ -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,