diff --git a/doc/COST-ACCOUNTING.md b/doc/COST-ACCOUNTING.md index 0f7275421a6..815cb07f830 100644 --- a/doc/COST-ACCOUNTING.md +++ b/doc/COST-ACCOUNTING.md @@ -108,9 +108,34 @@ A correction requires `idempotencyKey`, exact `expectedCents`, `correctedCents`, ## Scale and fault qualification -`pnpm benchmark:accounting` creates its own disposable PostgreSQL database, seeds one million events across twelve months, measures report latency and eight concurrent writers on one company, exercises two active budget policies, recovers 250 pending receipts, records an analyzed project-budget query plan, and runs the independent integrity checker. `PAPERCLIP_ACCOUNTING_BENCH_ROWS` optionally sets 1,000–10,000,000 events. Results go to `coverage/accounting/scale.json`. There are no machine-dependent latency assertions in CI. +`pnpm benchmark:accounting` creates its own disposable PostgreSQL database **and receipt-spool home**. It never uses the caller's database or pending receipts. `PAPERCLIP_ACCOUNTING_BENCH_ROWS` selects 1,000–10,000,000 events (one million by default), spread across twelve UTC months. Results go to `coverage/accounting/scale.json`; copy that file before running coverage, which cleans the same output directory. There are no machine-dependent latency assertions in CI. -The September 28 local run on macOS arm64 / Node 25.6.1 measured all-time summary p95 39 ms and project grouping p95 64 ms. With eight concurrent writers, throughput was about 325 writes/s without active policies and 22 writes/s with two active policies; p95 was 72 ms and 359 ms respectively. The budgeted 250-receipt backlog drained in three bounded batches in about 12.7 seconds, with zero integrity discrepancies. Active policies deliberately read the authoritative ledger, which is more expensive than updating a projection. These synthetic warm local results are evidence, not production service-level guarantees. +The benchmark measures all-time reports, eight concurrent writers on one company, two active budgets, admission, eight concurrent ledger/budget/health report bundles competing with writers, durable checkpoints, 100,000 historical run rows, 102 policies, 1,000/10,000-file receipt backlogs, and recovery of 250 completed runs. An independent integrity check must find no discrepancies. Report bundles measure service/database work, not browser rendering, HTTP latency, provider quota calls, or financial-event queries. The fixture concentrates spend in one agent; it is not a production traffic model. + +The October 4 local comparison on macOS arm64, 10 CPUs, Node 25.6.1 measured the following at one million ledger events. PostgreSQL, API code, and storage shared one developer machine. These are observations, not service-level guarantees. + +| Measurement | Before | After | +| --- | ---: | ---: | +| Budgeted writes/second, eight writers | 20.1 | 35.1 | +| Budgeted write p95 | 396 ms | 227 ms | +| All-time summary p95 | 44 ms | 48 ms | +| All-time project grouping p95 | 71 ms | 75 ms | + +The expanded run measured report-bundle p95 of 678 ms with eight viewers and eight writers, write p95 of 569 ms under that read load, admission p95 of 103 ms, and durable checkpoint p95 of 11 ms. A 102-policy overview had p95 of 73 ms. Health over 100,000 historical runs had p95 of 19 ms. Scanning 10,000 pending receipt files added about 580 ms to recorder startup; receipts are retained until safely persisted. The API process's event-loop delay p95 was 12 ms. The 250-run recovery finished in three batches in 11.5 seconds, with zero integrity findings. + +At ten million events, budgeted throughput rose from 2.8 to 3.8 writes/s and p95 fell from 2.80 to 2.09 seconds. Recovery of 250 runs fell from 90.4 to 71.8 seconds, with no integrity discrepancies. However, eight concurrent report bundles plus eight writers reached 12.2-second report p95 and 12.4-second write p95. Warm all-time summary p95 was 649 ms; the first all-time summary took 59.2 seconds (44.5 seconds before the change). Seeding took 7.2 minutes. These larger runs shared developer hardware with test activity; report code was unchanged, so their variability is not evidence of a report-query regression. They do show that this implementation is **not qualified for heavy traffic on a ten-million-row company**. The aggregation and storage capacity limit remains; the recovery guard prevents it from also creating overlapping sweeps. + +Performance protections: + +- Cost writes aggregate the affected company, agent, and project policies in one ledger scan and one pending-run scan. Each policy retains its own scope and UTC/lifetime window. Exact ledger values remain authoritative, including unknown pricing and native recovery guards. +- A budget operation reuses its observation for thresholds, incident amounts, and reasons. The observation also carries its UTC window across a midnight boundary. It does not cache observations across ledger mutations. PostgreSQL filters policies to the affected scopes before returning rows. +- Overview policy and incident reads run in batches of four, so a large policy list does not enqueue one query per policy at once. +- Health uses a read-only, repeatable-read snapshot. Its counts and details agree without taking the company accounting write lock. Inspection and repair retain their existing locking requirements. +- One heartbeat service shares a running receipt replay/recovery/policy sweep across overlapping scheduler ticks. Success or failure releases the guard; a later tick can retry. This is a per-process scheduler protection, not a distributed lease. Ledger locks and receipt idempotency still protect multiple processes. + +The regression suite checks SQL scan counts, mixed-scope/window totals, pending/unpriced/native recovery guards, snapshot consistency while a writer commits, and coalesced recovery through both success and failure. It also retains the concurrency, crash recovery, exact-money, and mutation tests. Wall-clock benchmarks stay separate from CI correctness gates. + +Capacity remains finite. Budget decisions still sum authoritative ledger rows, so a very large company or lifetime project costs more per write. First writes after a UTC rollover rebuild monthly projections. Cold all-time reports can be much slower than warm reports. A large receipt backlog still requires directory scanning. Qualify the expected company write rate, storage latency, database pool size, and viewer concurrency on deployment hardware before increasing traffic; monitor receipt age and recovery duration as well as request latency. Sustained traffic near measured saturation needs further database aggregation work with equivalent repair and consistency guarantees. Monthly projections carry an explicit UTC month marker. The first write after upgrade or rollover initializes from the ledger; subsequent same-month writes use exact atomic increments. Backdated events do not increment the current month. Reports and budget decisions continue using the ledger. Composite project/date and partial unpriced indexes support targeted queries; pending recovery is indexed in its actual update-time order. diff --git a/server/scripts/benchmark-accounting.ts b/server/scripts/benchmark-accounting.ts index a18a0ec8776..dc3c9e5d1f3 100644 --- a/server/scripts/benchmark-accounting.ts +++ b/server/scripts/benchmark-accounting.ts @@ -1,21 +1,34 @@ +import { randomUUID } from "node:crypto"; +import { reserveRunBudget } from "../src/services/budget-reservations.js"; +import { createRunUsageRecorder, indexPendingUsageReceipts, usageReceiptSpoolPath } from "../src/services/usage-receipts.js"; import { budgetService } from "../src/services/budgets.js"; import { sql } from "drizzle-orm"; /** Reproducible synthetic PostgreSQL benchmark. Never reads DATABASE_URL or * provider credentials; startEmbeddedPostgresTestDatabase owns all data. */ -import { performance } from "node:perf_hooks"; -import { mkdir, writeFile } from "node:fs/promises"; +import { performance, monitorEventLoopDelay } from "node:perf_hooks"; +import { mkdir, writeFile, mkdtemp, rm } from "node:fs/promises"; import os from "node:os"; import path from "node:path"; -import { agents, companies, projects, createDb, startEmbeddedPostgresTestDatabase } from "@paperclipai/db"; +import { agents, companies, projects, heartbeatRuns, budgetPolicies, createDb, startEmbeddedPostgresTestDatabase } from "@paperclipai/db"; import { costService } from "../src/services/costs.js"; import { withAccountingTransaction } from "../src/services/accounting-transaction.js"; import { accountingIntegrityService } from "../src/services/accounting-integrity.js"; import { reconcileRunCosts } from "../src/services/run-cost-accounting.js"; const count = Number(process.env.PAPERCLIP_ACCOUNTING_BENCH_ROWS ?? 1_000_000); if (!Number.isInteger(count) || count < 1000 || count > 10_000_000) throw new Error("Benchmark row count must be 1,000–10,000,000"); -const database = await startEmbeddedPostgresTestDatabase("accounting-scale-"); +// Recovery reads the instance receipt spool as well as the database. Always +// own both; a caller's real PAPERCLIP_HOME must never be scanned or replayed. +const previousHome = process.env.PAPERCLIP_HOME; +const benchmarkHome = await mkdtemp(path.join(os.tmpdir(), "accounting-benchmark-")); +process.env.PAPERCLIP_HOME = benchmarkHome; +const database = await startEmbeddedPostgresTestDatabase("accounting-scale-").catch(async error => { + if (previousHome === undefined) delete process.env.PAPERCLIP_HOME; else process.env.PAPERCLIP_HOME = previousHome; + await rm(benchmarkHome, { recursive: true, force: true }); + throw error; +}); const db = createDb(database.connectionString); const started = performance.now(); +const eventLoop = monitorEventLoopDelay({ resolution: 10 }); eventLoop.enable(); function distribution(values: number[]) { const sorted = [...values].sort((a,b) => a-b); const pick = (p: number) => Number(sorted[Math.min(sorted.length - 1, Math.ceil(p * sorted.length) - 1)].toFixed(2)); @@ -54,6 +67,34 @@ try { agentId: agent.id, provider: "fixture", model: "fixture", costCents: "0.0000001", idempotencyKey: `budgeted:${batch}:${index}`, occurredAt: new Date(), })).then(ms => budgetedWrites.push(ms)))); const budgetedMs = performance.now() - budgetedStart; + const admissions: number[] = []; + for (let i = 0; i < 20; i++) { + const [run] = await db.insert(heartbeatRuns).values({ companyId: company.id, agentId: agent.id, status: "running" }).returning(); + admissions.push(await time(() => reserveRunBudget(db, company.id, run.id, null))); + } + const dashboard: number[] = [], writesWhilePolling: number[] = []; + // Eight viewers each issue the ledger queries used by the Costs overview, + // together with budget and accounting health reads, while eight writes run. + for (let batch = 0; batch < 4; batch++) await Promise.all([ + ...Array.from({ length: 8 }, () => time(() => Promise.all([ + costService(db).summary(company.id), costService(db).byAgent(company.id), + costService(db).byProject(company.id), costService(db).byAgentModel(company.id), + budgetService(db).overview(company.id), accountingIntegrityService(db).health(company.id), + ])).then(ms => dashboard.push(ms))), + ...Array.from({ length: concurrency }, (_, index) => time(() => costService(db).createEvent(company.id, { + agentId: agent.id, provider: "fixture", model: "fixture", costCents: "0.0000001", idempotencyKey: `polling:${batch}:${index}`, occurredAt: new Date(), + })).then(ms => writesWhilePolling.push(ms))), + ]); + const [checkpointRun] = await db.insert(heartbeatRuns).values({ companyId: company.id, agentId: agent.id, status: "running" }).returning(); + const recorder = await createRunUsageRecorder(db, { companyId: company.id, runId: checkpointRun.id, adapterType: "process" }); + const checkpoints: number[] = []; + for (let i = 1; i <= 50; i++) checkpoints.push(await time(() => recorder.capture({ complete: false, usageBasis: "per_run", usage: { inputTokens: i, outputTokens: i }, costUsdExact: "0.000000001" }))); + // Include historical run-table size, independently of the pending index. + await db.$client`insert into heartbeat_runs (company_id,agent_id,status,cost_accounting_pending,finished_at) + select ${company.id},${agent.id},'succeeded',false, now() from generate_series(1,100000)`; + await db.$client`analyze heartbeat_runs`; + const health: number[] = []; + for (let i = 0; i < 20; i++) health.push(await time(() => accountingIntegrityService(db).health(company.id))); const waits: number[] = []; await Promise.all(Array.from({ length: concurrency }, () => { const start = performance.now(); @@ -66,6 +107,28 @@ try { const before = performance.now(); const result = await reconcileRunCosts(db); recovery.push({ milliseconds: performance.now()-before, ...result }); } if (recovery.reduce((sum, r) => sum+r.accounted,0) !== 250) throw new Error("Benchmark recovery did not conserve receipt count"); + const spool = usageReceiptSpoolPath(); await mkdir(spool, { recursive: true }); + const spoolBacklogs: Array<{ files: number; indexMs: number; recorderStartMs: number }> = []; + const foreignRunId = randomUUID(); + let published = 0; + for (const files of [1000, 10000]) { + // Publish synthetic immutable files directly; timed checkpoints above use + // the real fsync path. These belong to another run and must be retained. + for (; published < files; published += 50) await Promise.all(Array.from({ length: 50 }, (_, offset) => writeFile(path.join(spool, `${published + offset}.json`), JSON.stringify({ + schema: "paperclip/accounting-receipt/v1", id: randomUUID(), companyId: company.id, runId: foreignRunId, + sourceId: randomUUID(), sequence: 1, receivedAt: new Date().toISOString(), adapterType: "process", receipt: { complete: false }, + })))); + const indexMs = await time(() => indexPendingUsageReceipts()); + const [run] = await db.insert(heartbeatRuns).values({ companyId: company.id, agentId: agent.id, status: "running" }).returning(); + const recorderStartMs = await time(() => createRunUsageRecorder(db, { companyId: company.id, runId: run.id, adapterType: "process" })); + spoolBacklogs.push({ files, indexMs, recorderStartMs }); + } + await rm(spool, { recursive: true, force: true }); + // A large policy list must not make each write fetch every agent's policies. + const extraAgents = await db.insert(agents).values(Array.from({ length: 100 }, (_, i) => ({ companyId: company.id, name: `Idle ${i}`, role: "engineer", adapterType: "process" }))).returning(); + await db.insert(budgetPolicies).values(extraAgents.map(extra => ({ companyId: company.id, scopeType: "agent", scopeId: extra.id, windowKind: "calendar_month_utc", amount: 2_000_000_000 }))); + const largeOverview: number[] = []; + for (let i = 0; i < 10; i++) largeOverview.push(await time(() => budgetService(db).overview(company.id))); const integrity = await accountingIntegrityService(db).inspect(company.id); if (integrity.findings.length) throw new Error(`Benchmark violated accounting integrity: ${JSON.stringify(integrity.findings)}`); const plan = await db.$client`explain (analyze,buffers,format json) select sum(cost_cents) from cost_events where company_id = ${company.id} and project_id = ${project.id} @@ -74,7 +137,18 @@ try { seedRows: count, seedMs, concurrency, allTimeSummary: distribution(reports), allTimeProjects: distribution(projectReports), hotCompanyWrites: { activePolicies: 0, ...distribution(writes), writesPerSecond: writes.length / (writeMs / 1000) }, companyLockAcquisition: distribution(waits), budgetedCompanyWrites: { activePolicies: 2, ...distribution(budgetedWrites), writesPerSecond: budgetedWrites.length / (budgetedMs / 1000) }, + admission: distribution(admissions), concurrentDashboard: distribution(dashboard), writesWhilePolling: distribution(writesWhilePolling), + durableCheckpoint: distribution(checkpoints), healthWith100kHistoricalRuns: distribution(health), + overviewWith102Policies: distribution(largeOverview), spoolBacklogs, + eventLoopDelay: { p95Ms: eventLoop.percentile(95) / 1e6, maxMs: eventLoop.max / 1e6 }, recovery, projectBudgetQueryPlan: plan, integrityFindings: integrity.findings.length, totalMs: performance.now()-started }; const output = path.resolve("coverage/accounting/scale.json"); await mkdir(path.dirname(output), { recursive: true }); await writeFile(output, JSON.stringify(report,null,2)); process.stdout.write(JSON.stringify({ ...report, projectBudgetQueryPlan: "see coverage/accounting/scale.json" },null,2)+"\n"); -} finally { await db.$client.end(); await database.cleanup(); } +} finally { + eventLoop.disable(); + try { await db.$client.end(); await database.cleanup(); } + finally { + if (previousHome === undefined) delete process.env.PAPERCLIP_HOME; else process.env.PAPERCLIP_HOME = previousHome; + await rm(benchmarkHome, { recursive: true, force: true }); + } +} diff --git a/server/src/__tests__/cost-accounting-performance.test.ts b/server/src/__tests__/cost-accounting-performance.test.ts new file mode 100644 index 00000000000..68821f54066 --- /dev/null +++ b/server/src/__tests__/cost-accounting-performance.test.ts @@ -0,0 +1,205 @@ +import { randomUUID } from "node:crypto"; +import { promises as fs } from "node:fs"; +import { afterAll, afterEach, beforeAll, describe, expect, it, vi } from "vitest"; +import { eq, sql } from "drizzle-orm"; +import { drizzle } from "drizzle-orm/postgres-js"; +import { agents, budgetPolicies, companies, costEvents, createDb, heartbeatRuns, issues, nativeRunFinalizations, projects, type Db } from "@paperclipai/db"; +import { accountingIntegrityService } from "../services/accounting-integrity.js"; +import { budgetService, budgetServiceInTransaction } from "../services/budgets.js"; +import { withAccountingTransaction } from "../services/accounting-transaction.js"; +import { costService } from "../services/costs.js"; +import { createCostAccountingReconciler } from "../services/run-cost-accounting.js"; +import { startEmbeddedPostgresTestDatabase } from "./helpers/embedded-postgres.js"; + +function deferred() { + let resolve!: () => void; + const promise = new Promise(done => { resolve = done; }); + return { promise, resolve }; +} + +describe("accounting performance invariants (PostgreSQL)", () => { + let database: Awaited>, db: Db; + beforeAll(async () => { database = await startEmbeddedPostgresTestDatabase("accounting-performance-"); db = createDb(database.connectionString); }, 30_000); + afterEach(() => { vi.restoreAllMocks(); vi.useRealTimers(); }); + afterAll(async () => { await database?.cleanup(); }); + async function fixture() { + const [company] = await db.insert(companies).values({ name: "Performance", issuePrefix: `P${randomUUID().slice(0, 7)}` }).returning(); + const [agent] = await db.insert(agents).values({ companyId: company.id, name: "Worker", role: "engineer", adapterType: "process" }).returning(); + const [project] = await db.insert(projects).values({ companyId: company.id, name: "Work" }).returning(); + return { company, agent, project }; + } + function trace() { + const queries: string[] = []; + const traced = drizzle(db.$client, { logger: { logQuery(query) { queries.push(query); } } }) as unknown as Db; + return { traced, ledgerReads: () => queries.filter(query => query.startsWith("select") && query.includes('from "cost_events"')) }; + } + + it("reads health while a writer holds the company lock and keeps counts/items in one snapshot", async () => { + const f = await fixture(); + const writer = await db.$client.reserve(); + await writer`begin`; + try { + await writer`select id from companies where id = ${f.company.id} for no key update`; + const original = db.transaction.bind(db); + const transaction = vi.spyOn(db, "transaction").mockImplementation((callback, config) => original(async tx => { + // A lock attempt fails promptly. There is no elapsed-time pass/fail assertion. + await tx.execute(sql`set local lock_timeout = '200ms'`); + const execute = tx.execute.bind(tx); + let counted = false; + vi.spyOn(tx, "execute").mockImplementation(query => { + const result = execute(query); + const executeRaw = result.execute.bind(result); + vi.spyOn(result, "execute").mockImplementation(async () => { + const rows = await executeRaw(); + if (!counted) { + counted = true; + // Commit a pending row between health's counts and item list. + await writer`insert into heartbeat_runs (company_id,agent_id,status,cost_accounting_pending) + values (${f.company.id},${f.agent.id},'failed',true)`; + await writer`commit`; + } + return rows; + }); + return result; + }); + return callback(tx); + }, config)); + expect(await accountingIntegrityService(db).health(f.company.id)).toMatchObject({ pendingRunCount: 0, items: [] }); + expect(transaction).toHaveBeenCalledWith(expect.any(Function), { isolationLevel: "repeatable read", accessMode: "read only" }); + transaction.mockRestore(); + expect(await accountingIntegrityService(db).health(f.company.id)).toMatchObject({ pendingRunCount: 1, items: [expect.objectContaining({ state: "retryable" })] }); + } finally { await writer`rollback`; writer.release(); } + // Remove the fixture's pending debt before the global recovery tests. + await db.execute(sql`delete from heartbeat_runs where company_id = ${f.company.id}`); + await expect(accountingIntegrityService(db).health(randomUUID())).rejects.toThrow("Company not found"); + }); + + it("shares receipt replay across overlapping ticks and releases the guard after replay failure", async () => { + const reconcile = createCostAccountingReconciler(db); + const gate = deferred(); + const read = vi.spyOn(fs, "readdir").mockImplementationOnce(async () => { await gate.promise; throw new Error("Spool unavailable"); }); + const first = reconcile(), second = reconcile(); + expect(second).toBe(first); + expect(read).toHaveBeenCalledTimes(1); + const settled = Promise.allSettled([first, second]); + gate.resolve(); + expect(await settled).toEqual([expect.objectContaining({ status: "rejected" }), expect.objectContaining({ status: "rejected" })]); + read.mockRestore(); + expect(await reconcile()).toMatchObject({ scanned: 0, accounted: 0 }); + }); + + it("keeps the guard through budget delivery and retries after reconciliation failure", async () => { + const f = await fixture(); + await budgetService(db).upsertPolicy(f.company.id, { scopeType: "company", scopeId: f.company.id, amount: 1 }, "board"); + await costService(db).createEvent(f.company.id, { agentId: f.agent.id, provider: "fixture", model: "fixture", costCents: 1, occurredAt: new Date() }); + const entered = deferred(), gate = deferred(); + const cancelWorkForScope = vi.fn(async () => { entered.resolve(); await gate.promise; }); + const reconcile = createCostAccountingReconciler(db, { cancelWorkForScope }); + const first = reconcile(); + try { + await entered.promise; + expect(reconcile()).toBe(first); + expect(cancelWorkForScope).toHaveBeenCalledTimes(1); + } finally { gate.resolve(); } + await first; + const select = vi.spyOn(db, "select").mockImplementationOnce(() => { throw new Error("Database unavailable"); }); + await expect(reconcile()).rejects.toThrow("Database unavailable"); + select.mockRestore(); + const next = reconcile(); expect(next).not.toBe(first); + await next; + }); + + it("uses one ledger observation per overview policy and per admission policy", async () => { + const f = await fixture(); + await budgetService(db).upsertPolicy(f.company.id, { scopeType: "company", scopeId: f.company.id, amount: 100 }, "board"); + const { traced, ledgerReads } = trace(); + expect((await budgetService(traced).overview(f.company.id)).policies[0]).toMatchObject({ observedAmountExact: "0.0000000", status: "ok" }); + expect(ledgerReads()).toHaveLength(1); + await budgetService(traced).getInvocationBlock(f.company.id, f.agent.id); + expect(ledgerReads()).toHaveLength(2); + }); + + it("shares a scan across mixed policy scopes/windows without trusting monthly projections", async () => { + const f = await fixture(), foreign = await fixture(); + const [other] = await db.insert(agents).values({ companyId: f.company.id, name: "Other", role: "engineer", adapterType: "process" }).returning(); + const values = [ + { scopeType: "company", scopeId: f.company.id, windowKind: "calendar_month_utc", amount: 4 }, + { scopeType: "agent", scopeId: f.agent.id, windowKind: "lifetime", amount: 10 }, + { scopeType: "project", scopeId: f.project.id, windowKind: "lifetime", amount: 9 }, + { scopeType: "agent", scopeId: other.id, windowKind: "calendar_month_utc", amount: 1 }, + ]; + await db.insert(budgetPolicies).values(values.map(value => ({ ...value, companyId: f.company.id }))); + const base = { provider: "fixture", model: "fixture", costCents: 3, occurredAt: new Date() }; + const [event] = await db.insert(costEvents).values([ + { ...base, companyId: f.company.id, agentId: f.agent.id, projectId: f.project.id, costCents: 7, occurredAt: new Date("2000-01-01") }, + { ...base, companyId: f.company.id, agentId: f.agent.id, projectId: f.project.id }, + { ...base, companyId: f.company.id, agentId: other.id }, + { ...base, companyId: foreign.company.id, agentId: foreign.agent.id, costCents: 999 }, + ]).returning(); + await db.update(companies).set({ spentMonthlyCents: 0 }).where(eq(companies.id, f.company.id)); + const { traced, ledgerReads } = trace(); + await withAccountingTransaction(traced, f.company.id, tx => budgetServiceInTransaction(tx).evaluateCostEvent(event)); + expect(ledgerReads()).toHaveLength(1); + const incidents = await db.execute<{ scope_type: string; amount_observed: string }>(sql`select scope_type,amount_observed::text from budget_incidents where company_id = ${f.company.id} order by scope_type`); + expect(incidents).toEqual([ + { scope_type: "agent", amount_observed: "10.0000000" }, + { scope_type: "company", amount_observed: "6.0000000" }, + { scope_type: "project", amount_observed: "10.0000000" }, + ]); + }); + + it("keeps unpriced and pending-run guards scoped when batching native recovery observations", async () => { + const f = await fixture(); + const [other] = await db.insert(agents).values({ companyId: f.company.id, name: "Other", role: "engineer", adapterType: "process" }).returning(); + await db.insert(budgetPolicies).values([ + { companyId: f.company.id, scopeType: "company", scopeId: f.company.id, windowKind: "calendar_month_utc", amount: 100 }, + { companyId: f.company.id, scopeType: "agent", scopeId: f.agent.id, windowKind: "lifetime", amount: 100 }, + { companyId: f.company.id, scopeType: "project", scopeId: f.project.id, windowKind: "lifetime", amount: 100 }, + ]); + const [issue] = await db.insert(issues).values({ companyId: f.company.id, title: "Recover" }).returning(); + const [native] = await db.insert(heartbeatRuns).values({ companyId: f.company.id, agentId: f.agent.id, + status: "failed", runtimeMode: "native", nativeIssueId: issue.id, costAccountingPending: true, finishedAt: new Date() }).returning(); + await db.insert(nativeRunFinalizations).values({ companyId: f.company.id, runId: native.id, issueId: issue.id, phase: "retryable_failure" }); + const base = { companyId: f.company.id, provider: "fixture", model: "fixture", costCents: 0, occurredAt: new Date() }; + const [event] = await db.insert(costEvents).values([ + { ...base, agentId: f.agent.id, projectId: f.project.id }, + { ...base, agentId: other.id, projectId: f.project.id, occurredAt: new Date("2000-01-01"), costStatus: "unpriced" }, + ]).returning(); + await withAccountingTransaction(db, f.company.id, tx => budgetServiceInTransaction(tx).evaluateCostEvent(event)); + const overview = await budgetService(db).overview(f.company.id); + expect(overview.activeIncidents.map(row => row.scopeType)).toEqual(["project"]); + expect(overview.policies).toEqual(expect.arrayContaining([ + expect.objectContaining({ scopeType: "company", status: "ok", pendingRunCount: 1, unpricedEventCount: 0 }), + expect.objectContaining({ scopeType: "agent", status: "ok", pendingRunCount: 1, unpricedEventCount: 0 }), + expect.objectContaining({ scopeType: "project", status: "hard_stop", pendingRunCount: 0, unpricedEventCount: 1 }), + ])); + await db.insert(heartbeatRuns).values({ companyId: f.company.id, agentId: f.agent.id, status: "failed", + costAccountingPending: true, finishedAt: new Date(), usageJson: { ledgerScope: { projectId: f.project.id } } }); + await withAccountingTransaction(db, f.company.id, tx => budgetServiceInTransaction(tx).evaluateCostEvent(event)); + expect((await budgetService(db).overview(f.company.id)).activeIncidents.map(row => row.scopeType).sort()).toEqual(["agent", "company", "project"]); + }); + + + it("keeps incident amounts and overview windows tied to the observed UTC month across midnight", async () => { + const f = await fixture(); + await db.insert(budgetPolicies).values({ companyId: f.company.id, scopeType: "company", scopeId: f.company.id, + windowKind: "calendar_month_utc", amount: 1 }); + const [event] = await db.insert(costEvents).values({ companyId: f.company.id, agentId: f.agent.id, + provider: "fixture", model: "fixture", costCents: 1, occurredAt: new Date("2026-01-31T12:00:00Z") }).returning(); + const traced = drizzle(db.$client, { logger: { logQuery(query) { + if (query.startsWith("select") && query.includes('from "cost_events"')) vi.setSystemTime(new Date("2026-02-01T00:00:00Z")); + } } }) as unknown as Db; + vi.useFakeTimers({ toFake: ["Date"] }); + vi.setSystemTime(new Date("2026-01-31T23:59:59.999Z")); + await withAccountingTransaction(traced, f.company.id, tx => budgetServiceInTransaction(tx).evaluateCostEvent(event)); + const incident = (await budgetService(db).overview(f.company.id)).activeIncidents[0]; + expect(incident).toMatchObject({ amountObserved: 1, windowStart: new Date("2026-01-01T00:00:00Z"), windowEnd: new Date("2026-02-01T00:00:00Z") }); + vi.setSystemTime(new Date("2026-01-31T23:59:59.999Z")); + expect((await budgetService(traced).overview(f.company.id)).policies[0]).toMatchObject({ + observedAmountExact: "1.0000000", windowStart: new Date("2026-01-01T00:00:00Z"), windowEnd: new Date("2026-02-01T00:00:00Z"), + }); + // The next operation uses February and releases the prior month's pause. + expect(await budgetService(db).getInvocationBlock(f.company.id, f.agent.id)).toBeNull(); + }); + +}); diff --git a/server/src/services/accounting-integrity.ts b/server/src/services/accounting-integrity.ts index 56c0412169b..206eb73d190 100644 --- a/server/src/services/accounting-integrity.ts +++ b/server/src/services/accounting-integrity.ts @@ -98,7 +98,11 @@ export function accountingIntegrityService(db: Db, hooks: BudgetServiceHooks = { throw error; } }, - health: (companyId: string): Promise => withAccountingTransaction(db, companyId, async tx => { + // A consistent read snapshot must not queue behind (or delay) cost writes. + // Mutation paths still use the company accounting lock. + health: (companyId: string): Promise => db.transaction(async tx => { + const [company] = await tx.select({ id: companies.id }).from(companies).where(eq(companies.id, companyId)); + if (!company) throw notFound("Company not found"); const [counts] = await tx.execute<{ pending: number; unpriced: number; oldest: string | null; cancellations: number; reserved: string }>(sql` select (select count(*)::int from heartbeat_runs where company_id = ${companyId} and cost_accounting_pending and status in ('succeeded','failed','timed_out','cancelled','interrupted')) as pending, (select count(*)::int from cost_events where company_id = ${companyId} and cost_status = 'unpriced' and billing_type <> 'subscription_included') as unpriced, @@ -116,6 +120,6 @@ export function accountingIntegrityService(db: Db, hooks: BudgetServiceHooks = { where company_id = ${companyId} and cost_status = 'unpriced' and billing_type <> 'subscription_included' order by occurred_at,id limit 100`); return { companyId, pendingRunCount: counts.pending, unpricedEventCount: counts.unpriced, oldestPendingAt: counts.oldest, pendingCancellationCount: counts.cancellations, heldReservationCents: normalizeCents(counts.reserved), items: [...pending, ...unpriced.map(row => ({ ...row, state: "unpriced" as const, lastError: null, attempts: 0, lastAttemptAt: null }))] }; - }), + }, { isolationLevel: "repeatable read", accessMode: "read only" }), }; } diff --git a/server/src/services/budget-reservations.ts b/server/src/services/budget-reservations.ts index cefb4709da8..f1788a1d8b9 100644 --- a/server/src/services/budget-reservations.ts +++ b/server/src/services/budget-reservations.ts @@ -3,7 +3,7 @@ import { budgetPolicies, budgetReservations, heartbeatRuns, nativeRunFinalizatio import { centsToUnits, unitsToCents } from "@paperclipai/shared"; import { conflict, notFound } from "../errors.js"; import { withAccountingTransaction } from "./accounting-transaction.js"; -import { budgetServiceInTransaction, computeObservedSpend } from "./budgets.js"; +import { budgetServiceInTransaction, budgetPoliciesForRun, computeObservedSpend } from "./budgets.js"; /** Reserve before dispatch, under the same company lock as charges and policy * changes. Estimates constrain admission; they cannot cap a provider's bill. @@ -28,9 +28,11 @@ export async function reserveRunBudget(db: Db, companyId: string, runId: string, } const block = await budgetServiceInTransaction(tx, publications).getInvocationBlock(companyId, run.agentId, { projectId }); if (block) throw conflict(block.reason); - const candidates = await tx.select().from(budgetPolicies).where(and(eq(budgetPolicies.companyId, companyId), eq(budgetPolicies.isActive, true), eq(budgetPolicies.hardStopEnabled, true))); - const policies = candidates.filter(p => p.metric === "billed_cents" && p.amount > 0 && ( - p.scopeType === "company" && p.scopeId === companyId || p.scopeType === "agent" && p.scopeId === run.agentId || p.scopeType === "project" && p.scopeId === projectId)); + const policies = await tx.select().from(budgetPolicies).where(and( + eq(budgetPolicies.companyId, companyId), eq(budgetPolicies.isActive, true), eq(budgetPolicies.hardStopEnabled, true), + eq(budgetPolicies.metric, "billed_cents"), sql`${budgetPolicies.amount} > 0`, + budgetPoliciesForRun(companyId, run.agentId, projectId), + )); const amount = policies.reduce((max, policy) => { const next = centsToUnits(policy.reservationCents); return next > max ? next : max; }, 0n); for (const policy of policies) { const [held] = await tx.select({ amount: sql`coalesce(sum(${budgetReservations.amountCents}), 0)::text` }).from(budgetReservations).where(and( diff --git a/server/src/services/budgets.ts b/server/src/services/budgets.ts index 54d2675e97c..6fc08582000 100644 --- a/server/src/services/budgets.ts +++ b/server/src/services/budgets.ts @@ -1,5 +1,5 @@ import { compareCents, normalizeCents } from "@paperclipai/shared"; -import { and, desc, eq, gte, inArray, lt, ne, sql, or, isNull } from "drizzle-orm"; +import { and, desc, eq, gte, inArray, lt, ne, sql, or, isNull, type SQL } from "drizzle-orm"; import type { Db } from "@paperclipai/db"; import { agents, @@ -149,55 +149,77 @@ async function resolveScopeRecord(db: Db, scopeType: BudgetScopeType, scopeId: s }; } -export async function computeObservedSpend( - db: Db, - policy: Pick, -) { - if (policy.metric !== "billed_cents") return { total: 0, totalExact: "0.0000000", unpricedEventCount: 0, pendingRunCount: 0, recoveringRunCount: 0 }; - - const conditions = [eq(costEvents.companyId, policy.companyId)]; - if (policy.scopeType === "agent") conditions.push(eq(costEvents.agentId, policy.scopeId)); - if (policy.scopeType === "project") conditions.push(eq(costEvents.projectId, policy.scopeId)); - const { start, end } = resolveWindow(policy.windowKind as BudgetWindowKind); - if (policy.windowKind === "calendar_month_utc") { - conditions.push(gte(costEvents.occurredAt, start)); - conditions.push(lt(costEvents.occurredAt, end)); +// Keep a large dashboard from filling the database pool with queued reads. +async function mapBudgetReads(rows: T[], read: (row: T) => Promise): Promise { + const result: R[] = []; + for (let start = 0; start < rows.length; start += 4) { + result.push(...await Promise.all(rows.slice(start, start + 4).map(read))); } + return result; +} - const [row] = await db - .select({ - unpricedEventCount: sql`count(*) filter (where ${costEvents.costStatus} = 'unpriced' and ${costEvents.billingType} <> 'subscription_included')::int`, - total: sql`coalesce(sum(${costEvents.costCents}), 0)::text`, - }) - .from(costEvents) - .where(and(...conditions)); +/** Filter in PostgreSQL so one run never loads unrelated agents' policies. */ +export function budgetPoliciesForRun(companyId: string, agentId: string, projectId: string | null) { + return or( + and(eq(budgetPolicies.scopeType, "company"), eq(budgetPolicies.scopeId, companyId)), + and(eq(budgetPolicies.scopeType, "agent"), eq(budgetPolicies.scopeId, agentId)), + projectId ? and(eq(budgetPolicies.scopeType, "project"), eq(budgetPolicies.scopeId, projectId)) : undefined, + ); +} - const pendingConditions = [eq(heartbeatRuns.companyId, policy.companyId), eq(heartbeatRuns.costAccountingPending, true), - inArray(heartbeatRuns.status, ["succeeded", "failed", "timed_out", "cancelled", "interrupted"])]; - if (policy.scopeType === "agent") pendingConditions.push(eq(heartbeatRuns.agentId, policy.scopeId)); - if (policy.scopeType === "project") pendingConditions.push(sql`${heartbeatRuns.usageJson}->'ledgerScope'->>'projectId' = ${policy.scopeId}`); - if (policy.windowKind === "calendar_month_utc") { - pendingConditions.push(sql`coalesce(${heartbeatRuns.finishedAt}, ${heartbeatRuns.createdAt}) >= ${start.toISOString()}::timestamptz`, sql`coalesce(${heartbeatRuns.finishedAt}, ${heartbeatRuns.createdAt}) < ${end.toISOString()}::timestamptz`); +type SpendPolicy = Pick; +type ObservedSpend = { windowStart: Date; windowEnd: Date; total: number; totalExact: string; unpricedEventCount: number; pendingRunCount: number; recoveringRunCount: number }; + +/** One ledger scan and one pending-run scan for the affected policies. The + * company lock still protects decisions; nothing is cached across mutations. + * Each aggregate keeps its own scope/window, including mixed lifetime and + * monthly policies. An outer predicate excludes rows no policy needs. */ +async function computeObservedSpends(db: Db, policies: SpendPolicy[]): Promise { + const now = new Date(); + const windows = policies.map(policy => resolveWindow(policy.windowKind as BudgetWindowKind, now)); + const billable = policies.map((policy, index) => ({ policy, index })).filter(({ policy }) => policy.metric === "billed_cents"); + const ledgerColumns: Record = {}; + const pendingColumns: Record = {}; + const ledgerPredicates: SQL[] = [], pendingPredicates: SQL[] = []; + for (const { policy, index } of billable) { + const { start, end } = windows[index]; + const ledger = and( + eq(costEvents.companyId, policy.companyId), + policy.scopeType === "agent" ? eq(costEvents.agentId, policy.scopeId) : undefined, + policy.scopeType === "project" ? eq(costEvents.projectId, policy.scopeId) : undefined, + policy.windowKind === "calendar_month_utc" ? and(gte(costEvents.occurredAt, start), lt(costEvents.occurredAt, end)) : undefined, + )!; + const pending = and( + eq(heartbeatRuns.companyId, policy.companyId), eq(heartbeatRuns.costAccountingPending, true), + inArray(heartbeatRuns.status, ["succeeded", "failed", "timed_out", "cancelled", "interrupted"]), + policy.scopeType === "agent" ? eq(heartbeatRuns.agentId, policy.scopeId) : undefined, + policy.scopeType === "project" ? sql`${heartbeatRuns.usageJson}->'ledgerScope'->>'projectId' = ${policy.scopeId}` : undefined, + policy.windowKind === "calendar_month_utc" ? sql`coalesce(${heartbeatRuns.finishedAt}, ${heartbeatRuns.createdAt}) >= ${start.toISOString()}::timestamptz + and coalesce(${heartbeatRuns.finishedAt}, ${heartbeatRuns.createdAt}) < ${end.toISOString()}::timestamptz` : undefined, + )!; + ledgerPredicates.push(ledger); pendingPredicates.push(pending); + ledgerColumns[`total${index}`] = sql`coalesce(sum(${costEvents.costCents}) filter (where ${ledger}), 0)::text`; + ledgerColumns[`unpriced${index}`] = sql`count(*) filter (where ${ledger} and ${costEvents.costStatus} = 'unpriced' and ${costEvents.billingType} <> 'subscription_included')::int`; + pendingColumns[`pending${index}`] = sql`count(*) filter (where ${pending})::int`; + pendingColumns[`recovering${index}`] = sql`count(*) filter (where ${pending} and ${heartbeatRuns.runtimeMode} = 'native' + and exists (select 1 from ${nativeRunFinalizations} where ${nativeRunFinalizations.runId} = ${heartbeatRuns.id} + and ${nativeRunFinalizations.companyId} = ${heartbeatRuns.companyId} + and ${nativeRunFinalizations.resultId} is null and ${nativeRunFinalizations.phase} <> 'terminal_failure'))::int`; } - const [pending] = await db.select({ - count: sql`count(*)::int`, - recovering: sql`count(*) filter (where ${heartbeatRuns.runtimeMode} = 'native' - and exists (select 1 from ${nativeRunFinalizations} - where ${nativeRunFinalizations.runId} = ${heartbeatRuns.id} - and ${nativeRunFinalizations.companyId} = ${heartbeatRuns.companyId} - and ${nativeRunFinalizations.resultId} is null - and ${nativeRunFinalizations.phase} <> 'terminal_failure'))::int`, - }).from(heartbeatRuns).where(and(...pendingConditions)); - return { total: Number(row?.total ?? 0), totalExact: normalizeCents(row?.total ?? 0), unpricedEventCount: Number(row?.unpricedEventCount ?? 0), - pendingRunCount: pending?.count ?? 0, recoveringRunCount: pending?.recovering ?? 0 }; + const [ledger = {}] = billable.length ? await db.select(ledgerColumns).from(costEvents).where(or(...ledgerPredicates)) : []; + const [pending = {}] = billable.length ? await db.select(pendingColumns).from(heartbeatRuns).where(or(...pendingPredicates)) : []; + return policies.map((_, index) => ({ + windowStart: windows[index].start, windowEnd: windows[index].end, + total: Number(ledger[`total${index}`] ?? 0), totalExact: normalizeCents(String(ledger[`total${index}`] ?? 0)), + unpricedEventCount: Number(ledger[`unpriced${index}`] ?? 0), pendingRunCount: Number(pending[`pending${index}`] ?? 0), + recoveringRunCount: Number(pending[`recovering${index}`] ?? 0), + })); } -async function computeObservedAmount(db: Db, policy: PolicyRow) { - return (await computeObservedSpend(db, policy)).total; -} -async function reached(db: Db, policy: PolicyRow, threshold: number) { - return compareCents((await computeObservedSpend(db, policy)).totalExact, threshold) >= 0; +export async function computeObservedSpend(db: Db, policy: SpendPolicy) { + return (await computeObservedSpends(db, [policy]))[0]; } + function observedBlocks(policy: PolicyRow, observed: Awaited>) { if (!policy.isActive || !policy.hardStopEnabled || policy.amount <= 0) return false; return compareCents(observed.totalExact, policy.amount) >= 0 || (observed.pendingRunCount - observed.recoveringRunCount) > 0 || (policy.unpricedUsagePolicy !== "allow" && observed.unpricedEventCount > 0); @@ -367,8 +389,9 @@ export function budgetServiceInTransaction(db: Db, publications: ActivityPublica async function buildPolicySummary(policy: PolicyRow): Promise { const scope = await resolveScopeRecord(db, policy.scopeType as BudgetScopeType, policy.scopeId); - const { total: observedAmount, totalExact: observedAmountExact, unpricedEventCount, pendingRunCount } = await computeObservedSpend(db, policy); - const { start, end } = resolveWindow(policy.windowKind as BudgetWindowKind); + const observed = await computeObservedSpend(db, policy); + const { total: observedAmount, totalExact: observedAmountExact, unpricedEventCount, pendingRunCount } = observed; + const { windowStart: start, windowEnd: end } = observed; const amount = policy.isActive ? policy.amount : 0; const utilizationPercent = amount > 0 ? Number(((observedAmount / amount) * 100).toFixed(2)) : 0; @@ -393,7 +416,7 @@ export function budgetServiceInTransaction(db: Db, publications: ActivityPublica hardStopEnabled: policy.hardStopEnabled, notifyEnabled: policy.notifyEnabled, isActive: policy.isActive, - status: await policyBlocks(db, policy) ? "hard_stop" : policy.isActive + status: observedBlocks(policy, observed) ? "hard_stop" : policy.isActive ? budgetStatusFromObserved(observedAmountExact, amount, policy.warnPercent) : "ok", paused: scope.paused, @@ -406,9 +429,9 @@ export function budgetServiceInTransaction(db: Db, publications: ActivityPublica async function createIncidentIfNeeded( policy: PolicyRow, thresholdType: BudgetThresholdType, - amountObserved: number, + observed: Awaited>, ) { - const { start, end } = resolveWindow(policy.windowKind as BudgetWindowKind); + const { windowStart: start, windowEnd: end } = observed; const existing = await db .select() .from(budgetIncidents) @@ -424,7 +447,7 @@ export function budgetServiceInTransaction(db: Db, publications: ActivityPublica .then((rows) => rows[0] ?? null); if (existing) return { incident: existing, created: false }; - const amountObservedExact = (await computeObservedSpend(db, policy)).totalExact; + const { total: amountObserved, totalExact: amountObservedExact } = observed; const scope = await resolveScopeRecord(db, policy.scopeType as BudgetScopeType, policy.scopeId); const payload = buildApprovalPayload({ policy, @@ -528,32 +551,30 @@ export function budgetServiceInTransaction(db: Db, publications: ActivityPublica : []; const approvalStatusById = new Map(approvalRows.map((row) => [row.id, row.status])); - return Promise.all( - rows.map(async (row) => { - const scope = await resolveScopeRecord(db, row.scopeType as BudgetScopeType, row.scopeId); - return { - id: row.id, - companyId: row.companyId, - policyId: row.policyId, - scopeType: row.scopeType as BudgetScopeType, - scopeId: row.scopeId, - scopeName: normalizeScopeName(row.scopeType as BudgetScopeType, scope.name), - metric: row.metric as BudgetMetric, - windowKind: row.windowKind as BudgetWindowKind, - windowStart: row.windowStart, - windowEnd: row.windowEnd, - thresholdType: row.thresholdType as BudgetThresholdType, - amountLimit: row.amountLimit, - amountObserved: row.amountObserved, - status: row.status as BudgetIncident["status"], - approvalId: row.approvalId ?? null, - approvalStatus: row.approvalId ? approvalStatusById.get(row.approvalId) ?? null : null, - resolvedAt: row.resolvedAt ?? null, - createdAt: row.createdAt, - updatedAt: row.updatedAt, - }; - }), - ); + return mapBudgetReads(rows, async (row) => { + const scope = await resolveScopeRecord(db, row.scopeType as BudgetScopeType, row.scopeId); + return { + id: row.id, + companyId: row.companyId, + policyId: row.policyId, + scopeType: row.scopeType as BudgetScopeType, + scopeId: row.scopeId, + scopeName: normalizeScopeName(row.scopeType as BudgetScopeType, scope.name), + metric: row.metric as BudgetMetric, + windowKind: row.windowKind as BudgetWindowKind, + windowStart: row.windowStart, + windowEnd: row.windowEnd, + thresholdType: row.thresholdType as BudgetThresholdType, + amountLimit: row.amountLimit, + amountObserved: row.amountObserved, + status: row.status as BudgetIncident["status"], + approvalId: row.approvalId ?? null, + approvalStatus: row.approvalId ? approvalStatusById.get(row.approvalId) ?? null : null, + resolvedAt: row.resolvedAt ?? null, + createdAt: row.createdAt, + updatedAt: row.updatedAt, + }; + }); } async function reconcileScope(companyId: string, scopeType: BudgetScopeType, scopeId: string) { @@ -566,13 +587,13 @@ export function budgetServiceInTransaction(db: Db, publications: ActivityPublica for (const policy of policies) { const observed = await computeObservedSpend(db, policy); if (observedBlocks(policy, observed)) { - await createIncidentIfNeeded(policy, "hard", observed.total); + await createIncidentIfNeeded(policy, "hard", observed); await pauseAndCancelScopeForBudget(policy); const title = scopeType === "company" ? "Company" : scopeType === "agent" ? "Agent" : "Project"; - if ((await computeObservedSpend(db, policy)).pendingRunCount > 0) { + if (observed.pendingRunCount > 0) { return { scopeType, scopeId, scopeName: scope.name, reason: `${title} cannot start work while completed runs await accounting.` }; } - if (policy.unpricedUsagePolicy !== "allow" && (await computeObservedSpend(db, policy)).unpricedEventCount > 0) { + if (policy.unpricedUsagePolicy !== "allow" && observed.unpricedEventCount > 0) { return { scopeType, scopeId, scopeName: scope.name, reason: `${title} cannot start work because recorded usage has no reliable price.` }; } return { scopeType, scopeId, scopeName: scope.name, reason: scopeType === "project" @@ -585,7 +606,7 @@ export function budgetServiceInTransaction(db: Db, publications: ActivityPublica // A closed UTC window, inactive policy, or budget raise releases only // budget-owned pauses. Manual pauses remain operator decisions. await resolveOpenIncidentsForPolicy(policy.id, "approved", "budget_service", "hard"); - if (!policy.isActive || !policy.notifyEnabled || !(await reached(db, policy, Math.ceil(policy.amount * policy.warnPercent / 100)))) { + if (!policy.isActive || !policy.notifyEnabled || compareCents(observed.totalExact, Math.ceil(policy.amount * policy.warnPercent / 100)) < 0) { await resolveOpenIncidentsForPolicy(policy.id, "approved", "budget_service", "soft"); } } @@ -706,23 +727,23 @@ export function budgetServiceInTransaction(db: Db, publications: ActivityPublica } if (row.isActive && amount > 0) { - const observedAmount = await computeObservedAmount(db, row); - if (!(await policyBlocks(db, row))) { + const observed = await computeObservedSpend(db, row); + if (!observedBlocks(row, observed)) { await resumeScopeFromBudget(row); await resolveOpenIncidentsForPolicy(row.id, "approved", actorUserId ?? "budget_service", "hard"); - if (row.notifyEnabled && (await reached(db, row, Math.ceil(row.amount * row.warnPercent / 100)))) { - await createIncidentIfNeeded(row, "soft", observedAmount); + if (row.notifyEnabled && compareCents(observed.totalExact, Math.ceil(row.amount * row.warnPercent / 100)) >= 0) { + await createIncidentIfNeeded(row, "soft", observed); } else { await resolveOpenIncidentsForPolicy(row.id, null, null, "soft"); } } else { const softThreshold = Math.ceil((row.amount * row.warnPercent) / 100); - if (row.notifyEnabled && (await reached(db, row, softThreshold)) && (!row.hardStopEnabled || !(await reached(db, row, row.amount)))) { - await createIncidentIfNeeded(row, "soft", observedAmount); + if (row.notifyEnabled && compareCents(observed.totalExact, softThreshold) >= 0 && (!row.hardStopEnabled || compareCents(observed.totalExact, row.amount) < 0)) { + await createIncidentIfNeeded(row, "soft", observed); } - if (await policyBlocks(db, row)) { + if (observedBlocks(row, observed)) { await resolveOpenSoftIncidents(row.id); - await createIncidentIfNeeded(row, "hard", observedAmount); + await createIncidentIfNeeded(row, "hard", observed); await pauseAndCancelScopeForBudget(row); } } @@ -751,7 +772,7 @@ export function budgetServiceInTransaction(db: Db, publications: ActivityPublica overview: async (companyId: string): Promise => { const rows = await listPolicyRows(companyId); - const policies = await Promise.all(rows.map((row) => buildPolicySummary(row))); + const policies = await mapBudgetReads(rows, buildPolicySummary); const activeIncidentRows = await db .select() .from(budgetIncidents) @@ -776,25 +797,19 @@ export function budgetServiceInTransaction(db: Db, publications: ActivityPublica and( eq(budgetPolicies.companyId, event.companyId), eq(budgetPolicies.isActive, true), - inArray(budgetPolicies.scopeType, ["company", "agent", "project"]), + budgetPoliciesForRun(event.companyId, event.agentId, event.projectId), ), ); - const relevantPolicies = candidatePolicies.filter((policy) => { - if (policy.scopeType === "company") return policy.scopeId === event.companyId; - if (policy.scopeType === "agent") return policy.scopeId === event.agentId; - if (policy.scopeType === "project") return Boolean(event.projectId) && policy.scopeId === event.projectId; - return false; - }); - - for (const policy of relevantPolicies) { - if (policy.metric !== "billed_cents" || policy.amount <= 0) continue; - const observed = await computeObservedSpend(db, policy); + const policies = candidatePolicies.filter(policy => policy.metric === "billed_cents" && policy.amount > 0); + const observations = await computeObservedSpends(db, policies); + for (const [index, policy] of policies.entries()) { + const observed = observations[index]; const observedAmount = observed.total; const softThreshold = Math.ceil((policy.amount * policy.warnPercent) / 100); if (policy.notifyEnabled && compareCents(observed.totalExact, softThreshold) >= 0 && (!policy.hardStopEnabled || compareCents(observed.totalExact, policy.amount) < 0)) { - const softIncident = await createIncidentIfNeeded(policy, "soft", observedAmount); + const softIncident = await createIncidentIfNeeded(policy, "soft", observed); if (softIncident?.created) { await recordActivity({ companyId: policy.companyId, @@ -815,7 +830,7 @@ export function budgetServiceInTransaction(db: Db, publications: ActivityPublica if (observedBlocks(policy, observed)) { await resolveOpenSoftIncidents(policy.id); - const hardIncident = await createIncidentIfNeeded(policy, "hard", observedAmount); + const hardIncident = await createIncidentIfNeeded(policy, "hard", observed); await pauseAndCancelScopeForBudget(policy); if (hardIncident?.created) { await recordActivity({ @@ -874,12 +889,12 @@ export function budgetServiceInTransaction(db: Db, publications: ActivityPublica const policy = await getPolicyRow(incident.policyId); if (input.action === "raise_budget_and_resume") { const nextAmount = Math.max(0, Math.floor(input.amount ?? 0)); - const currentObserved = await computeObservedAmount(db, policy); - if (policy.unpricedUsagePolicy !== "allow" && (await computeObservedSpend(db, policy)).unpricedEventCount > 0) { + const observed = await computeObservedSpend(db, policy); + if (policy.unpricedUsagePolicy !== "allow" && observed.unpricedEventCount > 0) { throw unprocessable("Resolve unpriced usage or explicitly allow it in the budget policy before resuming"); } - if ((await computeObservedSpend(db, policy)).pendingRunCount > 0) throw unprocessable("Completed runs must finish accounting before resuming"); - if (await reached(db, policy, nextAmount)) { + if (observed.pendingRunCount > 0) throw unprocessable("Completed runs must finish accounting before resuming"); + if (compareCents(observed.totalExact, nextAmount) >= 0) { throw unprocessable("New budget must exceed current observed spend"); } diff --git a/server/src/services/heartbeat.ts b/server/src/services/heartbeat.ts index 224f0cb03a1..ff0a6ad49e0 100644 --- a/server/src/services/heartbeat.ts +++ b/server/src/services/heartbeat.ts @@ -15,8 +15,8 @@ import { } from "./native-runtime/native-workspace-finalization-ownership.js"; import { hasStopOnlyCleanup, settleStopOnlyCleanup } from "./sandbox-stop-and-retain.js"; import { reserveRunBudget } from "./budget-reservations.js"; -import { accountRunCost, reconcileRunCosts } from "./run-cost-accounting.js"; -import { createRunUsageRecorder, replayUsageReceipts } from "./usage-receipts.js"; +import { accountRunCost, createCostAccountingReconciler } from "./run-cost-accounting.js"; +import { createRunUsageRecorder } from "./usage-receipts.js"; import { applyWorkspaceRestoreFailure } from "@paperclipai/adapter-utils/workspace-restore-result"; import { hasWorkspaceRestoreFailure } from "@paperclipai/shared"; import { externalConversationStateSql, nonIdleSlackIssueCondition } from "./slack-conversation-state.js"; @@ -30345,10 +30345,7 @@ export function heartbeatService( scanSilentActiveRuns, reconcileTaskWatchdogs, - reconcileCostAccounting: async () => { - await replayUsageReceipts(db); - return reconcileRunCosts(db, budgetHooks); - }, + reconcileCostAccounting: createCostAccountingReconciler(db, budgetHooks), buildRunOutputSilence, diff --git a/server/src/services/run-cost-accounting.ts b/server/src/services/run-cost-accounting.ts index e5e218efb05..ea52ab1bb94 100644 --- a/server/src/services/run-cost-accounting.ts +++ b/server/src/services/run-cost-accounting.ts @@ -6,7 +6,7 @@ import { createCostEventInTransaction } from "./costs.js"; import { budgetService, deliverBudgetEnforcement, type BudgetServiceHooks } from "./budgets.js"; import { logger } from "../middleware/logger.js"; import { promises as fs } from "node:fs"; -import { indexPendingUsageReceipts, recoverPendingRunUsageReceipts, type UsageReceiptIndex } from "./usage-receipts.js"; +import { replayUsageReceipts, indexPendingUsageReceipts, recoverPendingRunUsageReceipts, type UsageReceiptIndex } from "./usage-receipts.js"; const terminalStatuses = ["succeeded", "failed", "timed_out", "cancelled", "interrupted"]; const object = (value: unknown): Record => value && typeof value === "object" && !Array.isArray(value) ? value as Record : {}; @@ -164,3 +164,22 @@ export async function reconcileRunCosts(db: Db, hooks: BudgetServiceHooks = {}) await budgetService(db, hooks).reconcilePolicies(); return { scanned: pending.length, accounted }; } + +/** Share one full sweep across overlapping scheduler ticks. Do not enqueue + * another sweep: the next tick retries after either success or failure. */ +export function createCostAccountingReconciler(db: Db, hooks: BudgetServiceHooks = {}) { + let inFlight: Promise>> | undefined; + return () => { + if (!inFlight) { + inFlight = (async () => { + try { + await replayUsageReceipts(db); + return await reconcileRunCosts(db, hooks); + } finally { + inFlight = undefined; + } + })(); + } + return inFlight; + }; +}