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.