From 0116e459583203fe616f4c63b9155008697fb7e3 Mon Sep 17 00:00:00 2001 From: Dotta <34892728+cryppadotta@users.noreply.github.com> Date: Mon, 9 Nov 2026 18:14:52 -0600 Subject: [PATCH] perf(db): index pending budget Stops instead of scanning all run results (#20391) 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 > - Budget enforcement stops agent runs when a budget hard limit is reached. The stop intent is saved in each run's `result_json` until run control delivers it. > - The budget worker, run control, and the idle-sleep check find these pending stops with a predicate on `result_json`. No index supports that predicate. > - Each check therefore reads every heartbeat run and its large result JSON from TOAST. On a hosted database, this statement family was among the most expensive: about 8,000 calls at about 1 s each, about 9,800 s in total. > - Pending stops are rare. The predicate is constant, so a partial index can contain only those rows. > - This pull request adds that partial index. The SQL that the server sends does not change. > - The benefit is that each check reads a few index entries instead of the full table. The database has more capacity for request work. ## Linked Issues or Issue Description **What existing behavior does this improve?** The check for committed budget Stops that are not delivered yet. The budget worker runs it each time it looks for work (`hasPendingBudgetEnforcement`), and again for each delivery pass (`listPendingBudgetCancellationScopes`). Run control and the idle-sleep check use the same predicate. **Subsystem affected** Budget enforcement in the server (`server/src/services/heartbeat/budget-cancellation.ts`, `server/src/services/budgets.ts`, `server/src/services/heartbeat/run-control.ts`, `server/src/services/idle-sleep-safety.ts`) and the `heartbeat_runs` schema in `packages/db`. **Current behavior** The query groups all rows of `heartbeat_runs` and reads `result_json->'startupCancellation'->'budgetEnforcement'` for each row. A limit of 1 does not help, because the group and sort occur before the limit. Cost grows with the run history, not with the number of pending stops. **Proposed behavior** A partial index, `heartbeat_runs_pending_budget_cancellation_idx`, contains only runs with a pending budget Stop. The planner uses it for all four callers. The results are the same. **Reason and benefit** Local SQL measurements show a 300x to 1,000x decrease in time and buffer reads (see Verification). **Breaking changes** None. ## What Changed - Added the partial index `heartbeat_runs_pending_budget_cancellation_idx` on `heartbeat_runs (company_id)`. Its predicate is identical to `pendingBudgetRunCancellationPredicate`. - Added migration `0455_pending_budget_cancellation_index.sql`. It uses `CREATE INDEX CONCURRENTLY IF NOT EXISTS`, and it is in the list of migrations that run outside a transaction, as `0436` is. - Moved the scoped run-control read and the idle-sleep check into exported builders in `budget-cancellation.ts`, next to the scope list query. The SQL does not change. The callers and the test now share the same statements. - Added a real PostgreSQL regression test. It checks that delivered, malformed, and missing intents are not pending. It also EXPLAINs the statements that the callers send (scope list with limit 1 and 100, scoped read, `EXISTS` check) and requires the new index in each plan. If a person changes a query shape, the server predicate, or the index predicate, the test fails. - Added `scripts/benchmark-budget-cancellation-scan.ts`. It captures the SQL that the real service sends, then compares the plans, timings, and complete rows with and without the index. - Updated `doc/DATABASE.md`. ## Verification Red/green regression test: - Before the migration, the new test failed: the plan was `Seq Scan` with the predicate as a filter. - After the migration: `npx vitest run src/services/heartbeat/budget-cancellation.test.ts` passes (17 tests). - Also passed: `src/__tests__/budgets-service.test.ts`, `src/__tests__/idle-sleep-diagnostics-postgres.test.ts`, `src/services/budget-resume-continuation.test.ts`. - `pnpm --filter @paperclipai/db check:migrations`, and typecheck of `@paperclipai/db` and `@paperclipai/server` pass. Performance and result equality (local SQL measurements only): ```sh pnpm --filter @paperclipai/server exec tsx ../scripts/benchmark-budget-cancellation-scan.ts --runs=20000 --payload-kib=16 [--reverse] ``` - Fixture: embedded PostgreSQL 18.1. 20,241 runs in 2 companies, with incompressible result JSON of about 16 KiB (table 379 MB). 40 pending stops in 6 scopes, 200 delivered or malformed decoys, and 1 cross-company intent. - The SQL is the same in both treatments. The script toggles only the new index. 2 warmups and 7 samples per query. Rows are compared outside the timed part. Ordered queries must match in order. The scoped run-control read has no `ORDER BY`, so it is compared as a multiset. | Query (caller) | Baseline median | With index median | Shared buffers | | --- | --- | --- | --- | | `listPendingBudgetCancellationScopes(db, 1)` (worker `hasPending`) | 688 ms | 2.1 ms | 243,500 → 1,158 | | `listPendingBudgetCancellationScopes(db)` (delivery) | 799 ms | 1.4 ms | 243,500 → 1,158 | | Scoped pending runs (run control) | 357 ms | 1.1 ms | 121,835 → 322 | | `EXISTS` pending (idle-sleep check) | 440 ms | 0.37 ms | 240,560 → 2 | - Reverse order (index first) gave the same result: 524 ms → 1.6 ms, 464 ms → 1.4 ms, 306 ms → 1.0 ms, 306 ms → 0.42 ms. - Plans change from `Seq Scan` (or a broad bitmap scan) to `Index Scan` or `Index Only Scan` on the new index. - All result rows are equal in every run. Migration duration: - 20,241 runs (379 MB): index build 2.0 s to 3.9 s. - 200,241 runs with 8 KiB result JSON (2.1 GB): index build 5.6 s. With this table size, the four queries change from 1,148 / 879 / 514 / 4,392 ms to 1.1 / 1.2 / 0.85 / 8.3 ms. - The build is concurrent, so it does not block run writes. ## Risks - Migration: the build reads every run's result JSON one time. It is concurrent and resumable (`IF NOT EXISTS`). The executor repairs an invalid concurrent index on retry. - Write cost: run writes evaluate the index predicate. The index contains very few rows, so the extra cost is small. - Predicate drift: if a later change edits only one of the two predicates, the planner cannot use the index. The new regression test detects this. - The hosted timings are not measured yet. The local numbers come from a synthetic fixture. ## Model Used - Claude Opus 5.5 (`claude-opus-5-5`), Anthropic, used through Claude Code with tool use (shell, file edits, local PostgreSQL). ## 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 - [ ] All Paperclip CI gates are green - [ ] 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 | 15 ++ packages/db/src/client.ts | 1 + ...0455_pending_budget_cancellation_index.sql | 2 + ...{0450_snapshot.json => 0455_snapshot.json} | 215 +++++++++++++++++- packages/db/src/migrations/meta/_journal.json | 7 + packages/db/src/schema/heartbeat_runs.ts | 9 + scripts/benchmark-budget-cancellation-scan.ts | 154 +++++++++++++ .../heartbeat/budget-cancellation.test.ts | 44 +++- .../services/heartbeat/budget-cancellation.ts | 22 +- server/src/services/heartbeat/run-control.ts | 6 +- server/src/services/idle-sleep-safety.ts | 4 +- 11 files changed, 466 insertions(+), 13 deletions(-) create mode 100644 packages/db/src/migrations/0455_pending_budget_cancellation_index.sql rename packages/db/src/migrations/meta/{0450_snapshot.json => 0455_snapshot.json} (99%) create mode 100644 scripts/benchmark-budget-cancellation-scan.ts diff --git a/doc/DATABASE.md b/doc/DATABASE.md index 6ee2badb769..fb0e9198acc 100644 --- a/doc/DATABASE.md +++ b/doc/DATABASE.md @@ -364,6 +364,21 @@ through the partial index `heartbeat_runs_tree_hold_suppression_idx`. The index covers only cancelled runs with a `treeHoldSuppression` marker, keyed by company and hold. Migration 0436 builds it concurrently. +The budget worker finds committed budget Stops that are not delivered yet +through the partial index `heartbeat_runs_pending_budget_cancellation_idx`. The +index covers only runs whose `result_json.startupCancellation.budgetEnforcement` +is an object without `deliveredAt`, keyed by company. Its predicate must stay +identical to `pendingBudgetRunCancellationPredicate` in +`server/src/services/heartbeat/budget-cancellation.ts`; a regression test checks +that the planner uses it. The worker check, run control, and the idle-sleep +check therefore read a few index entries, not every run's result JSON. The +migration builds the index concurrently. To compare both plans on a disposable +local fixture: + +```sh +pnpm --filter @paperclipai/server exec tsx ../scripts/benchmark-budget-cancellation-scan.ts --runs=20000 --payload-kib=16 +``` + The chat wake-notice sweep compares UUID columns directly against guarded, canonical lowercase UUID payload values. Malformed and noncanonical payloads still match nothing. The query can use existing primary keys without adding a diff --git a/packages/db/src/client.ts b/packages/db/src/client.ts index 2633c213f01..10743acd93c 100644 --- a/packages/db/src/client.ts +++ b/packages/db/src/client.ts @@ -617,6 +617,7 @@ const MIGRATIONS_WITH_TOP_LEVEL_STATEMENTS = new Set([ "0450_petite_frightful_four.sql", "0453_curious_ben_parker.sql", "0454_lumpy_ronan.sql", + "0455_pending_budget_cancellation_index.sql", "0384_funny_adam_warlock.sql", diff --git a/packages/db/src/migrations/0455_pending_budget_cancellation_index.sql b/packages/db/src/migrations/0455_pending_budget_cancellation_index.sql new file mode 100644 index 00000000000..fd5798a14c5 --- /dev/null +++ b/packages/db/src/migrations/0455_pending_budget_cancellation_index.sql @@ -0,0 +1,2 @@ +CREATE INDEX CONCURRENTLY IF NOT EXISTS "heartbeat_runs_pending_budget_cancellation_idx" ON "heartbeat_runs" USING btree ("company_id") WHERE jsonb_typeof("heartbeat_runs"."result_json"->'startupCancellation'->'budgetEnforcement') = 'object' + and "heartbeat_runs"."result_json"->'startupCancellation'->'budgetEnforcement'->>'deliveredAt' is null; \ No newline at end of file diff --git a/packages/db/src/migrations/meta/0450_snapshot.json b/packages/db/src/migrations/meta/0455_snapshot.json similarity index 99% rename from packages/db/src/migrations/meta/0450_snapshot.json rename to packages/db/src/migrations/meta/0455_snapshot.json index af424af109a..bc7b1347a0d 100644 --- a/packages/db/src/migrations/meta/0450_snapshot.json +++ b/packages/db/src/migrations/meta/0455_snapshot.json @@ -1,6 +1,6 @@ { - "id": "b8f43f1a-3230-4cda-991d-0e7b6323114b", - "prevId": "4bd9d289-9cb8-4262-9dd7-04ce2b6020c4", + "id": "2be53987-902a-41b6-8cb4-6e89a768f8cf", + "prevId": "f98af245-4a2e-4c3d-b89e-6ed3e81a074c", "version": "7", "dialect": "postgresql", "tables": { @@ -7882,6 +7882,150 @@ }, "isRLSEnabled": false }, + "public.budget_resume_claims": { + "name": "budget_resume_claims", + "schema": "", + "columns": { + "run_id": { + "name": "run_id", + "type": "uuid", + "primaryKey": true, + "notNull": true + }, + "company_id": { + "name": "company_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "incident_id": { + "name": "incident_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "actor_user_id": { + "name": "actor_user_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "actor_source": { + "name": "actor_source", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "state": { + "name": "state", + "type": "text", + "primaryKey": false, + "notNull": true, + "default": "'pending'" + }, + "claimed_at": { + "name": "claimed_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "settled_at": { + "name": "settled_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "reason": { + "name": "reason", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "request_id": { + "name": "request_id", + "type": "uuid", + "primaryKey": false, + "notNull": false + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": { + "budget_resume_claims_pending_idx": { + "name": "budget_resume_claims_pending_idx", + "columns": [ + { + "expression": "company_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "run_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "where": "\"budget_resume_claims\".\"state\" = 'pending'", + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": { + "budget_resume_claims_run_id_heartbeat_runs_id_fk": { + "name": "budget_resume_claims_run_id_heartbeat_runs_id_fk", + "tableFrom": "budget_resume_claims", + "tableTo": "heartbeat_runs", + "columnsFrom": [ + "run_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "cascade", + "onUpdate": "no action" + }, + "budget_resume_claims_company_id_companies_id_fk": { + "name": "budget_resume_claims_company_id_companies_id_fk", + "tableFrom": "budget_resume_claims", + "tableTo": "companies", + "columnsFrom": [ + "company_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "cascade", + "onUpdate": "no action" + }, + "budget_resume_claims_incident_id_budget_incidents_id_fk": { + "name": "budget_resume_claims_incident_id_budget_incidents_id_fk", + "tableFrom": "budget_resume_claims", + "tableTo": "budget_incidents", + "columnsFrom": [ + "incident_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "cascade", + "onUpdate": "no action" + } + }, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, "public.built_in_managed_resources": { "name": "built_in_managed_resources", "schema": "", @@ -17255,7 +17399,7 @@ "columnsTo": [ "id" ], - "onDelete": "set null", + "onDelete": "cascade", "onUpdate": "no action" }, "company_secret_proposals_secret_proposal_id_company_secret_proposals_id_fk": { @@ -34408,6 +34552,34 @@ "concurrently": false, "method": "btree", "with": {} + }, + "heartbeat_run_events_retry_exhausted_idx": { + "name": "heartbeat_run_events_retry_exhausted_idx", + "columns": [ + { + "expression": "company_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "run_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "id", + "isExpression": false, + "asc": false, + "nulls": "last" + } + ], + "isUnique": false, + "where": "\"heartbeat_run_events\".\"event_type\" = 'lifecycle' and \"heartbeat_run_events\".\"message\" like 'Bounded retry exhausted%'", + "concurrently": true, + "method": "btree", + "with": {} } }, "foreignKeys": { @@ -35249,6 +35421,22 @@ "method": "btree", "with": {} }, + "heartbeat_runs_pending_budget_cancellation_idx": { + "name": "heartbeat_runs_pending_budget_cancellation_idx", + "columns": [ + { + "expression": "company_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "where": "jsonb_typeof(\"heartbeat_runs\".\"result_json\"->'startupCancellation'->'budgetEnforcement') = 'object'\n and \"heartbeat_runs\".\"result_json\"->'startupCancellation'->'budgetEnforcement'->>'deliveredAt' is null", + "concurrently": true, + "method": "btree", + "with": {} + }, "heartbeat_runs_tree_hold_suppression_idx": { "name": "heartbeat_runs_tree_hold_suppression_idx", "columns": [ @@ -35418,6 +35606,27 @@ "method": "btree", "with": {} }, + "heartbeat_runs_company_native_workspace_idx": { + "name": "heartbeat_runs_company_native_workspace_idx", + "columns": [ + { + "expression": "company_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "(\"runner_profile_json\"->'nativeWorkspaceSync'->>'workspaceId')", + "asc": true, + "isExpression": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": true, + "method": "btree", + "with": {} + }, "heartbeat_runs_company_issue_created_idx": { "name": "heartbeat_runs_company_issue_created_idx", "columns": [ diff --git a/packages/db/src/migrations/meta/_journal.json b/packages/db/src/migrations/meta/_journal.json index 29424d35338..f6f36f4fd59 100644 --- a/packages/db/src/migrations/meta/_journal.json +++ b/packages/db/src/migrations/meta/_journal.json @@ -3172,6 +3172,13 @@ "when": 1794266236755, "tag": "0454_lumpy_ronan", "breakpoints": true + }, + { + "idx": 455, + "version": "7", + "when": 1794268118742, + "tag": "0455_pending_budget_cancellation_index", + "breakpoints": true } ] } \ No newline at end of file diff --git a/packages/db/src/schema/heartbeat_runs.ts b/packages/db/src/schema/heartbeat_runs.ts index 5b88537f32a..aaaf4fc2e74 100644 --- a/packages/db/src/schema/heartbeat_runs.ts +++ b/packages/db/src/schema/heartbeat_runs.ts @@ -152,6 +152,15 @@ export const heartbeatRuns = pgTable( ), scheduledRetryCancellationIdx: index("heartbeat_runs_scheduled_retry_cancellation_idx").on(table.id) .where(sql`${table.status} = 'scheduled_retry' and jsonb_typeof(${table.resultJson}->'startupCancellation') = 'object'`), + // Committed budget Stops that are not delivered yet. The predicate must stay + // identical to `pendingBudgetRunCancellationPredicate` in the server, so the + // budget worker, run control, and the idle check read only these few rows + // instead of every run's result JSON. + pendingBudgetCancellationIdx: index("heartbeat_runs_pending_budget_cancellation_idx") + .on(table.companyId) + .where(sql`jsonb_typeof(${table.resultJson}->'startupCancellation'->'budgetEnforcement') = 'object' + and ${table.resultJson}->'startupCancellation'->'budgetEnforcement'->>'deliveredAt' is null`) + .concurrently(), // Runs a dispatch gate refused for a pause; that hold's release resumes them. treeHoldSuppressionIdx: index("heartbeat_runs_tree_hold_suppression_idx") .on(table.companyId, sql`(${table.resultJson}->'treeHoldSuppression'->>'holdId')`) diff --git a/scripts/benchmark-budget-cancellation-scan.ts b/scripts/benchmark-budget-cancellation-scan.ts new file mode 100644 index 00000000000..69f48f0b408 --- /dev/null +++ b/scripts/benchmark-budget-cancellation-scan.ts @@ -0,0 +1,154 @@ +// Measures the committed budget Stop scans with and without the partial index +// `heartbeat_runs_pending_budget_cancellation_idx`. The SQL is the same in both +// treatments; only that index is toggled, in a disposable database. +// +// pnpm --filter @paperclipai/server exec tsx ../scripts/benchmark-budget-cancellation-scan.ts [--runs=20000] [--payload-kib=16] [--reverse] +import assert from "node:assert/strict"; +import { randomBytes, randomUUID } from "node:crypto"; +import { readFileSync, readdirSync } from "node:fs"; +import { createRequire } from "node:module"; +import { fileURLToPath } from "node:url"; +import { agents, companies, createDb, heartbeatRuns, startEmbeddedPostgresTestDatabase, type Db } from "../packages/db/src/index.js"; +import { + listPendingBudgetCancellationScopes, + pendingBudgetRunCancellationExistsQuery, + pendingBudgetRunIdsForScopeQuery, +} from "../server/src/services/heartbeat/budget-cancellation.js"; + +const { sql } = createRequire(new URL("../server/package.json", import.meta.url))("drizzle-orm") as typeof import("drizzle-orm"); +const arg = (name: string, fallback: number) => Number(process.argv.find(value => value.startsWith(`--${name}=`))?.split("=")[1] ?? fallback); +const runCount = arg("runs", 20_000); +const payloadKib = arg("payload-kib", 16); +const samples = arg("samples", 7); +const indexName = "heartbeat_runs_pending_budget_cancellation_idx"; +const migrationDir = fileURLToPath(new URL("../packages/db/src/migrations/", import.meta.url)); +const migrationFile = readdirSync(migrationDir).find(file => file.endsWith(".sql") + && readFileSync(`${migrationDir}${file}`, "utf8").includes(`"${indexName}"`)); +assert.ok(migrationFile, "The index migration was not found"); +const createIndexSql = readFileSync(`${migrationDir}${migrationFile}`, "utf8").trim(); + +type Query = { sql: string; params: unknown[] }; +const keepAlive = setInterval(() => {}, 1000); +let database: Awaited> | undefined; +let db: Db | undefined; +try { + database = await startEmbeddedPostgresTestDatabase("budget-cancellation-scan-benchmark-"); + db = createDb(database.connectionString); + const [version] = await db.$client.unsafe("select version()"); + + // Two tenants. Most runs carry an incompressible result JSON, so a scan must + // read TOAST. A few runs carry committed Stops: pending, delivered, decoys. + const companyIds = [randomUUID(), randomUUID()]; + const agentIds = [randomUUID(), randomUUID(), randomUUID()]; + await db.insert(companies).values(companyIds.map((id, n) => ({ id, name: `Budget scan ${n}`, issuePrefix: `BS${n}` }))); + await db.insert(agents).values(agentIds.map((id, n) => ({ id, companyId: companyIds[n % 2]!, name: `Agent ${n}`, adapterType: "process" }))); + const payload = randomBytes(payloadKib * 768).toString("base64"); + await db.$client.unsafe(`insert into heartbeat_runs (company_id, agent_id, status, runtime_mode, created_at, updated_at, result_json) + select case when n % 2 = 0 then $1::uuid else $2::uuid end, case when n % 2 = 0 then $3::uuid else $4::uuid end, + (array['succeeded','failed','cancelled','timed_out'])[1 + n % 4], 'legacy', + timestamptz '2026-10-01' + n * interval '1 second', timestamptz '2026-10-01' + n * interval '1 second', + jsonb_build_object('summary', $5::text, 'stdoutExcerpt', left($5::text, 512)) + from generate_series(1, $6::int) n`, [companyIds[0], companyIds[1], agentIds[0], agentIds[1], payload, runCount]); + const policyId = randomUUID(); + const intent = (companyId: string, scopeType: string, scopeId: string, extra: Record = {}) => + ({ companyId, scopeType, scopeId, policyId, version: 3, ...extra }); + const special: Array<{ companyId: string; agentId: string; status: string; updatedAt: Date; resultJson: Record }> = []; + for (let n = 0; n < 40; n++) { + const companyId = companyIds[n % 2]!, agentId = n % 2 === 0 ? agentIds[0]! : agentIds[1]!; + const scopes: Array<[string, string]> = [["company", companyId], ["agent", agentId], ["agent", agentIds[2]!]]; + const [scopeType, scopeId] = scopes[n % 3]!; + const extra = n % 5 === 0 ? { lastAttemptAt: new Date(Date.UTC(2026, 10, 1, 0, n)).toISOString() } : {}; + special.push({ companyId, agentId, status: n % 4 === 0 ? "cancelled" : "running", updatedAt: new Date(Date.UTC(2026, 10, 2, 0, n)), + resultJson: { summary: payload, startupCancellation: { + ...(n % 7 === 0 ? {} : { requestedAt: new Date(Date.UTC(2026, 10, 1, 1, 40 - n)).toISOString() }), + budgetEnforcement: intent(companyId, scopeType, scopeId, extra) } } }); + } + // Delivered Stops and malformed intents are not pending. + for (let n = 0; n < 200; n++) special.push({ companyId: companyIds[n % 2]!, agentId: agentIds[n % 2]!, status: "cancelled", + updatedAt: new Date(Date.UTC(2026, 10, 3, 0, 0, n)), resultJson: { summary: payload, startupCancellation: { + budgetEnforcement: n % 10 === 0 ? "legacy-string" : intent(companyIds[n % 2]!, "company", companyIds[n % 2]!, + { deliveredAt: new Date(Date.UTC(2026, 10, 3)).toISOString() }) } } }); + // An intent whose company does not match its run is returned by SQL and + // filtered by the service; it must stay visible to both treatments. + special.push({ companyId: companyIds[0]!, agentId: agentIds[0]!, status: "running", updatedAt: new Date(Date.UTC(2026, 10, 4)), + resultJson: { startupCancellation: { budgetEnforcement: intent(companyIds[1]!, "company", companyIds[1]!) } } }); + await db.insert(heartbeatRuns).values(special.map(row => ({ ...row, runtimeMode: "legacy" }))); + await db.$client.unsafe("vacuum analyze heartbeat_runs"); + const [{ total_bytes: tableBytes }] = await db.$client.unsafe(`select pg_total_relation_size('heartbeat_runs')::bigint as total_bytes`); + + // Capture the SQL that the real service and its callers send. + const captured: Query[] = []; + let capturing = true; + const session = Reflect.get(db, "session"); + const logger = { logQuery(query: string, params: unknown[]) { if (capturing) captured.push({ sql: query, params }); } }; + session.logger = logger; session.options.logger = logger; + const queries: Array<{ name: string; query: Query }> = []; + const captureOne = async (name: string, invoke: () => Promise) => { + captured.length = 0; await invoke(); + assert.equal(captured.length, 1, `${name} sent ${captured.length} statements`); + queries.push({ name, query: captured[0]! }); + }; + await captureOne("hasPendingBudgetEnforcement(limit 1)", () => listPendingBudgetCancellationScopes(db!, 1)); + await captureOne("deliverBudgetEnforcement(limit 100)", () => listPendingBudgetCancellationScopes(db!)); + await captureOne("runControl.cancelBudgetScopeWork", async () => pendingBudgetRunIdsForScopeQuery(db!, + { companyId: companyIds[0]!, scopeType: "agent", scopeId: agentIds[2]! })); + await captureOne("idleSleep.budget_cancellation", () => db!.execute(sql`select exists (${pendingBudgetRunCancellationExistsQuery()}) as pending`)); + capturing = false; + + async function measure(query: Query) { + const timings: number[] = []; let rows: unknown[] = []; + for (let n = 0; n < samples + 2; n++) { + const start = performance.now(); + rows = await db!.$client.unsafe(query.sql, query.params as never[]).values(); + if (n > 1) timings.push(performance.now() - start); + } + const [plan] = await db!.$client.unsafe(`EXPLAIN (ANALYZE, BUFFERS, FORMAT JSON) ${query.sql}`, query.params as never[]); + const root = plan["QUERY PLAN"][0]; + const nodes: string[] = []; + (function walk(node: Record) { nodes.push(String(node["Node Type"]) + (node["Index Name"] ? `(${node["Index Name"]})` : "")); + for (const child of (node.Plans as Record[] | undefined) ?? []) walk(child); })(root.Plan); + timings.sort((a, b) => a - b); + return { medianMs: timings[Math.floor(timings.length / 2)]!, minMs: timings[0]!, maxMs: timings.at(-1)!, rows, + sharedBuffers: Number(root.Plan["Shared Hit Blocks"] ?? 0) + Number(root.Plan["Shared Read Blocks"] ?? 0), nodes }; + } + const indexed = async (present: boolean) => { + const existing = await db!.$client.unsafe(`select 1 from pg_indexes where indexname = $1`, [indexName]); + if (existing.length) await db!.$client.unsafe(`drop index "${indexName}"`); + if (present) { + const start = performance.now(); + await db!.$client.unsafe(createIndexSql); + await db!.$client.unsafe("analyze heartbeat_runs"); + return performance.now() - start; + } + return null; + }; + const order: Array<"baseline" | "treatment"> = process.argv.includes("--reverse") ? ["treatment", "baseline"] : ["baseline", "treatment"]; + const results: Record>>> = {}; + let indexBuildMs: number | null = null; + for (const variant of order) { + const built = await indexed(variant === "treatment"); + if (built !== null) indexBuildMs = built; + for (const { name, query } of queries) (results[name] ??= {})[variant] = await measure(query); + } + const report = Object.entries(results).map(([name, { baseline, treatment }]) => { + // Run control's scoped read has no ORDER BY; its caller treats the rows as a + // set. Every other read promises its order, so compare it exactly. + const comparable = (rows: unknown[]) => name === "runControl.cancelBudgetScopeWork" + ? rows.map(row => JSON.stringify(row)).sort() : rows; + assert.deepEqual(comparable(treatment!.rows), comparable(baseline!.rows), `${name} returned different rows`); + if (name !== "idleSleep.budget_cancellation") assert.ok(treatment!.nodes.some(node => node.includes(indexName)), `${name} did not use the index`); + return { name, equalRows: baseline!.rows.length, baselineMs: baseline!.medianMs, baselineRangeMs: [baseline!.minMs, baseline!.maxMs], + treatmentMs: treatment!.medianMs, treatmentRangeMs: [treatment!.minMs, treatment!.maxMs], + speedup: Number((baseline!.medianMs / treatment!.medianMs).toFixed(1)), + baselineBuffers: baseline!.sharedBuffers, treatmentBuffers: treatment!.sharedBuffers, + baselinePlan: baseline!.nodes.join(" > "), treatmentPlan: treatment!.nodes.join(" > ") }; + }); + // The service's own post-processing must agree too. + const scopes = await listPendingBudgetCancellationScopes(db); + assert.ok(scopes.length > 0); + console.log(JSON.stringify({ postgres: version.version, runs: runCount + special.length, payloadKib, tableBytes: Number(tableBytes), + order, samples, indexBuildMs, createIndexSql, servicesScopes: scopes.length, report }, null, 2)); +} finally { + try { await db?.$client.end({ timeout: 0 }); await database?.cleanup(); } + finally { clearInterval(keepAlive); } +} diff --git a/server/src/services/heartbeat/budget-cancellation.test.ts b/server/src/services/heartbeat/budget-cancellation.test.ts index 30d0b23dcec..07d68c14bb3 100644 --- a/server/src/services/heartbeat/budget-cancellation.test.ts +++ b/server/src/services/heartbeat/budget-cancellation.test.ts @@ -3,7 +3,7 @@ import { spawn } from "node:child_process"; import { once } from "node:events"; import type { Db } from "@paperclipai/db"; import { agents, agentWakeupRequests, browserUseRuns, browserUseSessions, budgetPolicies, companies, connectionGrants, costEvents, environmentLeases, heartbeatRuns, issues, projects, toolApplications, toolConnections } from "@paperclipai/db"; -import { eq, inArray } from "drizzle-orm"; +import { eq, inArray, sql } from "drizzle-orm"; import { describe, expect, it, vi } from "vitest"; import { runningProcesses } from "../../adapters/index.js"; import { terminateLocalService } from "../local-service-supervisor.js"; @@ -14,7 +14,13 @@ import { DELIVERY_QUEUES, subscribeDeliveryWork } from "../delivery-work-notific import { PROCESS_START_REQUESTED } from "../native-local-process-stop.js"; import { createPostgresRunDispatchAdapter } from "../../modules/run-dispatch/adapters/postgres.js"; import { createHeartbeatRunControl } from "./run-control.js"; -import { listPendingBudgetCancellationScopes } from "./budget-cancellation.js"; +import { + listPendingBudgetCancellationScopes, + pendingBudgetCancellationScopesQuery, + pendingBudgetRunCancellationExistsQuery, + pendingBudgetRunCancellationPredicate, + pendingBudgetRunIdsForScopeQuery, +} from "./budget-cancellation.js"; import { callbacks, runControlPostgresSupport, setupRunControlDatabase } from "./__tests__/run-control-fixtures.js"; vi.mock("../live-events.js", () => ({ publishLiveEvent: vi.fn() })); @@ -466,4 +472,38 @@ describe.skipIf(!runControlPostgresSupport.supported)("scoped budget cancellatio expect(deps.setRunStatusFromLive).not.toHaveBeenCalled(); } finally { unsubscribe(); } }); + + it("reads pending Stops through the partial index and ignores delivered or malformed intents", async () => { + const { company, agent } = await fixture(); + const intent = { companyId: company.id, scopeType: "agent" as const, scopeId: agent.id, policyId: randomUUID(), version: 1 }; + const [pending] = await db.insert(heartbeatRuns).values([ + { resultJson: { startupCancellation: { budgetEnforcement: intent } } }, + { resultJson: { startupCancellation: { budgetEnforcement: { ...intent, deliveredAt: new Date().toISOString() } } } }, + { resultJson: { startupCancellation: { budgetEnforcement: "legacy" } } }, + { resultJson: { startupCancellation: { requestedAt: new Date().toISOString() } } }, + { resultJson: null }, + ].map(row => ({ ...row, companyId: company.id, agentId: agent.id, status: "cancelled", runtimeMode: "legacy" as const }))).returning(); + expect(await db.select({ id: heartbeatRuns.id }).from(heartbeatRuns).where(pendingBudgetRunCancellationPredicate())) + .toEqual([{ id: pending.id }]); + expect(await listPendingBudgetCancellationScopes(db)).toEqual([{ companyId: company.id, scopeType: "agent", scopeId: agent.id, + enforcement: { policyId: intent.policyId, version: 1 } }]); + // The schema index predicate must stay identical to the server predicate. + // Otherwise every worker pass falls back to reading all run result JSON. + // EXPLAIN the statements the callers send. This table is tiny, so a + // sequential scan is cheaper here; with it disabled, the planner can use + // the partial index only when each statement implies the index predicate. + const plans = await db.transaction(async tx => { + await tx.execute(sql`set local enable_seqscan = off`); + const explained: string[] = []; + for (const query of [ + pendingBudgetCancellationScopesQuery(db, 1).getSQL(), + pendingBudgetCancellationScopesQuery(db).getSQL(), + pendingBudgetRunIdsForScopeQuery(db, intent).getSQL(), + sql`select exists (${pendingBudgetRunCancellationExistsQuery()})`, + ]) explained.push(JSON.stringify(await tx.execute(sql`explain (format json) ${query}`))); + return explained; + }); + expect(plans).toHaveLength(4); + for (const plan of plans) expect(plan).toContain("heartbeat_runs_pending_budget_cancellation_idx"); + }); }); diff --git a/server/src/services/heartbeat/budget-cancellation.ts b/server/src/services/heartbeat/budget-cancellation.ts index a2d6c31b082..abcd48ed250 100644 --- a/server/src/services/heartbeat/budget-cancellation.ts +++ b/server/src/services/heartbeat/budget-cancellation.ts @@ -19,17 +19,35 @@ export function budgetCancellationScopePredicate(scope: Pick'startupCancellation'->'budgetEnforcement'->>'scopeId' = ${scope.scopeId}`); } -export async function listPendingBudgetCancellationScopes(db: Db, limit = 100): Promise { +/** The run IDs that still owe one scope's committed Stop (run control). */ +export function pendingBudgetRunIdsForScopeQuery(db: Db, scope: Pick) { + return db.select({ id: heartbeatRuns.id }).from(heartbeatRuns) + .where(and(pendingBudgetRunCancellationPredicate(), budgetCancellationScopePredicate(scope))); +} + +/** The idle-sleep check: any committed Stop that is not delivered yet. */ +export function pendingBudgetRunCancellationExistsQuery() { + return sql`SELECT 1 FROM heartbeat_runs WHERE ${pendingBudgetRunCancellationPredicate()}`; +} + +/** Oldest pending scopes first. The partial index + * `heartbeat_runs_pending_budget_cancellation_idx` keeps this off the full + * run history; its predicate must match `pendingBudgetRunCancellationPredicate`. */ +export function pendingBudgetCancellationScopesQuery(db: Db, limit = 100) { const intent = sql`${heartbeatRuns.resultJson}->'startupCancellation'->'budgetEnforcement'`; const scopeType = sql`${intent}->>'scopeType'`, scopeId = sql`${intent}->>'scopeId'`; const attemptAt = sql`coalesce((${intent}->>'lastAttemptAt')::timestamptz, (${heartbeatRuns.resultJson}->'startupCancellation'->>'requestedAt')::timestamptz, ${heartbeatRuns.updatedAt})`; - const rows = await db.select({ companyId: heartbeatRuns.companyId, + return db.select({ companyId: heartbeatRuns.companyId, intent: sql>`(array_agg(${intent} order by ${attemptAt}, ${heartbeatRuns.id}))[1]`, }).from(heartbeatRuns).where(pendingBudgetRunCancellationPredicate()) .groupBy(heartbeatRuns.companyId, scopeType, scopeId) .orderBy(sql`min(${attemptAt})`, heartbeatRuns.companyId, scopeType, scopeId) .limit(limit); +} + +export async function listPendingBudgetCancellationScopes(db: Db, limit = 100): Promise { + const rows = await pendingBudgetCancellationScopesQuery(db, limit); return rows.flatMap(({ companyId, intent }) => { if (intent.companyId !== companyId || !["company", "agent", "project"].includes(String(intent.scopeType)) || typeof intent.scopeId !== "string" || !isUuidLike(intent.scopeId) diff --git a/server/src/services/heartbeat/run-control.ts b/server/src/services/heartbeat/run-control.ts index 6cae0c38ab3..bd91d9ef3fc 100644 --- a/server/src/services/heartbeat/run-control.ts +++ b/server/src/services/heartbeat/run-control.ts @@ -72,7 +72,7 @@ import { waitForAdapterStop, } from "../adapter-execution-control.js"; import { AdapterStopTimeoutError } from "../adapter-stop-timeout.js"; -import { budgetCancellationScopePredicate, pendingBudgetRunCancellationPredicate } from "./budget-cancellation.js"; +import { pendingBudgetRunCancellationPredicate, pendingBudgetRunIdsForScopeQuery } from "./budget-cancellation.js"; import { notifyBudgetStopsSettled } from "../budget-resume-continuation.js"; import { withBudgetSpan } from "../budget-tracing.js"; import { captureException } from "../../sentry.js"; @@ -1818,9 +1818,7 @@ export function createHeartbeatRunControl(db: Db, dependencies: HeartbeatRunCont }); // A later policy change revokes permission to fence new work, not the Stop // already committed for older work. Recovery must deliver those intents. - const pending = await db.select({ id: heartbeatRuns.id }).from(heartbeatRuns).where(and( - pendingBudgetRunCancellationPredicate(), budgetCancellationScopePredicate(scope), - )); + const pending = await pendingBudgetRunIdsForScopeQuery(db, scope); const runIds = [...new Set([...(newlyFenced ?? []), ...pending.map(run => run.id)])]; if (!runIds.length) return; const failures: unknown[] = []; diff --git a/server/src/services/idle-sleep-safety.ts b/server/src/services/idle-sleep-safety.ts index c0f2794b4ad..c560331574c 100644 --- a/server/src/services/idle-sleep-safety.ts +++ b/server/src/services/idle-sleep-safety.ts @@ -17,7 +17,7 @@ import { sql, type SQLWrapper } from "drizzle-orm"; import { createIdleSleepDiagnostics, type IdleSleepDiagnosticSink, type IdleSleepDiagnostics } from "./idle-sleep-diagnostics.js"; import { verifiedResourceCleanup } from "./cleanup-verification-proof.js"; import { companyHoldMismatch } from "../modules/agent-lifecycle/index.js"; -import { pendingBudgetRunCancellationPredicate } from "./heartbeat/budget-cancellation.js"; +import { pendingBudgetRunCancellationExistsQuery } from "./heartbeat/budget-cancellation.js"; import { nativeResponsePresentationPendingQuery } from "./native-runtime/native-response-presentation-query.js"; import { passiveIdleConversation } from "./conversation-idle-proof.js"; import { idleChatHistoryChecks } from "./chat-channels/idle-history-proof.js"; @@ -443,7 +443,7 @@ export async function readIdleSleepSafety( }); checks.push(workCheck("agents.company_hold", sql`SELECT 1 FROM agents JOIN companies ON companies.id = agents.company_id WHERE ${companyHoldMismatch()}`)); - checks.push(workCheck("heartbeat_runs.budget_cancellation", sql`SELECT 1 FROM heartbeat_runs WHERE ${pendingBudgetRunCancellationPredicate()}`)); + checks.push(workCheck("heartbeat_runs.budget_cancellation", pendingBudgetRunCancellationExistsQuery())); checks.push(workCheck("native_response_presentation", nativeResponsePresentationPendingQuery())); // Only the current worker's owned runtime drain can exempt a plugin. // Missing workers, old SDKs, and newly installed plugins still block.