From 08dafde1d65a8f3299237af04df5e2543bb4382d Mon Sep 17 00:00:00 2001 From: Devin Foley Date: Sat, 24 Oct 2026 23:58:43 -0700 Subject: [PATCH] perf(accounting): narrow locked usage receipt run reads (#17559) Select only the five run fields used by usage receipt persistence under the existing company and row locks. Keep receipt history, replay checks, budget updates, and stored execution evidence intact. Validation: 37 focused tests, full build and typecheck, exact-head Greptile 5/5 (medium), and full green CI with no unresolved threads. The first CI attempt failed an unchanged Runner shutdown test; its qualified local case and same-head failed-lane rerun passed. Co-Authored-By: Paperclip --- .../src/services/usage-receipts-read.test.ts | 81 +++++++++++++++++++ server/src/services/usage-receipts.ts | 5 +- 2 files changed, 85 insertions(+), 1 deletion(-) create mode 100644 server/src/services/usage-receipts-read.test.ts diff --git a/server/src/services/usage-receipts-read.test.ts b/server/src/services/usage-receipts-read.test.ts new file mode 100644 index 00000000000..83ddbed9364 --- /dev/null +++ b/server/src/services/usage-receipts-read.test.ts @@ -0,0 +1,81 @@ +import { randomBytes, randomUUID } from "node:crypto"; +import { eq } from "drizzle-orm"; +import { afterAll, beforeAll, describe, expect, it } from "vitest"; +import { agents, companies, createDb, heartbeatRuns, runUsageReceipts, type Db } from "@paperclipai/db"; +import { getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase } from "../__tests__/helpers/embedded-postgres.js"; +import { persistUsageReceipt, type UsageReceiptEnvelope } from "./usage-receipts.js"; + +const support = await getEmbeddedPostgresTestSupport(); +describe.skipIf(!support.supported)("usage receipt PostgreSQL reads", () => { + let database: Awaited>; + let db: Db; + beforeAll(async () => { + database = await startEmbeddedPostgresTestDatabase("usage-receipt-read-"); + db = createDb(database.connectionString); + }, 60_000); + afterAll(async () => { await db?.$client.end(); await database?.cleanup(); }); + + async function fixture() { + const companyId = randomUUID(), agentId = randomUUID(), runId = randomUUID(), sourceId = randomUUID(); + await db.insert(companies).values({ id: companyId, name: "Receipt fixture", issuePrefix: companyId.slice(0, 8) }); + await db.insert(agents).values({ id: agentId, companyId, name: "Receipt agent", adapterType: "process", status: "active" }); + const payload = randomBytes(128 * 1024).toString("base64"); + const evidence = { contextSnapshot: { prompt: payload }, runnerProfileJson: { instructions: payload }, + resultJson: { stdout: payload }, stdoutExcerpt: payload, stderrExcerpt: payload }; + await db.insert(heartbeatRuns).values({ id: runId, companyId, agentId, status: "running", ...evidence, + usageJson: { accountingReceiptSourceId: sourceId, accountingReceiptSequence: 0, preserved: "receipt metadata" } }); + const receipt = (sequence: number): UsageReceiptEnvelope => ({ schema: "paperclip/accounting-receipt/v1", + id: randomUUID(), companyId, runId, sourceId, sequence, receivedAt: new Date().toISOString(), adapterType: "process", + receipt: { complete: false, usageBasis: "per_run", costUsd: sequence / 100, billingType: "metered_api" } }); + return { companyId, runId, evidence, receipt }; + } + + it("updates live usage with a small locked read and preserves all execution evidence", async () => { + const f = await fixture(), reads: unknown[][] = []; + const observed = new Proxy(db, { get(target, key) { + if (key !== "transaction") return Reflect.get(target, key); + return (callback: Parameters[0]) => target.transaction(tx => callback(new Proxy(tx, { + get(transaction, property) { + if (property !== "select") return Reflect.get(transaction, property); + return (...fields: Parameters) => { + const builder = transaction.select(...fields), from = builder.from.bind(builder); + builder.from = ((...args: Parameters) => { + const query = from(...args); + if (args[0] === heartbeatRuns) { + const execute = query.execute.bind(query); + query.execute = async (...executeArgs: Parameters) => { + expect(query.toSQL().sql).toContain("for update"); + const rows = await execute(...executeArgs); reads.push(rows); return rows; + }; + } + return query; + }) as typeof builder.from; + return builder; + }; + }, + }))); + } }); + await persistUsageReceipt(observed, f.receipt(1)); + expect(reads).toHaveLength(1); + expect(reads[0]).toHaveLength(1); + expect(Buffer.byteLength(JSON.stringify(reads))).toBeLessThan(2048); + const [stored] = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, f.runId)); + expect(stored).toMatchObject(f.evidence); + expect(stored.usageJson).toMatchObject({ preserved: "receipt metadata", accountingReceiptSequence: 1, + budgetExposure: { costCents: "1.0000000" } }); + expect(Buffer.byteLength(JSON.stringify(stored))).toBeGreaterThan(512 * 1024); + }); + + it("serializes competing checkpoints and retains replay conflicts and company isolation", async () => { + const f = await fixture(), first = f.receipt(1), second = f.receipt(2); + await Promise.all([persistUsageReceipt(db, first), persistUsageReceipt(db, second)]); + await persistUsageReceipt(db, first); + const [stored] = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, f.runId)); + expect(stored.usageJson).toMatchObject({ accountingReceiptSequence: 2, budgetExposure: { costCents: "2.0000000" } }); + expect(await db.select().from(runUsageReceipts).where(eq(runUsageReceipts.runId, f.runId))).toHaveLength(2); + await expect(persistUsageReceipt(db, { ...first, receipt: { ...first.receipt, costUsd: 10 } })).rejects.toThrow("Conflicting"); + const foreign = await fixture(); + await expect(persistUsageReceipt(db, { ...f.receipt(3), companyId: foreign.companyId })).rejects.toThrow("not found"); + expect((await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, f.runId)))[0]).toEqual(stored); + }); +}); diff --git a/server/src/services/usage-receipts.ts b/server/src/services/usage-receipts.ts index 52bae89f8fb..0f99401d9da 100644 --- a/server/src/services/usage-receipts.ts +++ b/server/src/services/usage-receipts.ts @@ -68,7 +68,10 @@ export async function persistUsageReceipt(db: Db, raw: UsageReceiptEnvelope) { const envelope = envelopeSchema.parse(raw); const hash = receiptFingerprint(envelope); return withAccountingTransaction(db, envelope.companyId, async (tx, publications) => { - const [run] = await tx.select().from(heartbeatRuns).where(and(eq(heartbeatRuns.id, envelope.runId), eq(heartbeatRuns.companyId, envelope.companyId))).for("update"); + const [run] = await tx.select({ + id: heartbeatRuns.id, agentId: heartbeatRuns.agentId, status: heartbeatRuns.status, + usageJson: heartbeatRuns.usageJson, costAccountedAt: heartbeatRuns.costAccountedAt, + }).from(heartbeatRuns).where(and(eq(heartbeatRuns.id, envelope.runId), eq(heartbeatRuns.companyId, envelope.companyId))).for("update"); if (!run) throw notFound("Accounting run not found"); const [existing] = await tx.select().from(runUsageReceipts).where(and(eq(runUsageReceipts.companyId, envelope.companyId), eq(runUsageReceipts.runId, envelope.runId), eq(runUsageReceipts.sourceId, envelope.sourceId), eq(runUsageReceipts.sequence, envelope.sequence))); if (existing && existing.receiptHash !== hash) throw conflict("Conflicting usage checkpoint replay");