mirror of
https://github.com/paperclipai/paperclip.git
synced 2026-11-11 05:11:49 +01:00
perf(db): index pending budget Stops instead of scanning all run results (#20391)
## 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 <noreply@anthropic.com>
This commit is contained in:
DottaandClaude Opus 5.5 authored and GitHub committed
1 parent
ab80fea62f
commit
0116e45958
11 files changed
+466
-13
No files matched your search
@@ -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
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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;
|
||||
+212
-3
@@ -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": [
|
||||
+7
@@ -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
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -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')`)
|
||||
|
||||
@@ -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<ReturnType<typeof startEmbeddedPostgresTestDatabase>> | 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<string, unknown> = {}) =>
|
||||
({ companyId, scopeType, scopeId, policyId, version: 3, ...extra });
|
||||
const special: Array<{ companyId: string; agentId: string; status: string; updatedAt: Date; resultJson: Record<string, unknown> }> = [];
|
||||
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<unknown>) => {
|
||||
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<string, unknown>) { nodes.push(String(node["Node Type"]) + (node["Index Name"] ? `(${node["Index Name"]})` : ""));
|
||||
for (const child of (node.Plans as Record<string, unknown>[] | 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<string, Record<string, Awaited<ReturnType<typeof measure>>>> = {};
|
||||
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); }
|
||||
}
|
||||
@@ -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");
|
||||
});
|
||||
});
|
||||
@@ -19,17 +19,35 @@ export function budgetCancellationScopePredicate(scope: Pick<BudgetEnforcementSc
|
||||
sql`${heartbeatRuns.resultJson}->'startupCancellation'->'budgetEnforcement'->>'scopeId' = ${scope.scopeId}`);
|
||||
}
|
||||
|
||||
export async function listPendingBudgetCancellationScopes(db: Db, limit = 100): Promise<BudgetEnforcementScope[]> {
|
||||
/** The run IDs that still owe one scope's committed Stop (run control). */
|
||||
export function pendingBudgetRunIdsForScopeQuery(db: Db, scope: Pick<BudgetEnforcementScope, "companyId" | "scopeType" | "scopeId">) {
|
||||
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<string>`${intent}->>'scopeType'`, scopeId = sql<string>`${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<Record<string, unknown>>`(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<BudgetEnforcementScope[]> {
|
||||
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)
|
||||
|
||||
@@ -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[] = [];
|
||||
|
||||
@@ -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.
|
||||
|
||||
Reference in new issue
Block a user