diff --git a/doc/DATABASE.md b/doc/DATABASE.md index f89eaae0b52..fa124cd343a 100644 --- a/doc/DATABASE.md +++ b/doc/DATABASE.md @@ -422,6 +422,23 @@ local fixture: pnpm --filter @paperclipai/server exec tsx ../scripts/benchmark-budget-cancellation-scan.ts --runs=20000 --payload-kib=16 ``` +Workspace activity and cleanup checks (`expirationScope` in +`server/src/services/git-workspace-expiration.ts`) find a workspace's runs with +one index probe per binding: `issue_id`, `native_issue_id`, the bound issues' +checkout and execution runs, and the legacy context keys `issueId`, `taskId` +and `executionWorkspaceId`. Migration 0460 builds +`heartbeat_runs_company_ctx_execution_workspace_idx` concurrently for the last +key. The index is partial (`context_snapshot ? 'executionWorkspaceId'`), so it +serves only queries that repeat that predicate and does not compete with other +company-scoped indexes. The context expressions and the predicate must stay +textually equal to the index definitions; a regression test checks that no +binding scans `heartbeat_runs`. +To compare both plans on a disposable local fixture: + +```sh +pnpm --filter @paperclipai/server exec tsx ../scripts/benchmark-workspace-activity.ts --baseline-ref= --runs-per-workspace=3000 +``` + GitHub review lookups use four indexes from migration 0456, which builds them concurrently. `chat_message_links_company_comment_idx` finds a source comment's link for inherited review waits and run-origin checks. diff --git a/packages/db/src/client.ts b/packages/db/src/client.ts index 0e0f80b846d..c5fdc77313c 100644 --- a/packages/db/src/client.ts +++ b/packages/db/src/client.ts @@ -622,6 +622,7 @@ const MIGRATIONS_WITH_TOP_LEVEL_STATEMENTS = new Set([ "0457_native_session_run_index.sql", "0458_chat_admission_issue_index.sql", "0459_environment_lease_backup_scope_index.sql", + "0460_workspace_activity_context_index.sql", "0384_funny_adam_warlock.sql", diff --git a/packages/db/src/migrations/0460_workspace_activity_context_index.sql b/packages/db/src/migrations/0460_workspace_activity_context_index.sql new file mode 100644 index 00000000000..36dd943c317 --- /dev/null +++ b/packages/db/src/migrations/0460_workspace_activity_context_index.sql @@ -0,0 +1 @@ +CREATE INDEX CONCURRENTLY IF NOT EXISTS "heartbeat_runs_company_ctx_execution_workspace_idx" ON "heartbeat_runs" USING btree ("company_id",("context_snapshot" ->> 'executionWorkspaceId')) WHERE "heartbeat_runs"."context_snapshot" ? 'executionWorkspaceId'; \ No newline at end of file diff --git a/packages/db/src/migrations/meta/0455_snapshot.json b/packages/db/src/migrations/meta/0460_snapshot.json similarity index 99% rename from packages/db/src/migrations/meta/0455_snapshot.json rename to packages/db/src/migrations/meta/0460_snapshot.json index bc7b1347a0d..38297d445c0 100644 --- a/packages/db/src/migrations/meta/0455_snapshot.json +++ b/packages/db/src/migrations/meta/0460_snapshot.json @@ -1,6 +1,6 @@ { - "id": "2be53987-902a-41b6-8cb4-6e89a768f8cf", - "prevId": "f98af245-4a2e-4c3d-b89e-6ed3e81a074c", + "id": "f93ccf18-7e86-4788-be77-b17430c86b2f", + "prevId": "e070677d-cfb8-4438-847c-29c4da7ca8c5", "version": "7", "dialect": "postgresql", "tables": { @@ -10188,6 +10188,40 @@ "method": "btree", "with": {} }, + "chat_actions_company_issue_kind_idx": { + "name": "chat_actions_company_issue_kind_idx", + "columns": [ + { + "expression": "company_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "(\"payload\"->>'issueId')", + "asc": true, + "isExpression": true, + "nulls": "last" + }, + { + "expression": "kind", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "status", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "where": "\"chat_actions\".\"kind\" IN ('inbound_wakeup', 'failed_run_retry')", + "concurrently": true, + "method": "btree", + "with": {} + }, "chat_actions_github_stale_head_idx": { "name": "chat_actions_github_stale_head_idx", "columns": [ @@ -10955,6 +10989,28 @@ "concurrently": false, "method": "btree", "with": {} + }, + "chat_deliveries_github_automatic_delivery_idx": { + "name": "chat_deliveries_github_automatic_delivery_idx", + "columns": [ + { + "expression": "endpoint_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "(\"normalized_event\"->'githubAutomatic'->'context'->>'deliveryId')", + "asc": true, + "isExpression": true, + "nulls": "last" + } + ], + "isUnique": false, + "where": "(\"chat_deliveries\".\"normalized_event\"->'githubAutomatic'->'context'->>'deliveryId') is not null", + "concurrently": true, + "method": "btree", + "with": {} } }, "foreignKeys": { @@ -12393,6 +12449,44 @@ "concurrently": false, "method": "btree", "with": {} + }, + "chat_message_links_company_comment_idx": { + "name": "chat_message_links_company_comment_idx", + "columns": [ + { + "expression": "company_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "comment_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "where": "\"chat_message_links\".\"comment_id\" is not null", + "concurrently": true, + "method": "btree", + "with": {} + }, + "chat_message_links_delivery_idx": { + "name": "chat_message_links_delivery_idx", + "columns": [ + { + "expression": "delivery_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "where": "\"chat_message_links\".\"delivery_id\" is not null", + "concurrently": true, + "method": "btree", + "with": {} } }, "foreignKeys": { @@ -29073,6 +29167,22 @@ "concurrently": false, "method": "btree", "with": {} + }, + "environment_leases_native_backup_scope_idx": { + "name": "environment_leases_native_backup_scope_idx", + "columns": [ + { + "expression": "(\"metadata\"->'nativeHarnessBackup'->>'sessionScopeSha256')", + "asc": true, + "isExpression": true, + "nulls": "last" + } + ], + "isUnique": false, + "where": "(\"environment_leases\".\"metadata\"->'nativeHarnessBackup'->>'sessionScopeSha256') is not null", + "concurrently": true, + "method": "btree", + "with": {} } }, "foreignKeys": { @@ -35475,6 +35585,27 @@ "method": "btree", "with": {} }, + "heartbeat_runs_company_wake_comment_idx": { + "name": "heartbeat_runs_company_wake_comment_idx", + "columns": [ + { + "expression": "company_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "(\"context_snapshot\"->>'wakeCommentId')", + "asc": true, + "isExpression": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": true, + "method": "btree", + "with": {} + }, "heartbeat_runs_cost_accounting_pending_idx": { "name": "heartbeat_runs_cost_accounting_pending_idx", "columns": [ @@ -35627,6 +35758,28 @@ "method": "btree", "with": {} }, + "heartbeat_runs_company_native_session_idx": { + "name": "heartbeat_runs_company_native_session_idx", + "columns": [ + { + "expression": "company_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "native_session_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "where": "\"heartbeat_runs\".\"native_session_id\" is not null", + "concurrently": true, + "method": "btree", + "with": {} + }, "heartbeat_runs_company_issue_created_idx": { "name": "heartbeat_runs_company_issue_created_idx", "columns": [ @@ -35864,6 +36017,28 @@ "method": "btree", "with": {} }, + "heartbeat_runs_company_ctx_execution_workspace_idx": { + "name": "heartbeat_runs_company_ctx_execution_workspace_idx", + "columns": [ + { + "expression": "company_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "(\"context_snapshot\" ->> 'executionWorkspaceId')", + "asc": true, + "isExpression": true, + "nulls": "last" + } + ], + "isUnique": false, + "where": "\"heartbeat_runs\".\"context_snapshot\" ? 'executionWorkspaceId'", + "concurrently": true, + "method": "btree", + "with": {} + }, "heartbeat_runs_company_ctx_paperclip_issue_created_idx": { "name": "heartbeat_runs_company_ctx_paperclip_issue_created_idx", "columns": [ diff --git a/packages/db/src/migrations/meta/_journal.json b/packages/db/src/migrations/meta/_journal.json index 45a14d43c57..a3932712d56 100644 --- a/packages/db/src/migrations/meta/_journal.json +++ b/packages/db/src/migrations/meta/_journal.json @@ -3207,6 +3207,13 @@ "when": 1794329161946, "tag": "0459_environment_lease_backup_scope_index", "breakpoints": true + }, + { + "idx": 460, + "version": "7", + "when": 1794332988056, + "tag": "0460_workspace_activity_context_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 59fc9150286..b10c7ca8d61 100644 --- a/packages/db/src/schema/heartbeat_runs.ts +++ b/packages/db/src/schema/heartbeat_runs.ts @@ -258,6 +258,14 @@ export const heartbeatRuns = pgTable( sql`(${table.contextSnapshot} ->> 'taskKey')`, table.createdAt.desc(), ), + // Runs that recorded an execution workspace in their legacy context. Workspace + // activity and cleanup checks probe it instead of decoding every run context. + // Partial, so it serves only queries that repeat the predicate and never + // competes with other company-scoped indexes. + companyCtxExecutionWorkspaceIdx: index("heartbeat_runs_company_ctx_execution_workspace_idx") + .on(table.companyId, sql`(${table.contextSnapshot} ->> 'executionWorkspaceId')`) + .where(sql`${table.contextSnapshot} ? 'executionWorkspaceId'`) + .concurrently(), companyCtxPaperclipIssueCreatedIdx: index("heartbeat_runs_company_ctx_paperclip_issue_created_idx") .on(table.companyId, sql`(${table.contextSnapshot} -> 'paperclipIssue' ->> 'id')`, table.createdAt.desc()) .concurrently(), diff --git a/scripts/benchmark-workspace-activity.ts b/scripts/benchmark-workspace-activity.ts index fad3007ab1d..95e5bbf7a48 100644 --- a/scripts/benchmark-workspace-activity.ts +++ b/scripts/benchmark-workspace-activity.ts @@ -24,7 +24,9 @@ const sha256 = (source: string) => createHash("sha256").update(source).digest("h const baselineOnly = process.argv.includes("--baseline-only"); const dockerPostgres = process.argv.includes("--docker-postgres=17"); const small = process.argv.includes("--small"); -const runsPerWorkspace = small ? 16 : 600; +const runsPerWorkspace = Number(process.argv.find(arg => arg.startsWith("--runs-per-workspace="))?.slice(21) ?? (small ? 16 : 600)); +assert.ok(Number.isInteger(runsPerWorkspace) && runsPerWorkspace >= 16 && runsPerWorkspace <= 20000, + "--runs-per-workspace must be an integer between 16 and 20000"); const bindingRows = Math.max(1, Math.floor(runsPerWorkspace / 20)); const descendantsPerWorkspace = Number(process.argv.find(arg => arg.startsWith("--descendants="))?.slice(14) ?? 1); assert.ok(Number.isInteger(descendantsPerWorkspace) && descendantsPerWorkspace >= 1 && descendantsPerWorkspace <= 500, @@ -114,7 +116,8 @@ try { title: `Child ${n}/${child}`, status: "done", completedAt: old, updatedAt: old }))); } const padding = randomBytes(24 * 1024).toString("base64"); - // Most history is unrelated, as in the reported scan; contexts and results are + // Most history is unrelated, as in the reported scan, and unrelated runs carry + // distinct legacy issue, task and workspace IDs as real history does; contexts and results are // incompressible and wide enough to live in TOAST. Existing indexes are retained. for (let n = 0; n < 8; n++) { const workspace = workspaces[n]; @@ -122,9 +125,9 @@ try { SELECT $1::uuid, $2::uuid, 'succeeded', CASE WHEN i <= ${bindingRows} THEN 'issue' ELSE 'company' END, CASE WHEN i <= ${bindingRows} THEN $3::uuid END, CASE WHEN i > ${bindingRows} AND i <= ${bindingRows * 2} THEN $3::uuid END, $6::timestamptz + i * interval '1 second', $6::timestamptz + i * interval '2 seconds', $6::timestamptz + i * interval '3 seconds', - jsonb_build_object('issueId', CASE WHEN i > ${bindingRows * 2} AND i <= ${bindingRows * 3} THEN $3::text ELSE $5::text END, - 'taskId', CASE WHEN i > ${bindingRows * 3} AND i <= ${bindingRows * 4} THEN $3::text ELSE $5::text END, - 'executionWorkspaceId', CASE WHEN i > ${bindingRows * 4} AND i <= ${bindingRows * 5} THEN $4::text ELSE 'unrelated' END, + jsonb_build_object('issueId', CASE WHEN i > ${bindingRows * 2} AND i <= ${bindingRows * 3} THEN $3::text WHEN i % 50 = 0 THEN $5::text ELSE gen_random_uuid()::text END, + 'taskId', CASE WHEN i > ${bindingRows * 3} AND i <= ${bindingRows * 4} THEN $3::text ELSE gen_random_uuid()::text END, + 'executionWorkspaceId', CASE WHEN i > ${bindingRows * 4} AND i <= ${bindingRows * 5} THEN $4::text ELSE gen_random_uuid()::text END, 'fixtureRun', ${n * runsPerWorkspace} + i, 'padding', $7::text), jsonb_build_object('stdout', $7::text) FROM generate_series(1, ${runsPerWorkspace}) i`, [companyId, agentId, workspace.sourceIssueId, workspace.id, unrelatedIssueId, old.toISOString(), padding]); diff --git a/server/src/__tests__/git-workspace-activity.test.ts b/server/src/__tests__/git-workspace-activity.test.ts index cb677f477e9..8c3b3cada52 100644 --- a/server/src/__tests__/git-workspace-activity.test.ts +++ b/server/src/__tests__/git-workspace-activity.test.ts @@ -72,6 +72,32 @@ describeDatabase("Git workspace activity bindings", () => { expect(await gitWorkspaceLastActivityAt(db, f.workspace)).toEqual(recent); }); + it.each(["issueId", "executionWorkspaceId"])("keeps a legacy %s binding on a run bound to another issue", async key => { + const f = await fixture(); + const [unrelated] = await db.insert(issues).values({ companyId: f.companyId, title: "Unrelated", updatedAt: old }).returning(); + await db.insert(heartbeatRuns).values({ companyId: f.companyId, agentId: f.agentId, scopeKind: "issue", issueId: unrelated.id, + nativeIssueId: unrelated.id, status: "running", createdAt: recent, + contextSnapshot: { [key]: key === "issueId" ? f.child.id : f.workspace.id } }); + expect(await f.eligible()).toBe(false); + expect(await workspaceMoveEligible(db, f.workspace)).toBe(false); + expect(await gitWorkspaceLastActivityAt(db, f.workspace)).toEqual(recent); + }); + + it("reaches every run binding through an index", async () => { + const f = await fixture(); + // With sequential scans disabled, a binding expression that no index + // matches still forces a sequential scan of heartbeat_runs. + const plan = await db.transaction(async tx => { + await tx.execute(sql`SET LOCAL enable_seqscan = off`); + return tx.execute<{ "QUERY PLAN": string }>(sql`EXPLAIN ${expirationScope(f.workspace)} SELECT id FROM relevant_runs`); + }); + const text = plan.map(row => row["QUERY PLAN"]).join("\n"); + expect(text).not.toMatch(/Seq Scan on heartbeat_runs/); + expect(text).toMatch(/heartbeat_runs_company_ctx_execution_workspace_idx/); + expect(text).toMatch(/heartbeat_runs_company_ctx_issue_created_idx/); + expect(text).toMatch(/heartbeat_runs_company_ctx_task_created_idx/); + }); + it("includes a multiply bound run exactly once", async () => { const f = await fixture(), runId = randomUUID(); await db.insert(heartbeatRuns).values({ id: runId, companyId: f.companyId, agentId: f.agentId, scopeKind: "issue", issueId: f.issueId, diff --git a/server/src/services/git-workspace-expiration.ts b/server/src/services/git-workspace-expiration.ts index 2c64fc50fca..2edfdba0ab0 100644 --- a/server/src/services/git-workspace-expiration.ts +++ b/server/src/services/git-workspace-expiration.ts @@ -104,9 +104,16 @@ async function restorePointBoundary(cwd: string, repoUrl: string, headOnRemote: /** Re-read the whole tree and every runtime/recovery consumer immediately before deletion. */ export function expirationScope(workspace: Workspace) { - // Decode the three legacy keys together only after cheap native/pointer - // matches fail. Strict JSONPath leaves non-object inputs unmatched, as ->> - // did, while those inputs can still match an independent native binding. + // Each binding is a separate index probe. A single OR over all bindings made + // PostgreSQL read every company run and decode its context JSON. Each key set + // is passed as an array so the planner probes the index instead of hashing a + // full scan; it does not count the cost of reading TOASTed contexts. The legacy + // context expressions match the heartbeat_runs expression indexes exactly, + // and the key test matches the partial workspace index predicate; + // ->> yields NULL for non-object contexts, so those rows stay unmatched there + // and can still match through an independent native binding. The UNION + // removes duplicate IDs, so the outer primary-key probe keeps each run once. + const companyRuns = sql`SELECT id FROM ${heartbeatRuns} WHERE company_id = ${workspace.companyId}`; return sql` WITH RECURSIVE tree AS ( SELECT id, status, completed_at, cancelled_at, updated_at FROM ${issues} @@ -117,23 +124,17 @@ export function expirationScope(workspace: Workspace) { ), bound AS ( SELECT id, updated_at, checkout_run_id, execution_run_id FROM ${issues} WHERE company_id = ${workspace.companyId} AND execution_workspace_id = ${workspace.id} - ), legacy_issue_ids AS MATERIALIZED ( - SELECT unnest(array_agg(id::text)) AS id FROM tree ), relevant_runs AS ( SELECT r.id, r.status, r.created_at, r.started_at, r.finished_at, r.controller_lease_expires_at - FROM ${heartbeatRuns} r WHERE r.company_id = ${workspace.companyId} AND ( - r.issue_id IN (SELECT id FROM tree) OR r.native_issue_id IN (SELECT id FROM tree) - OR r.issue_id IN (SELECT id FROM bound) OR r.native_issue_id IN (SELECT id FROM bound) - OR r.id IN (SELECT unnest(ARRAY[checkout_run_id, execution_run_id]) FROM bound) - OR EXISTS ( - SELECT 1 FROM jsonb_to_record( - jsonb_path_query_first(r.context_snapshot, 'strict $ ? (@.type() == "object")') - ) context("issueId" text, "taskId" text, "executionWorkspaceId" text) - WHERE context."issueId" IN (SELECT id FROM legacy_issue_ids) - OR context."taskId" IN (SELECT id FROM legacy_issue_ids) - OR context."executionWorkspaceId" = ${workspace.id} - ) - ) + FROM ${heartbeatRuns} r WHERE r.company_id = ${workspace.companyId} AND r.id = ANY(ARRAY( + ${companyRuns} AND issue_id = ANY(ARRAY(SELECT id FROM tree UNION SELECT id FROM bound)) + UNION ${companyRuns} AND native_issue_id = ANY(ARRAY(SELECT id FROM tree UNION SELECT id FROM bound)) + UNION SELECT unnest(ARRAY[checkout_run_id, execution_run_id]) FROM bound + UNION ${companyRuns} AND (context_snapshot ->> 'issueId') = ANY(ARRAY(SELECT id::text FROM tree)) + UNION ${companyRuns} AND (context_snapshot ->> 'taskId') = ANY(ARRAY(SELECT id::text FROM tree)) + UNION ${companyRuns} AND context_snapshot ? 'executionWorkspaceId' + AND (context_snapshot ->> 'executionWorkspaceId') = ${workspace.id} + )) )`; }