diff --git a/doc/DATABASE.md b/doc/DATABASE.md index 265f48d560d..ee653377291 100644 --- a/doc/DATABASE.md +++ b/doc/DATABASE.md @@ -438,6 +438,15 @@ indexes also serve the comment and delivery `on delete set null` actions. Issue lists share one external review-wait read per issue between the review and blocked-inbox attention maps. +Migration 0457 builds `heartbeat_runs_company_native_session_idx` concurrently. +It keys runs by company and `native_session_id`, partial on a non-null session. +Native retention, recovery, resume and archive reads find a session's runs +through it instead of reading the agent's whole run history. To measure it: + +```sh +pnpm --filter @paperclipai/server exec tsx ../scripts/benchmark-native-session-run-index.ts +``` + 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 b6524aa1b51..add40c1c478 100644 --- a/packages/db/src/client.ts +++ b/packages/db/src/client.ts @@ -619,6 +619,7 @@ const MIGRATIONS_WITH_TOP_LEVEL_STATEMENTS = new Set([ "0454_lumpy_ronan.sql", "0455_pending_budget_cancellation_index.sql", "0456_chat_review_lookup_indexes.sql", + "0457_native_session_run_index.sql", "0384_funny_adam_warlock.sql", diff --git a/packages/db/src/migrations/0457_native_session_run_index.sql b/packages/db/src/migrations/0457_native_session_run_index.sql new file mode 100644 index 00000000000..89a9ca08e19 --- /dev/null +++ b/packages/db/src/migrations/0457_native_session_run_index.sql @@ -0,0 +1 @@ +CREATE INDEX CONCURRENTLY IF NOT EXISTS "heartbeat_runs_company_native_session_idx" ON "heartbeat_runs" USING btree ("company_id","native_session_id") WHERE "heartbeat_runs"."native_session_id" is not null; \ No newline at end of file diff --git a/packages/db/src/migrations/meta/0452_snapshot.json b/packages/db/src/migrations/meta/0457_snapshot.json similarity index 99% rename from packages/db/src/migrations/meta/0452_snapshot.json rename to packages/db/src/migrations/meta/0457_snapshot.json index 082e8419e84..7372eea79c5 100644 --- a/packages/db/src/migrations/meta/0452_snapshot.json +++ b/packages/db/src/migrations/meta/0457_snapshot.json @@ -1,6 +1,6 @@ { - "id": "64e8be1e-e0cb-4cd8-85e0-0fdcce5eafde", - "prevId": "b5699366-77c2-48c8-b7dc-658a75de7658", + "id": "109f5f96-1597-40da-9e08-6ed488f8ad27", + "prevId": "a2606b41-36c2-4eb5-a59d-28cc730bc455", "version": "7", "dialect": "postgresql", "tables": { @@ -10955,6 +10955,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 +12415,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": { @@ -34552,6 +34612,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": { @@ -35393,6 +35481,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": [ @@ -35431,6 +35535,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": [ @@ -35562,6 +35687,49 @@ "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_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": [ diff --git a/packages/db/src/migrations/meta/_journal.json b/packages/db/src/migrations/meta/_journal.json index 6da1ec71b79..006da97123e 100644 --- a/packages/db/src/migrations/meta/_journal.json +++ b/packages/db/src/migrations/meta/_journal.json @@ -3186,6 +3186,13 @@ "when": 1794270045155, "tag": "0456_chat_review_lookup_indexes", "breakpoints": true + }, + { + "idx": 457, + "version": "7", + "when": 1794323383634, + "tag": "0457_native_session_run_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 819aaa20a0c..59fc9150286 100644 --- a/packages/db/src/schema/heartbeat_runs.ts +++ b/packages/db/src/schema/heartbeat_runs.ts @@ -208,6 +208,12 @@ export const heartbeatRuns = pgTable( companyNativeWorkspaceIdx: index("heartbeat_runs_company_native_workspace_idx") .on(table.companyId, sql`(${table.runnerProfileJson}->'nativeWorkspaceSync'->>'workspaceId')`) .concurrently(), + // Runs of one native session. Retention, recovery and resume reads filter + // by company and session; without this they walk the agent's run history. + companyNativeSessionIdx: index("heartbeat_runs_company_native_session_idx") + .on(table.companyId, table.nativeSessionId) + .where(sql`${table.nativeSessionId} is not null`) + .concurrently(), companyIssueCreatedIdx: index("heartbeat_runs_company_issue_created_idx").on( table.companyId, table.issueId, diff --git a/scripts/benchmark-native-session-run-index.ts b/scripts/benchmark-native-session-run-index.ts new file mode 100644 index 00000000000..4f887bc590c --- /dev/null +++ b/scripts/benchmark-native-session-run-index.ts @@ -0,0 +1,182 @@ +/** + * Local proof for the native-session run lookup (migration 0457). + * + * Runs the real checkpoint-retention pass (dry run) against a synthetic agent + * history, captures the session-run statement that `sessionRetentionBlocker` + * sends, and replays that exact SQL with and without + * `heartbeat_runs_company_native_session_idx`. Order: no index, index, no + * index, index, to expose cache-order bias. Fails unless every pass returns the + * same rows and the same service result. Then times the migration through the + * normal migration runner, including a rerun. + * + * pnpm --filter @paperclipai/server exec tsx ../scripts/benchmark-native-session-run-index.ts + * + * Optional sizes: RUNS (default 50000 runs of the measured agent), OTHER_RUNS + * (50000 runs of a second agent), RUNS_PER_SESSION (4), PROFILE_KIB (4). + */ +import assert from "node:assert/strict"; +import { createHash } from "node:crypto"; +import { promises as fs, readFileSync } from "node:fs"; +import os from "node:os"; +import path from "node:path"; +import { agents, applyPendingMigrations, companies, createDb, issues, nativeRunFinalizations, + nativeRunnerCheckpoints, startEmbeddedPostgresTestDatabase } from "../packages/db/src/index.js"; +import { nativeHarnessBackupExpirationService } from "../server/src/services/native-runtime/native-harness-backup-expiration.js"; + +const size = (name: string, fallback: number) => Number(process.env[name] ?? fallback); +const RUNS = size("RUNS", 50_000), OTHER_RUNS = size("OTHER_RUNS", 50_000); +const RUNS_PER_SESSION = size("RUNS_PER_SESSION", 4), PROFILE_KIB = size("PROFILE_KIB", 4); +const DAY_MS = 24 * 60 * 60 * 1000; +const migrationName = "0457_native_session_run_index"; +const migration = readFileSync(new URL(`../packages/db/src/migrations/${migrationName}.sql`, import.meta.url), "utf8"); +const statement = migration.trim(); +// The migration runner records each migration by the SHA-256 of its file. +const migrationHash = createHash("sha256").update(migration).digest("hex"); +const indexName = statement.match(/IF NOT EXISTS "([a-z_]+)"/)![1]!; + +const keepAlive = setInterval(() => {}, 1_000); +let temporary: Awaited> | undefined; +let db: ReturnType | undefined; +let stateBase: string | undefined; +try { + temporary = await startEmbeddedPostgresTestDatabase("native-session-run-index-"); + db = createDb(temporary.connectionString); + const sql = db.$client; + await sql.unsafe("create extension if not exists pgcrypto"); + const now = new Date(); + const finishedAt = new Date(now.getTime() - 10 * DAY_MS); + const [company] = await db.insert(companies).values({ name: "Index", issuePrefix: "IDX" }).returning(); + const [agent] = await db.insert(agents).values({ companyId: company!.id, name: "Measured" }).returning(); + const [other] = await db.insert(agents).values({ companyId: company!.id, name: "Other" }).returning(); + const [issue] = await db.insert(issues).values({ companyId: company!.id, title: "Closed", status: "done", + completedAt: finishedAt }).returning(); + + const loadStart = performance.now(); + // Incompressible profiles; one run in 100 has an incomplete recovery marker. + const seedRuns = (agentId: string, count: number) => sql.unsafe(` + with sessions as (select s, gen_random_uuid() as id from generate_series(0, ($3::int - 1) / $4::int) s) + insert into heartbeat_runs (company_id, agent_id, status, issue_id, native_issue_id, native_session_id, + runtime_mode, finished_at, created_at, started_at, runner_profile_json) + select $1::uuid, $2::uuid, 'succeeded', $5::uuid, $5::uuid, sessions.id, 'native', $6::timestamptz, + $6::timestamptz - make_interval(secs => g), $6::timestamptz - make_interval(secs => g), + jsonb_build_object('padding', (select string_agg(encode(gen_random_bytes(768), 'base64'), '') + from generate_series(1, $7::int) i where g > 0), + 'nativeControllerRecovery', case when g % 100 = 0 then jsonb_build_object( + 'schema', 'paperclip.native-controller-recovery.v1', 'checkpoint', null) end) + from generate_series(0, $3::int - 1) g join sessions on sessions.s = g / $4::int`, + [company!.id, agentId, count, RUNS_PER_SESSION, issue!.id, finishedAt.toISOString(), PROFILE_KIB]); + await seedRuns(agent!.id, RUNS); + await seedRuns(other!.id, OTHER_RUNS); + // The retention candidate: one settled run whose session has other runs. + const [target] = await sql.unsafe(`select id, native_session_id from heartbeat_runs where agent_id = $1::uuid + order by created_at limit 1 offset $2::int`, [agent!.id, Math.floor(RUNS / 2)]); + const sessionRuns = await sql.unsafe(`select id from heartbeat_runs where native_session_id = $1::uuid`, + [target!.native_session_id as string]); + for (const run of sessionRuns) await db.insert(nativeRunFinalizations).values({ runId: run.id, + companyId: company!.id, issueId: issue!.id, phase: "committed" }); + await db.insert(nativeRunnerCheckpoints).values({ companyId: company!.id, issueId: issue!.id, + runId: target!.id as string, part: "header", material: { encrypted: true } }); + await sql.unsafe("vacuum analyze"); + const loadMs = performance.now() - loadStart; + const tableSizes = await sql.unsafe(`select relname as table, n_live_tup::bigint::text as rows, + pg_size_pretty(pg_total_relation_size(relid)) as total_size from pg_stat_user_tables + where relname = 'heartbeat_runs'`); + + stateBase = await fs.mkdtemp(path.join(os.tmpdir(), "native-session-index-state-")); + const service = nativeHarnessBackupExpirationService(db, { stateBase, closedForDays: 7, now, dryRun: true, + isSessionBusy: () => false }); + type Captured = { sql: string; params: unknown[] }; + async function runPass() { + const captured: Captured[] = []; + const prior = sql.options.debug; + sql.options.debug = (_connection, query, params) => { captured.push({ sql: query, params: [...params] }); }; + try { + const result = await service.tickCheckpoints(); + const lookup = captured.filter(query => query.sql.includes("nativeControllerRecovery") + && query.sql.includes('"heartbeat_runs"."native_session_id" = ')); + assert.equal(lookup.length, 1, "One session-run lookup per retention candidate"); + return { result, lookup: lookup[0]! }; + } finally { sql.options.debug = prior; } + } + const median = (values: number[]) => [...values].sort((a, b) => a - b)[Math.floor(values.length / 2)]!; + const round = (value: number) => Number(value.toFixed(3)); + async function explain(query: Captured) { + const [result] = await sql.unsafe(`EXPLAIN (ANALYZE, BUFFERS, FORMAT JSON) ${query.sql}`, query.params as never[]); + return (result!["QUERY PLAN"] as Array<{ Plan: Record; "Execution Time": number }>)[0]!; + } + async function measure(withIndex: boolean) { + let buildMs: number | null = null; + if (withIndex) { + const start = performance.now(); + await sql.unsafe(statement); + buildMs = round(performance.now() - start); + } else await sql.unsafe(`DROP INDEX CONCURRENTLY IF EXISTS "${indexName}"`); + await sql.unsafe("analyze heartbeat_runs"); + const pass = await runPass(); + const rows = (await sql.unsafe(pass.lookup.sql, pass.lookup.params as never[]).values()) + .map(row => row.map(value => value instanceof Date ? value.toISOString() : value?.toString() ?? null)) + .sort((a, b) => String(a[0]).localeCompare(String(b[0]))); + for (let n = 0; n < 3; n++) await explain(pass.lookup); + const samples: number[] = []; + for (let n = 0; n < 15; n++) samples.push((await explain(pass.lookup))["Execution Time"]); + const plan = await explain(pass.lookup); + const nodes: string[] = []; + const walk = (node: Record) => { + nodes.push(`${node["Node Type"]}${node["Index Name"] ? ` using ${node["Index Name"]}` : ""}`); + for (const child of (node.Plans as Array> | undefined) ?? []) walk(child); + }; + walk(plan.Plan); + const wall: number[] = []; + for (let n = 0; n < 9; n++) { + const start = performance.now(); + await runPass(); + wall.push(performance.now() - start); + } + return { withIndex, buildMs, sql: pass.lookup.sql, rows, result: pass.result, + executionMs: { median: round(median(samples)), min: round(Math.min(...samples)), max: round(Math.max(...samples)) }, + retentionPassWallMs: { median: round(median(wall)), min: round(Math.min(...wall)), max: round(Math.max(...wall)) }, + sharedHitBlocks: plan.Plan["Shared Hit Blocks"], sharedReadBlocks: plan.Plan["Shared Read Blocks"], nodes }; + } + const runs = [await measure(false), await measure(true), await measure(false), await measure(true)]; + const reference = runs[0]!; + for (const run of runs) { + assert.equal(run.sql, reference.sql, "The service SQL must not change"); + assert.deepEqual(run.rows, reference.rows, "Session-run rows must match"); + assert.deepEqual(run.result, reference.result, "Retention results must match"); + } + assert.equal(reference.rows.length, RUNS_PER_SESSION); + assert.equal(reference.result.checkpointCandidates, 1); + assert.ok(runs[1]!.nodes.some(node => node.includes(indexName)), "The index must serve the lookup"); + + async function migrate() { + const removed = await sql.unsafe(`delete from drizzle.__drizzle_migrations where hash = $1 returning hash`, [migrationHash]); + assert.equal(removed.length, 1, `Exactly one ${migrationName} receipt`); + const start = performance.now(); + await applyPendingMigrations(temporary!.connectionString); + return round(performance.now() - start); + } + await sql.unsafe(`DROP INDEX CONCURRENTLY IF EXISTS "${indexName}"`); + const firstRunMs = await migrate(); + const [valid] = await sql.unsafe(`select i.indisvalid, pg_size_pretty(pg_relation_size(c.oid)) as size + from pg_index i join pg_class c on c.oid = i.indexrelid where c.relname = $1`, [indexName]); + assert.equal(valid?.indisvalid, true); + const rerunMs = await migrate(); + + console.log(JSON.stringify({ + postgresVersion: (await sql.unsafe("select version()"))[0]!.version, + settings: await sql.unsafe("select name, setting from pg_settings where name in ('shared_buffers','work_mem','jit') order by name"), + fixture: { synthetic: true, measuredAgentRuns: RUNS, otherAgentRuns: OTHER_RUNS, runsPerSession: RUNS_PER_SESSION, + profileKiB: PROFILE_KIB, loadMs: Math.round(loadMs), tableSizes }, + correctness: { sqlUnchanged: true, rowsEqual: true, retentionResultEqual: true, rows: reference.rows.length }, + results: runs.map(({ sql: _sql, rows: _rows, result: _result, ...run }) => run), + sql: reference.sql, + migration: { firstRunMs, rerunMs, index: valid }, + }, null, 2)); +} finally { + try { + await db?.$client.end({ timeout: 0 }); + await temporary?.cleanup(); + if (stateBase) await fs.rm(stateBase, { recursive: true, force: true }); + } finally { clearInterval(keepAlive); } +} + diff --git a/server/src/__tests__/native-harness-backup-expiration.test.ts b/server/src/__tests__/native-harness-backup-expiration.test.ts index aa3d5ee3197..1251a29bb67 100644 --- a/server/src/__tests__/native-harness-backup-expiration.test.ts +++ b/server/src/__tests__/native-harness-backup-expiration.test.ts @@ -345,6 +345,40 @@ describeEmbeddedPostgres("native harness backup expiration", () => { }, ); + it("reads one company's session runs through the company/session index", async () => { + const session = await seedSession({ finalization: { phase: "committed" } }); + await seedCheckpoint(session); + await db.delete(agentTaskSessions); + await fs.rm(stateBase, { recursive: true, force: true }); + // Another company's live run with the same session ID must not protect it. + const foreign = await seedSession({ runStatus: "running" }); + await db.update(heartbeatRuns).set({ nativeSessionId: session.sessionId }).where(eq(heartbeatRuns.id, foreign.runId)); + await db.delete(agentTaskSessions); + // Another live run of the same agent in a different session must not either. + await db.insert(heartbeatRuns).values({ companyId: session.companyId, agentId: session.agentId, status: "running", + issueId: session.issueId, nativeIssueId: session.issueId, nativeSessionId: randomUUID() }); + const statements: Array<{ sql: string; params: unknown[] }> = []; + const prior = db.$client.options.debug; + db.$client.options.debug = (_connection, query, params) => { statements.push({ sql: query, params: [...params] }); }; + try { + expect(await service().tickCheckpoints({ dryRun: true })).toMatchObject({ checkpointCandidates: 1, failed: 0 }); + } finally { db.$client.options.debug = prior; } + const lookup = statements.find(query => query.sql.includes("nativeControllerRecovery") + && query.sql.includes('"heartbeat_runs"."native_session_id" = ')); + expect(lookup).toBeTruthy(); + const plan = await db.$client.begin(async sql => { + await sql.unsafe("set local enable_seqscan = off"); + return JSON.stringify(await sql.unsafe(`explain (format json) ${lookup!.sql}`, lookup!.params as never[])); + }); + expect(plan).toContain("heartbeat_runs_company_native_session_idx"); + + // A second live run of the same session, company and agent protects it. + await db.insert(heartbeatRuns).values({ companyId: session.companyId, agentId: session.agentId, status: "running", + issueId: session.issueId, nativeIssueId: session.issueId, nativeSessionId: session.sessionId }); + expect(await service().tickCheckpoints({ dryRun: true })).toMatchObject({ checkpointCandidates: 0, failed: 0 }); + expect(await db.select().from(nativeRunnerCheckpoints)).toHaveLength(3); + }); + it("continues checkpoint expiration past retained sessions within bounded pages", async () => { for (let index = 0; index < 3; index++) { const session = await seedSession({ finalization: { phase: "committed" } });