diff --git a/doc/DATABASE.md b/doc/DATABASE.md index 89db9c4acdc..265f48d560d 100644 --- a/doc/DATABASE.md +++ b/doc/DATABASE.md @@ -2166,6 +2166,23 @@ keeps the normal late-receipt reconciliation path. Other independent accounting operations retain their repeatable-read snapshot. Receipt identity, replay conflicts, company scoping and budget-work publication remain transactional. +Run budget admission (`reserveRunBudget`), the final dispatch gate's accounting +phase (resolved-interaction dispatch only), the usage recorder's receipt-source handoff and its capture-failure fence +use the same pattern. The dispatch gate's ownership phase already ran at READ +COMMITTED under the same issue and run locks. Retry promotion and stale-run +cancellation keep their repeatable-read snapshot because they combine many facts +in one decision. Admission reads all +mutable state after the shared admission fences, the agent share lock and the +run-row lock. Its exposure is one statement, and the per-run reservation index +plus the run lock prevent a second hold. Run-log appends and output progress that +update a preparing run therefore cannot exhaust admission retries. Concurrent +admissions of different runs remain best-effort, as before. When any operation +transaction still exhausts its retries, it raises `AccountingContentionError`. The +error carries a bounded message, keeps the SQLSTATE in `code` and the database +error in `cause`, and is logged with both. A run that fails this way keeps the +generic `adapter_failed` code, so existing retry paths apply, and records that +bounded message with the SQLSTATE as its failure cause. + Upgrade all controllers/workers to this accounting contract and drain incompatible pre-upgrade database transactions before enabling independent writers. Do not run pre-upgrade synchronous counter/evaluation writers beside unlocked writers. diff --git a/server/src/modules/run-dispatch/adapters/postgres.ts b/server/src/modules/run-dispatch/adapters/postgres.ts index 5174c71137e..5618daa66ef 100644 --- a/server/src/modules/run-dispatch/adapters/postgres.ts +++ b/server/src/modules/run-dispatch/adapters/postgres.ts @@ -227,6 +227,7 @@ export function createPostgresRunDispatchAdapter( onMissing: () => T, operation: (tx: Db, run: HeartbeatRun, publications: ActivityPublication[]) => Promise, mode: "accounting" | "ownership", + isolation: "snapshot" | "read committed" = "snapshot", ): Promise { // Queue editing and claiming already use issue -> wake -> run. Read the // immutable issue reference without a lock first, then acquire issue -> @@ -282,7 +283,8 @@ export function createPostgresRunDispatchAdapter( return { kind: "value" as const, value: await operation(typedTx, run, publications) }; }; const result: { kind: "missing" | "retry" } | { kind: "value"; value: T } = mode === "accounting" - ? await withOperationAccountingTransaction(db, input.companyId, locked) + ? await withOperationAccountingTransaction(db, input.companyId, locked, + isolation === "read committed" ? { isolationLevel: "read committed" } : {}) : await db.transaction(tx => locked(tx as unknown as Db, [])); if (result.kind === "missing") return onMissing(); @@ -294,12 +296,21 @@ export function createPostgresRunDispatchAdapter( ); } + /** Accounting phase under issue -> run locks. Gate decisions that combine + * many facts (retry promotion, stale-run cancellation) keep one repeatable- + * read snapshot. The final dispatch gate passes READ COMMITTED: run-log + * appends and output progress update a running row while it waits, and a + * snapshot taken before the run lock went stale on each one until every + * retry was spent. Its mutable reads all follow the admission fences and the + * issue/run locks, budget exposure is one statement, and its ownership phase + * re-reads the staleness facts at READ COMMITTED under the same locks. */ function withIssueThenRunLocks( input: { runId: string; companyId: string; recoveryLeaseOwner?: string; enforceContinuationOwnership?: boolean }, onMissing: () => T, operation: (tx: Db, run: HeartbeatRun, publications: ActivityPublication[]) => Promise, + isolation: "snapshot" | "read committed" = "snapshot", ) { - return withRunOwnershipLocks(input, onMissing, operation, "accounting"); + return withRunOwnershipLocks(input, onMissing, operation, "accounting", isolation); } async function loadGateFacts( @@ -1474,6 +1485,7 @@ export function createPostgresRunDispatchAdapter( ); }, dispatchLockedRun, + "read committed", ); if (!admission.dispatched) return admission; // Accounting has committed. Close the separate issue/run ownership race diff --git a/server/src/services/accounting-transaction.ts b/server/src/services/accounting-transaction.ts index 505c9983260..e87eef79aa0 100644 --- a/server/src/services/accounting-transaction.ts +++ b/server/src/services/accounting-transaction.ts @@ -1,7 +1,7 @@ import { eq } from "drizzle-orm"; import { companies, type Db } from "@paperclipai/db"; import { publishActivity, type ActivityPublication } from "./activity-log.js"; -import { notFound } from "../errors.js"; +import { HttpError, notFound } from "../errors.js"; import { logger } from "../middleware/logger.js"; import { withBudgetSpan } from "./budget-tracing.js"; @@ -42,20 +42,55 @@ export async function withOperationAccountingTransaction(db: Db, companyId: s publishAccountingActivities(companyId, publications); return result; } catch (error) { - if (attempt >= 3 || !isAccountingRetryable(error)) throw error; + if (!isAccountingRetryable(error)) throw error; + if (attempt >= 3) { + const contention = new AccountingContentionError(error, attempt + 1); + logger.warn({ companyId, sqlState: contention.code, attempts: contention.attempts, err: error }, + "Accounting transaction exhausted its serialization retries"); + throw contention; + } await new Promise(resolve => setTimeout(resolve, 10 * (attempt + 1))); } } } -export function isAccountingRetryable(error: unknown): boolean { +export const ACCOUNTING_CONTENTION_MESSAGE = "Paperclip could not save this run's budget or usage records because other " + + "updates to the same run kept conflicting with them. Retry the run."; + +/** Bounded, SQL-free failure for exhausted serialization retries. `code` + * keeps the SQLSTATE (40001 or 40P01) for retry classifiers and diagnostics; + * the database error, with its statement text, stays only in `cause`. */ +export class AccountingContentionError extends HttpError { + readonly code: string; + constructor(cause: unknown, readonly attempts: number) { + super(503, ACCOUNTING_CONTENTION_MESSAGE); + this.name = "AccountingContentionError"; + this.code = accountingRetryCode(cause) ?? "40001"; + this.cause = cause; + } +} + +export function isAccountingContentionError(error: unknown): error is AccountingContentionError { + let current = error; + for (let depth = 0; current && typeof current === "object" && depth < 4; depth++) { + if (current instanceof AccountingContentionError) return true; + current = (current as { cause?: unknown }).cause; + } + return false; +} + +function accountingRetryCode(error: unknown): string | null { let current = error; for (let depth = 0; current && typeof current === "object" && depth < 4; depth++) { const row = current as { code?: string; cause?: unknown }; - if (row.code === "40001" || row.code === "40P01") return true; + if (row.code === "40001" || row.code === "40P01") return row.code; current = row.cause; } - return false; + return null; +} + +export function isAccountingRetryable(error: unknown): boolean { + return accountingRetryCode(error) !== null; } /** Consistent operator reports must not wait for or delay accounting writers. diff --git a/server/src/services/budget-reservations.ts b/server/src/services/budget-reservations.ts index c9fec7555c5..c417a4d0cd4 100644 --- a/server/src/services/budget-reservations.ts +++ b/server/src/services/budget-reservations.ts @@ -78,7 +78,15 @@ export async function recheckRunBudgetReservationInTransaction( /** Reserve before dispatch, under the same company lock as charges and policy * changes. Estimates constrain admission; they cannot cap a provider's bill. - * A reservation survives timeouts and restarts until accounting proves closure. */ + * A reservation survives timeouts and restarts until accounting proves closure. + * + * READ COMMITTED: every mutable read follows the shared admission fences, the + * agent share lock and the run-row lock, so it sees each commit that preceded + * those locks. Run-log appends and output progress update the run row while a + * run prepares; a repeatable-read snapshot taken before the run lock went stale + * on each such commit and could exhaust every retry. Exposure is one statement + * (hold-to-ledger atomic), and the per-run unique index plus the run lock + * fence a duplicate reservation. */ export async function reserveRunBudget(db: Db, companyId: string, runId: string, projectId: string | null, ledgerScope: Record = {}, recoveryLeaseOwner?: string) { const result = await withBudgetSpan("budget.admission.reserve", () => withOperationAccountingTransaction(db, companyId, async (tx, publications) => { await lockBudgetAdmissionControls(tx, companyId, projectId); @@ -142,7 +150,7 @@ export async function reserveRunBudget(db: Db, companyId: string, runId: string, }).where(eq(heartbeatRuns.id, runId)); await queueBudgetEvaluation(tx, { companyId, agentId: run.agentId, projectId }); return { ...reservation, reused: false }; - })); + }, { isolationLevel: "read committed" })); // Commit required block/resume transitions even when dispatch is denied. // Throwing inside the transaction would discard those durable obligations. if ("admissionBlockReason" in result) throw conflict(result.admissionBlockReason ?? "Budget admission blocked"); diff --git a/server/src/services/heartbeat/run-completion.test.ts b/server/src/services/heartbeat/run-completion.test.ts index 7ce45027e9f..bd9c3a4ce69 100644 --- a/server/src/services/heartbeat/run-completion.test.ts +++ b/server/src/services/heartbeat/run-completion.test.ts @@ -38,6 +38,7 @@ import { withS3FailureDiagnostics } from "../../storage/s3-failure.js"; import { workspaceProviderCallFailure } from "../workspace-provider-failure.js"; import { logger } from "../../middleware/logger.js"; import { unprocessable } from "../../errors.js"; +import { ACCOUNTING_CONTENTION_MESSAGE, AccountingContentionError } from "../accounting-transaction.js"; import { createHeartbeatLifecycle, type HeartbeatLifecycleDependencies } from "./run-lifecycle.js"; import { createHeartbeatRunPreparation, ConfigurationIncompleteFailure, classifyManagedProjectAuthFailure } from "./run-preparation.js"; import { createHeartbeatRunState, getAdapterSessionCodec } from "./run-state.js"; @@ -1414,6 +1415,24 @@ describe.skipIf(!support.supported && !process.env.PAPERCLIP_TEST_DATABASE_URL)( expect(f.deps.scheduleInteractionContinuationInfrastructureRetryIfEligible).not.toHaveBeenCalled(); }); + it("saves a bounded cause and keeps the retryable generic code when accounting retries are exhausted", async () => { + const f = await fixture(); + const logged = vi.spyOn(logger, "error"); + onTestFinished(() => logged.mockRestore()); + const driverError = Object.assign(new Error("could not serialize access due to concurrent update"), + { name: "PostgresError", code: "40001" }); + const queryError = new Error("Failed query: select \"id\" from \"heartbeat_runs\" where \"id\" = $1 for update\nparams: secret-run-id", + { cause: driverError }); + await f.completion.failRun(f.failure(new AccountingContentionError(queryError, 4))); + const failed = await read(f.run.id); + // adapter_failed keeps chat and recovery retry paths that key on it. + expect(failed).toMatchObject({ status: "failed", errorCode: "adapter_failed", error: ACCOUNTING_CONTENTION_MESSAGE, + resultJson: { failureCause: "40001" } }); + expect(JSON.stringify(failed)).not.toMatch(/heartbeat_runs|secret-run-id|Failed query/); + const failureLog = logged.mock.calls.find(call => call[1] === "heartbeat execution failed"); + expect(failureLog?.[0]).toMatchObject({ runId: f.run.id, errorCode: "adapter_failed", failureCause: "40001" }); + }); + it("records a workspace provider's reason on the run and in the setup failure log without retrying", async () => { const f = await fixture(); const logged = vi.spyOn(logger, "error"); diff --git a/server/src/services/heartbeat/run-completion.ts b/server/src/services/heartbeat/run-completion.ts index 0aae8d1725f..d62d1a38911 100644 --- a/server/src/services/heartbeat/run-completion.ts +++ b/server/src/services/heartbeat/run-completion.ts @@ -65,6 +65,7 @@ import type { BudgetServiceHooks } from "../budgets.js"; import type { providerTraceStore } from "../provider-trace-store.js"; import type { createAgentIdentityRedactor } from "../agent-identity-redaction.js"; import type { CurrentUserRedactionOptions } from "../../log-redaction.js"; +import { ACCOUNTING_CONTENTION_MESSAGE, isAccountingContentionError } from "../accounting-transaction.js"; import type { AdapterExecutionResult } from "@paperclipai/adapter-utils"; import type { getRunLogStore, RunLogHandle } from "../run-log-store.js"; import type { createHeartbeatRunState, getAdapterSessionCodec } from "./run-state.js"; @@ -1110,7 +1111,11 @@ export function createHeartbeatRunCompletion(db: Db, dependencies: HeartbeatRunC redactCurrentUserText(identityRedactor.redact(text), currentUserRedactionOptions); // Bounded code from the error or its first three causes; never message text. const thrownFailureCause = failureCauseFromError(err, providerFailure); - const thrownMessage = withoutSqlStatementText( + // Exhausted accounting retries carry their own bounded message. The code + // stays generic (adapter_failed) so existing retry paths still apply; + // failureCause records the SQLSTATE. + const accountingContention = isAccountingContentionError(err); + const thrownMessage = accountingContention ? ACCOUNTING_CONTENTION_MESSAGE : withoutSqlStatementText( redactFailureText(err instanceof Error ? err.message : "Unknown adapter failure"), thrownFailureCause); const workspaceValidationFailure = isWorkspaceValidationFailure(err) ? err diff --git a/server/src/services/run-accounting-contention.test.ts b/server/src/services/run-accounting-contention.test.ts new file mode 100644 index 00000000000..6a9e4ef2b10 --- /dev/null +++ b/server/src/services/run-accounting-contention.test.ts @@ -0,0 +1,303 @@ +/** + * Run-row accounting writers under concurrent run-row updates (real PostgreSQL). + * + * Streamed run-log appends, output progress and lifecycle events update the + * same `heartbeat_runs` row that budget admission and the usage recorder lock. + * A repeatable-read snapshot taken before that lock goes stale whenever one of + * those writes commits first, and four retries could not outlast a busy stream + * (restart-chaos L1 with `--db-pressure`). The deterministic cases commit a real + * run-event append after the writer's transaction validates the company and + * before it locks the run, on every attempt, which is the race a stream produces. + */ +import { randomUUID } from "node:crypto"; +import fs from "node:fs/promises"; +import os from "node:os"; +import path from "node:path"; +import { and, eq, sql } from "drizzle-orm"; +import { afterAll, beforeAll, describe, expect, it } from "vitest"; +import { + agents, budgetEvaluationWork, budgetPolicies, budgetReservations, companies, createDb, heartbeatRuns, runUsageReceipts, type Db, +} from "@paperclipai/db"; +import { getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase } from "../__tests__/helpers/embedded-postgres.js"; +import { + AccountingContentionError, isAccountingContentionError, withOperationAccountingTransaction, +} from "./accounting-transaction.js"; +import { createPostgresRunDispatchAdapter } from "../modules/run-dispatch/adapters/postgres.js"; +import { reserveRunBudget } from "./budget-reservations.js"; +import { appendHeartbeatRunEvent } from "./heartbeat-run-events.js"; +import { createRunUsageRecorder, spoolUsageReceipt, type UsageReceiptEnvelope } from "./usage-receipts.js"; + +const support = await getEmbeddedPostgresTestSupport(); + +describe.skipIf(!support.supported)("run accounting writers under concurrent run-row updates", () => { + let database: Awaited>; + let db: Db; + beforeAll(async () => { + database = await startEmbeddedPostgresTestDatabase("run-accounting-contention-"); + db = createDb(database.connectionString); + }, 120_000); + afterAll(async () => { await db?.$client.end(); await database?.cleanup(); }); + + async function fixture(options: { status?: "queued" | "running"; limitCents?: number; reservationCents?: string } = {}) { + const companyId = randomUUID(), agentId = randomUUID(); + await db.insert(companies).values({ id: companyId, name: "Contention lab", issuePrefix: `C${companyId.slice(0, 5)}` }); + await db.insert(agents).values({ id: agentId, companyId, name: "Worker", role: "engineer", status: "idle", adapterType: "process" }); + if (options.limitCents !== undefined) { + await db.insert(budgetPolicies).values({ companyId, scopeType: "company", scopeId: companyId, metric: "billed_cents", + windowKind: "calendar_month_utc", amount: options.limitCents, reservationCents: options.reservationCents ?? "0", + notifyEnabled: false, unpricedUsagePolicy: "allow" }); + } + const run = async (status = options.status ?? "queued") => { + const [row] = await db.insert(heartbeatRuns).values({ companyId, agentId, invocationSource: "assignment", status }).returning(); + return row.id; + }; + return { companyId, agentId, run }; + } + + /** Commit a real run-event append after the accounting transaction checks + * the company (its first statement, which starts the snapshot) and before + * it locks the run. Fires on every attempt, so a snapshot-sensitive writer + * exhausts its retries deterministically. */ + function interleaveRunEvents(f: { companyId: string; agentId: string }, runId: string) { + let interleavings = 0; + const observed = new Proxy(db, { get(target, key) { + if (key !== "transaction") return Reflect.get(target, key); + return (callback: Parameters[0], options: Parameters[1]) => { + // Only accounting transactions (which name an isolation level). A plain + // transaction can already hold the run lock where it reads companies. + let first = Boolean(options?.isolationLevel); + return 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] === companies && first) { + first = false; + const execute = query.execute.bind(query); + query.execute = async (...executeArgs: Parameters) => { + const rows = await execute(...executeArgs); + await appendHeartbeatRunEvent(db, { companyId: f.companyId, runId, agentId: f.agentId, + eventType: "lifecycle", message: "stream advanced" }); + interleavings++; + return rows; + }; + } + return query; + }) as typeof builder.from; + return builder; + }; + } })), options); + }; + } }); + return { db: observed, interleavings: () => interleavings }; + } + + /** Back-to-back single-row writes, like a fast stream, until stopped. */ + function pressure(runId: string, writers = 4) { + const client = createDb(database.connectionString).$client; + let running = true, updates = 0; + const loops = Array.from({ length: writers }, async () => { + while (running) { + await client.unsafe("update heartbeat_runs set updated_at = updated_at where id = $1", [runId]); + updates++; + } + }); + return { + updates: () => updates, + async stop() { running = false; await Promise.all(loops); await client.end(); }, + }; + } + + async function steady(writers: ReturnType) { + while (writers.updates() < 50) await new Promise(resolve => setTimeout(resolve, 5)); + } + + describe("reserveRunBudget", () => { + it("reserves once when run events commit between its company check and its run lock", async () => { + const f = await fixture({ limitCents: 1_000, reservationCents: "100" }); + const runId = await f.run(); + const hooked = interleaveRunEvents(f, runId); + await expect(reserveRunBudget(hooked.db, f.companyId, runId, null)).resolves.toMatchObject({ reused: false }); + expect(hooked.interleavings()).toBe(1); + expect(await db.select().from(budgetReservations).where(eq(budgetReservations.runId, runId))).toHaveLength(1); + const [run] = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, runId)); + // The admission write and the concurrent stream write both survive. + expect(run.costAccountingPending).toBe(true); + expect(run.nextEventSeq).toBe(2); + }); + + it("reserves every run while back-to-back writers update the run row", async () => { + const f = await fixture({ limitCents: 100_000, reservationCents: "1" }); + for (let i = 0; i < 6; i++) { + const runId = await f.run(); + const writers = pressure(runId); + await steady(writers); + try { + await expect(reserveRunBudget(db, f.companyId, runId, null)).resolves.toMatchObject({ reused: false }); + } finally { await writers.stop(); } + } + }, 60_000); + + it("never creates a second reservation for one run when two dispatches race", async () => { + const f = await fixture({ limitCents: 1_000, reservationCents: "100" }); + const runId = await f.run(); + const outcomes = await Promise.allSettled([ + reserveRunBudget(db, f.companyId, runId, null), reserveRunBudget(db, f.companyId, runId, null), + ]); + expect(outcomes.filter(outcome => outcome.status === "fulfilled")).toHaveLength(1); + const rejected = outcomes.find(outcome => outcome.status === "rejected") as PromiseRejectedResult; + expect(rejected.reason).toMatchObject({ status: 409, message: "Provider dispatch has already reserved this run" }); + expect(await db.select().from(budgetReservations).where(eq(budgetReservations.runId, runId))).toHaveLength(1); + }); + + it("counts every committed hold, so successive admissions never exceed the budget", async () => { + // Each reservation fits alone; the third would exceed 250 cents. + const f = await fixture({ limitCents: 250, reservationCents: "100" }); + const results: string[] = []; + for (let i = 0; i < 3; i++) { + const runId = await f.run(); + const hooked = interleaveRunEvents(f, runId); + results.push(await reserveRunBudget(hooked.db, f.companyId, runId, null).then(() => "reserved", + (error: Error) => error.message)); + } + expect(results).toEqual(["reserved", "reserved", "Available budget is reserved by unfinished runs"]); + const [held] = await db.select({ cents: sql`sum(${budgetReservations.amountCents})::text` }).from(budgetReservations) + .where(and(eq(budgetReservations.companyId, f.companyId), eq(budgetReservations.state, "held"))); + expect(Number(held.cents)).toBeLessThanOrEqual(250); + }); + + it("still refuses a run once the hard stop is reached", async () => { + const f = await fixture({ limitCents: 100 }); + // Another run's observed spend already exceeds the limit. + await db.insert(heartbeatRuns).values({ companyId: f.companyId, agentId: f.agentId, invocationSource: "assignment", + status: "running", costAccountingPending: true, usageJson: { costUsd: 2, billingType: "metered_api" } }); + const runId = await f.run(); + const hooked = interleaveRunEvents(f, runId); + await expect(reserveRunBudget(hooked.db, f.companyId, runId, null)).rejects.toMatchObject({ status: 409 }); + expect(await db.select().from(budgetReservations).where(eq(budgetReservations.runId, runId))).toHaveLength(0); + }); + }); + + describe("final dispatch gate", () => { + it("admits and hands off a reserved run when run events commit between its company check and its run lock", async () => { + const f = await fixture({ limitCents: 1_000, reservationCents: "100" }); + const runId = await f.run(); + const reservation = await reserveRunBudget(db, f.companyId, runId, null); + await db.update(heartbeatRuns).set({ status: "running" }).where(eq(heartbeatRuns.id, runId)); + const hooked = interleaveRunEvents(f, runId); + let dispatched = 0; + const outcome = await createPostgresRunDispatchAdapter(hooked.db).dispatchResolvedInteractionIfCurrent({ + companyId: f.companyId, runId, reservationId: (reservation as { id: string }).id, expectedStatus: "running", + now: new Date(), dispatch: async () => { dispatched++; } }); + expect(outcome).toMatchObject({ dispatched: true }); + expect(dispatched).toBe(1); + expect(hooked.interleavings()).toBe(1); + const [run] = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, runId)); + expect(run.usageJson?.budgetDispatch).toMatchObject({ reservationId: (reservation as { id: string }).id }); + expect((run.usageJson?.budgetDispatch as { handedOffAt?: string }).handedOffAt).toBeTruthy(); + }); + + it("admits while back-to-back writers update the run row", async () => { + const f = await fixture({ limitCents: 100_000, reservationCents: "1" }); + for (let i = 0; i < 6; i++) { + const runId = await f.run(); + const reservation = await reserveRunBudget(db, f.companyId, runId, null); + await db.update(heartbeatRuns).set({ status: "running" }).where(eq(heartbeatRuns.id, runId)); + const writers = pressure(runId); + await steady(writers); + try { + await expect(createPostgresRunDispatchAdapter(db).dispatchResolvedInteractionIfCurrent({ companyId: f.companyId, runId, + reservationId: (reservation as { id: string }).id, expectedStatus: "running", now: new Date(), dispatch: async () => {} })) + .resolves.toMatchObject({ dispatched: true }); + } finally { await writers.stop(); } + } + }, 60_000); + }); + + describe("createRunUsageRecorder", () => { + const spool = () => fs.mkdtemp(path.join(os.tmpdir(), "run-accounting-contention-")); + + it("binds its receipt source when run events commit between its company check and its run lock", async () => { + const f = await fixture({ status: "running" }); + const runId = await f.run(); + const directory = await spool(); + // A receipt the previous controller spooled but never committed: the + // handoff journals it exactly once. + const previousSource = randomUUID(); + await db.update(heartbeatRuns).set({ usageJson: { accountingReceiptSourceId: previousSource, accountingReceiptSequence: 0 } }) + .where(eq(heartbeatRuns.id, runId)); + const pending: UsageReceiptEnvelope = { schema: "paperclip/accounting-receipt/v1", id: randomUUID(), companyId: f.companyId, + runId, sourceId: previousSource, sequence: 1, receivedAt: new Date().toISOString(), adapterType: "process", + receipt: { complete: false, usageBasis: "per_run", costUsd: 0.01, billingType: "metered_api" } }; + await spoolUsageReceipt(pending, directory); + try { + const hooked = interleaveRunEvents(f, runId); + await expect(createRunUsageRecorder(hooked.db, { companyId: f.companyId, runId, adapterType: "process" }, directory)).resolves.toBeTruthy(); + expect(hooked.interleavings()).toBe(1); + const [run] = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, runId)); + expect(run.usageJson).toMatchObject({ accountingReceiptSequence: 0, accountingReceiptReady: false, + accountingReceiptSourceFences: { [previousSource]: 1 } }); + expect(run.usageJson?.accountingReceiptSourceId).not.toBe(previousSource); + expect(run.nextEventSeq).toBe(2); + expect(await db.select().from(runUsageReceipts).where(eq(runUsageReceipts.runId, runId))).toHaveLength(1); + } finally { await fs.rm(directory, { recursive: true, force: true }); } + }); + + it("persists its capture-failure fence once when run events commit between its company check and its run lock", async () => { + const f = await fixture({ status: "running" }); + const runId = await f.run(); + const directory = await spool(); + try { + const hooked = interleaveRunEvents(f, runId); + const recorder = await createRunUsageRecorder(hooked.db, { companyId: f.companyId, runId, adapterType: "process" }, directory); + // A checkpoint the schema rejects fails the capture; the recorder then + // fences the run as capture-failed in its own accounting transaction. + await expect(recorder.capture(null as never)).rejects.toThrow(); + await recorder.persistFailure(); + // Handoff, fence and the repeated fence call each saw one concurrent + // run-event append. + expect(hooked.interleavings()).toBe(3); + const [run] = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, runId)); + expect(run.usageJson).toMatchObject({ accountingCaptureFailed: true, accountingReceiptReady: false }); + expect(run).toMatchObject({ costAccountingPending: true, accountingLastError: "usage_capture_failed", nextEventSeq: 4 }); + // One evaluation input for the handoff and one for the fence; the + // repeated persistFailure call found the fence already committed. + expect(await db.select().from(budgetEvaluationWork).where(eq(budgetEvaluationWork.companyId, f.companyId))).toHaveLength(2); + } finally { await fs.rm(directory, { recursive: true, force: true }); } + }); + + it("binds while back-to-back writers update the run row", async () => { + const f = await fixture({ status: "running" }); + for (let i = 0; i < 6; i++) { + const runId = await f.run(); + const directory = await spool(); + const writers = pressure(runId); + await steady(writers); + try { + await expect(createRunUsageRecorder(db, { companyId: f.companyId, runId, adapterType: "process" }, directory)).resolves.toBeTruthy(); + } finally { await writers.stop(); await fs.rm(directory, { recursive: true, force: true }); } + } + }, 60_000); + }); + + describe("exhausted retries", () => { + it("fail with a bounded, SQL-free cause that keeps the SQLSTATE", async () => { + const f = await fixture({ status: "running" }); + const runId = await f.run(); + const hooked = interleaveRunEvents(f, runId); + // A writer that still uses a repeatable-read snapshot loses every attempt. + const error = await withOperationAccountingTransaction(hooked.db, f.companyId, async tx => { + await tx.select({ id: heartbeatRuns.id }).from(heartbeatRuns).where(eq(heartbeatRuns.id, runId)).for("update"); + }).then(() => null, (failure: unknown) => failure); + expect(hooked.interleavings()).toBe(4); + expect(error).toBeInstanceOf(AccountingContentionError); + expect(isAccountingContentionError(error)).toBe(true); + const contention = error as AccountingContentionError; + expect(contention).toMatchObject({ code: "40001", attempts: 4, status: 503 }); + expect(contention.message).not.toMatch(/select|heartbeat_runs|Failed query|params/i); + // The database error stays attached for the server log only. + expect(String((contention.cause as Error).message)).toMatch(/^Failed query: select/); + }); + }); +}); diff --git a/server/src/services/usage-receipts.ts b/server/src/services/usage-receipts.ts index fdf2fe894ca..afa9af450b5 100644 --- a/server/src/services/usage-receipts.ts +++ b/server/src/services/usage-receipts.ts @@ -568,6 +568,9 @@ export async function createRunUsageRecorder(db: Db, input: { companyId: string; state?.accountingReceiptSourceId, state?.accountingReceiptSequence, state?.accountingReceiptId, ]); const pending = await readPendingRunUsageReceipts(input, directory); + // READ COMMITTED under the run-row lock, like persistUsageReceipt: every + // mutable read below follows that lock, so a run-log append committed after + // the company check cannot make the handoff's snapshot stale. const bound = await withOperationAccountingTransaction(db, input.companyId, async tx => { const [run] = await tx.select({ id: heartbeatRuns.id, agentId: heartbeatRuns.agentId, usageJson: heartbeatRuns.usageJson, costAccountedAt: heartbeatRuns.costAccountedAt }).from(heartbeatRuns) @@ -587,7 +590,7 @@ export async function createRunUsageRecorder(db: Db, input: { companyId: string; return tx.update(heartbeatRuns).set({ usageJson: { ...old, accountingReceiptSourceFences: fences, accountingReceiptSourceId: sourceId, accountingReceiptSequence: 0, accountingReceiptReady: false } }) .where(and(eq(heartbeatRuns.id, input.runId), eq(heartbeatRuns.companyId, input.companyId), isNull(heartbeatRuns.costAccountedAt))).returning({ id: heartbeatRuns.id, runtimeMode: heartbeatRuns.runtimeMode, usageJson: heartbeatRuns.usageJson }); - }); + }, { isolationLevel: "read committed" }); for (const item of pending) await fs.rm(item.file, { force: true }); if (bound.length !== 1) throw conflict("Run cannot accept a new usage recorder"); let sequence = 0; @@ -630,7 +633,7 @@ export async function createRunUsageRecorder(db: Db, input: { companyId: string; }).where(and(eq(heartbeatRuns.id, input.runId), eq(heartbeatRuns.companyId, input.companyId), isNull(heartbeatRuns.costAccountedAt), sql`${heartbeatRuns.usageJson}->>'accountingReceiptSourceId' = ${sourceId}`)).returning({ id: heartbeatRuns.id }); return rows.length === 1; - }); + }, { isolationLevel: "read committed" }); // Run-row lock first, as above. if (persisted) { finishFailedCapture?.(); finishFailedCapture = undefined; } } function enqueue(action: () => Promise) {