From 0b20ea6fe927b7bfa7bf461dcb87d5a95ead5f92 Mon Sep 17 00:00:00 2001 From: Dotta <34892728+cryppadotta@users.noreply.github.com> Date: Tue, 10 Nov 2026 03:47:48 -0600 Subject: [PATCH] fix(accounting): keep budget admission and usage handoff alive under run-row updates (#20484) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Thinking Path > - Paperclip is the open source app people use to manage AI agents for work. > - Before a run starts provider work, Paperclip reserves its budget and binds a usage recorder. Both write the run's `heartbeat_runs` row. > - Run-log appends, output progress and lifecycle events update that same row while the run prepares and streams. > - #20385 fixed this race for the standalone usage-receipt write only. Budget admission and the recorder handoff still take a repeatable-read snapshot before they lock the run row. > - Each concurrent row update makes that snapshot stale. Four retries run out with SQLSTATE `40001`, and the run ends `adapter_failed` with raw SQL as its error. > - This pull request applies the #20385 pattern (READ COMMITTED under the existing locks) to those writers and to the final dispatch gate. It also gives exhausted retries a bounded cause. > - The benefit is that busy runs start and account normally, and a remaining failure names its cause. ## Linked Issues or Issue Description Refs #20385. Builds on #20483 (bounded failure cause and `withoutSqlStatementText`). **What happened?** The restart-chaos L1 workload with `--db-pressure 4` (harness writers that update each running run row, one row per statement, like a fast stream) failed on master. In one wave, 0 of 4 tasks passed. Every first run failed in `reserveRunBudget` before spawn. In an earlier wave, a run failed in `createRunUsageRecorder`. The stack in both cases is `could not serialize access due to concurrent update` (`40001`) on the run-row `FOR UPDATE`. The board `runError` showed the full `select ... from heartbeat_runs ... for update` statement. The server log line `heartbeat execution failed` dropped the database cause. Realism: the harness writers are synthetic, but the product has the same writers on the same row. `appendHeartbeatRunEvent` increments `next_event_seq`, output progress updates `last_output_*`, and lifecycle events land while a run prepares. On a hosted database the window between snapshot and row lock is wider than locally, so a busy run reaches the same race without the harness. **Expected behavior** Budget admission and the recorder handoff commit while other writers update the run row. Budget invariants stay intact: one reservation per run, a reached hard stop still refuses work, and committed holds count against later admissions. Usage receipts are journaled once. If retries still run out, the run shows a bounded cause without SQL, and the server log records the SQLSTATE. **Steps to reproduce** Run `server/src/services/run-accounting-contention.test.ts` against the parent commit. The deterministic cases commit a real `appendHeartbeatRunEvent` after the writer's transaction checks the company and before it locks the run. Six of the eight cases fail with `40001`. ## What Changed - `reserveRunBudget` runs at READ COMMITTED. All mutable reads follow the shared admission fences, the agent `FOR SHARE` lock and the run-row `FOR UPDATE`. Exposure stays one SQL statement (hold-to-ledger atomic). The `budget_reservations_run_idx` unique index and the run lock still prevent a second hold. - The final dispatch gate's accounting phase (`dispatchResolvedInteractionIfCurrent` through `withIssueThenRunLocks` in `run-dispatch/adapters/postgres.ts`) runs at READ COMMITTED. Its ownership phase already read the same staleness facts at READ COMMITTED under the same issue and run locks. Retry promotion and stale-run cancellation keep their repeatable-read snapshot. The first L1 proof wave found this writer exhausting `40001` once the other three were fixed. - `createRunUsageRecorder` runs its receipt-source handoff and its capture-failure fence at READ COMMITTED. Both lock the run row first, the same as `persistUsageReceipt` after #20385. - `withOperationAccountingTransaction` raises `AccountingContentionError` when retries run out. It has a bounded message, `code` keeps the SQLSTATE (`40001`/`40P01`), `cause` keeps the database error, and a warning logs both. - `failRun` shows that bounded message. The code stays `adapter_failed`, so chat and recovery retry paths still apply. #20483's `failureCause` records `40001`. All other SQL text handling comes from #20483 (`withoutSqlStatementText`), unchanged here. - `doc/DATABASE.md` documents the isolation choice and the failure contract. ### Inventory of REPEATABLE READ / SERIALIZABLE transactions that write run-scoped rows | Writer | Status | Notes | | --- | --- | --- | | `persistUsageReceipt` (`usage-receipts.ts`) | Fixed in #20385 | READ COMMITTED under the run lock. | | `reserveRunBudget` (`budget-reservations.ts`) | Fixed here | Seen in waves (`budget-reservations.ts:90`). | | `createRunUsageRecorder` handoff (`usage-receipts.ts`) | Fixed here | Seen in waves (`usage-receipts.ts:572`). | | `createRunUsageRecorder` capture-failure fence (`persistFailure`) | Fixed here | Same shape: run lock first, runs mid-stream. | | Native finalizer `claimCoordinator` (`native-run-finalizer.ts`) | Safe | A default READ COMMITTED transaction with `FOR UPDATE` on the coordinator row. It cannot raise `40001`. The wave errors there were `native_finalization_missing`, a symptom of the failed reservation (no coordinator existed yet). | | Final dispatch gate, accounting phase (`dispatchResolvedInteractionIfCurrent`) | Fixed here | Seen in proof wave `w-20261110T075857-s20386` (runs `140ddb10`, `c028c6c4`, `fea0ee5a`). Its ownership phase already used READ COMMITTED under the same locks with the same staleness reads. | | Retry promotion and stale-run cancellation (`promoteOrCancelDueRetry`, `cancelStaleQueuedRun`, same lock helper) | Follow-up | Kept at repeatable read: `loadGateFacts` combines wakes, comments, dependencies, tree holds and repair state in one decision. They act on queued or scheduled runs, not streaming ones. Not seen failing in waves. If retries run out, they now fail with the bounded cause. | | Run settlement `accountRunCost` (`run-cost-accounting.ts`) | Follow-up | Seen with `40001` in a wave log. Leaves the run pending; `reconcileRunCosts` retries it durably. Multi-aggregate (coordinator, receipts, images, closure). | | Tool cost events (`paperclip-runner-tool-authority.ts`, `costs.ts`) | Safe | Insert `cost_events` only; the run row is referenced, not locked or updated. | | Budget enforcement and evaluation (`budgets.ts`), billing reconciliation, accounting-integrity repair | Safe | Lock policy, invoice or company rows, not a streaming run row. Run reads are plain selects. | | Browser use, decisions, prompt-injection, fast responses | Safe | Write their own invocation/reservation rows, not `heartbeat_runs`. | | Read-only RR snapshots (`withAccountingReadSnapshot`, idle-sleep, tool gateway, customer success, `routes/decisions.ts`) | Safe | Read only, or no run-scoped writes. | ## Verification - Red on the parent commit, green on this branch: `server/src/services/run-accounting-contention.test.ts` (real embedded PostgreSQL). - `reserveRunBudget > reserves once when run events commit between its company check and its run lock` (red: `40001`) - `reserveRunBudget > reserves every run while back-to-back writers update the run row` (red: `40001`) - `reserveRunBudget > counts every committed hold, so successive admissions never exceed the budget` (red) - `reserveRunBudget > still refuses a run once the hard stop is reached` (red: `40001` instead of `409`) - `createRunUsageRecorder > binds its receipt source when run events commit between its company check and its run lock` (red: `40001`; green asserts the spooled receipt is journaled exactly once) - `createRunUsageRecorder > binds while back-to-back writers update the run row` (red: `40001`) - `reserveRunBudget > never creates a second reservation for one run when two dispatches race` (invariant; passes before and after) - `final dispatch gate > admits and hands off a reserved run when run events commit between its company check and its run lock` (red: `AccountingContentionError` after 4 × `40001`) - `createRunUsageRecorder > persists its capture-failure fence once when run events commit between its company check and its run lock` - `final dispatch gate > admits while back-to-back writers update the run row` (red: `40001`) - `exhausted retries > fail with a bounded, SQL-free cause that keeps the SQLSTATE` (new contract) - `run-completion.test.ts > saves a bounded cause and keeps the retryable generic code when accounting retries are exhausted` (saved code, message and `failureCause`, and log fields, through `failRun`) - Related suites pass: `usage-receipts-read`, `admission-independent-of-cleanup`, `budget-live-attempts`, `idle-accounting-work`, `budget-enforcement-scoping`, `closed-native-receipt-accounting`, `budget-evaluation-work`, `run-dispatch/adapters/postgres`, `heartbeat/run-completion` (all 126 run-completion cases pass; its `afterAll` embedded-PostgreSQL teardown exceeded 30 s on this loaded host, the known local flake). - `pnpm --filter @paperclipai/server typecheck` passes. - Latency of `reserveRunBudget` on the normal path (no contention, embedded PostgreSQL, 350 measured reservations per run, 3 runs each): before p50 14.7 / 15.0 / 18.6 ms, p95 25.6 / 31.4 / 39.4 ms; after p50 14.2 / 14.0 / 16.6 ms, p95 29.7 / 24.0 / 29.3 ms. No added latency: the statements are the same, only the isolation level changes. Under contention it removes the retry backoff. - Live restart-chaos L1 waves (`--seed 20386 --db-pressure 4`, codex and claude, K1): - Master before this change: `w-20261110T050818-s20386`, 0/4. Every first run failed in `reserveRunBudget` with `40001`. - Head `90530789b0` (first three writers fixed): `w-20261110T075857-s20386`, 2/4. Both codex tasks passed. Both claude failures (t02, t03) trace to the dispatch gate, now fixed. Those runs already showed the bounded cause, with no SQL on the board. The failure was logged with `sqlState` 40001. No `reserveRunBudget` or recorder `40001` remained. - Head `d5ed21926f` (dispatch gate fixed): `w-20261110T090531-s20386`. All 4 tasks reached done with no retry and R1–R6 passed. There were no failed runs and zero `40001` in the server log, under 310k pressure updates. Each task still failed the T-oracle `company_spend_settled` with `pending_runs=1`. That pending run is t02's continuation run `601ae80a`, cancelled before provider work as `issue_terminal_status`, whose accounting debt was not yet settled when graded about 15 s later. It is not a contention failure. Caveat: I edited the worktree during this wave (`worktree_changed_mid_wave`), so a clean wave on the final head follows. - Final head `9f1c480717`, clean worktree: `w-20261110T094417-s20386`. All 4 tasks reached done with no retry, and R1–R6 passed for each. There were no failed runs and zero serialization failures in the server log, under about 523k pressure updates. Each task still fails only the T-oracle `company_spend_settled` (`pending_runs=1`). It is the same cancelled pre-provider continuation run (`issue_terminal_status`), whose accounting debt was still unsettled at grading. That is a separate settlement-timing follow-up, not database contention. Master was 0/4 on the same seed. ## Risks - READ COMMITTED gives each statement its own snapshot. Admission reads policy, scope and exposure in separate statements. The shared company/project fences and the agent share lock exclude concurrent pause and policy edits, and exposure stays one statement, so each read is at least as fresh as the old snapshot. - Concurrent admissions of different runs stay best-effort, as `doc/SPEC-implementation.md` already states ("concurrent starts can exceed the same observed balance"). This change does not widen that window: each statement snapshot is newer than the old transaction snapshot. It does not add a company mutex. - A run that fails after exhausted retries now records `accounting_contention` instead of `adapter_failed`. Retry rules that key on `adapter_failed` spawn-like messages did not match this failure before, so no automatic retry changes. - No schema, migration or persisted-format change. ## Model Used - Claude Opus 5.5 (Anthropic), model ID `claude-opus-5-5`, used through Claude Code with shell, file-edit and code-execution tools. ## Checklist - [x] I have included a thinking path that traces from project context to this change - [x] I have specified the model used (with version and capability details) - [x] I have checked ROADMAP.md and confirmed this PR does not duplicate planned core work - [x] I have searched GitHub for duplicate or related PRs and linked them above - [x] I have either (a) linked existing issues with `Fixes: #` / `Closes #` / `Refs #` OR (b) described the issue in-PR following the relevant issue template - [x] I have not referenced internal/instance-local Paperclip issues or links (only public GitHub `#NNN` / `github.com/paperclipai/paperclip` URLs) - [x] My branch name describes the change (e.g. `docs/...`, `fix/...`) and contains no internal Paperclip ticket id or instance-derived details - [x] I have run tests locally and they pass - [x] I have added or updated tests where applicable - [x] I have updated relevant documentation to reflect my changes - [x] I have considered and documented any risks above - [x] All Paperclip CI gates are green - [x] Greptile is 5/5 with no open P2s, recommendations, or follow-ups - [x] I will address all Greptile and reviewer comments before requesting merge 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 5.5 --- doc/DATABASE.md | 17 + .../modules/run-dispatch/adapters/postgres.ts | 16 +- server/src/services/accounting-transaction.ts | 45 ++- server/src/services/budget-reservations.ts | 12 +- .../services/heartbeat/run-completion.test.ts | 19 ++ .../src/services/heartbeat/run-completion.ts | 7 +- .../run-accounting-contention.test.ts | 303 ++++++++++++++++++ server/src/services/usage-receipts.ts | 7 +- 8 files changed, 414 insertions(+), 12 deletions(-) create mode 100644 server/src/services/run-accounting-contention.test.ts 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) {