diff --git a/server/src/services/heartbeat-run-events.ts b/server/src/services/heartbeat-run-events.ts index cbc23b0e7a7..1c3325ebd05 100644 --- a/server/src/services/heartbeat-run-events.ts +++ b/server/src/services/heartbeat-run-events.ts @@ -36,8 +36,19 @@ export interface AppendHeartbeatRunEventInput { protocolSchemaVersion: number; canonicalPayload: Record; }; + /** + * Run columns to change in the same UPDATE that reserves the sequence. A + * second UPDATE of a run row this transaction already changed makes + * PostgreSQL re-check every foreign key on heartbeat_runs (KEY SHARE on the + * agent, issue, predecessor run, ...), so it can wait behind, or deadlock + * with, an unrelated transaction that holds one of those rows FOR UPDATE. + */ + runUpdate?: HeartbeatRunEventRunUpdate; } +export type HeartbeatRunEventRunUpdate = Partial>; + export interface AppendHeartbeatRunEventResult { row: typeof heartbeatRunEvents.$inferSelect; disposition: "committed" | "duplicate"; @@ -61,10 +72,12 @@ export class HeartbeatRunEventConflictError extends Error { export async function allocateHeartbeatRunEventSeq( db: Db, runId: string, + runUpdate: HeartbeatRunEventRunUpdate = {}, ): Promise { const [updated] = await db .update(heartbeatRuns) .set({ + ...runUpdate, nextEventSeq: sql`${heartbeatRuns.nextEventSeq} + 1`, updatedAt: new Date(), }) @@ -86,7 +99,9 @@ export async function appendHeartbeatRunEvent( const projectsReassessment = readNativeEventEvidenceOutcome(input.payload?.prpEvent) !== null; const projectsProgress = isSafeNativeChatProgressEventType(input.eventType); const projectsResponse = isNativeResponseEvidence(input.eventType, input.payload?.prpEvent); - if (!input.nativeSource && !input.retryExhaustion && !projectsProgress && !projectsResponse && !projectsReassessment && !projectsRecovery && !projectsRetention) { + // Deduplicated receipts return without allocating, so they cannot carry run changes. + if (input.runUpdate && (input.nativeSource || input.retryExhaustion)) throw new Error("heartbeat_run_event_run_update_unsupported"); + if (!input.runUpdate && !input.nativeSource && !input.retryExhaustion && !projectsProgress && !projectsResponse && !projectsReassessment && !projectsRecovery && !projectsRetention) { const allocated = db.$with("allocated_run_event", {}).as(sql` update ${heartbeatRuns} set next_event_seq = ${heartbeatRuns.nextEventSeq} + 1, @@ -202,6 +217,7 @@ export async function appendHeartbeatRunEvent( const seq = await allocateHeartbeatRunEventSeq( tx as unknown as Db, input.runId, + input.runUpdate, ); const [row] = await tx.insert(heartbeatRunEvents).values({ companyId: input.companyId, diff --git a/server/src/services/heartbeat.ts b/server/src/services/heartbeat.ts index 52f70ed5ac2..a9fd5f8837d 100644 --- a/server/src/services/heartbeat.ts +++ b/server/src/services/heartbeat.ts @@ -82,6 +82,7 @@ import { decorateHeartbeatRunRuntimeStatus, createHeartbeatLifecycle, persistHeartbeatRunProcessMetadata, + retryTransientRunWrite, } from "./heartbeat/run-lifecycle.js"; export { redactDetectedSuccessfulRunProgressSummaryForBoard, @@ -3547,13 +3548,13 @@ export function heartbeatService( } // Bind cleanup intent to this exact process generation in the // same commit as its host-observed PID/group/start identity. - return db.transaction(async transaction => { + return retryTransientRunWrite(() => db.transaction(async transaction => { const tx = transaction as unknown as Db; const persisted = await persistHeartbeatRunProcessMetadata(tx, processRunId, meta); await recordLegacyCaseHandoffCleanupIntent(tx, { runId: processRunId, issueId, controllerBootId: legacyControllerBootId, scratch: runScratch }); return persisted; - }); + })); }, recordCurrentHeartbeatRunRuntimeProgress, dispatchPendingNativeStatusWakeups, diff --git a/server/src/services/heartbeat/process-identity-persistence.test.ts b/server/src/services/heartbeat/process-identity-persistence.test.ts new file mode 100644 index 00000000000..2571acf24a3 --- /dev/null +++ b/server/src/services/heartbeat/process-identity-persistence.test.ts @@ -0,0 +1,99 @@ +import { and, eq } from "drizzle-orm"; +import { afterAll, beforeAll, describe, expect, it } from "vitest"; +import { agents, companies, createDb, heartbeatRunEvents, heartbeatRuns, issues } from "@paperclipai/db"; +import { getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase } from "../../__tests__/helpers/embedded-postgres.js"; +import { PROCESS_IDENTITY_RECORDED } from "../native-local-process-stop.js"; +import { persistHeartbeatRunProcessMetadata, retryTransientRunWrite } from "./run-lifecycle.js"; + +const support = await getEmbeddedPostgresTestSupport(); +const companyId = "30000000-0000-4000-8000-000000000001"; +const agentId = "30000000-0000-4000-8000-000000000002"; +const issueId = "30000000-0000-4000-8000-000000000003"; +const predecessorRunId = "30000000-0000-4000-8000-000000000004"; +const runId = "30000000-0000-4000-8000-000000000005"; +type Tx = ReturnType; +const sleep = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms)); + +describe.skipIf(!support.supported)("spawn process identity persistence", () => { + let temporary: Awaited>; + let db: ReturnType; + let other: ReturnType; + + beforeAll(async () => { + temporary = await startEmbeddedPostgresTestDatabase("paperclip-process-identity-"); + db = createDb(temporary.connectionString); + other = createDb(temporary.connectionString); + await db.insert(companies).values({ id: companyId, name: "Identity", issuePrefix: "PID" }); + await db.insert(agents).values({ id: agentId, companyId, name: "Identity agent" }); + await db.insert(issues).values({ id: issueId, companyId, title: "Retry me" }); + await db.insert(heartbeatRuns).values([ + { id: predecessorRunId, companyId, agentId, issueId, scopeKind: "issue", status: "failed" }, + { id: runId, companyId, agentId, issueId, scopeKind: "issue", status: "running", retryOfRunId: predecessorRunId }, + ]); + }); + + afterAll(async () => { + await temporary?.cleanup(); + }); + + // A user Retry spawned the provider while another transaction held a parent + // row (the issue or predecessor run) FOR UPDATE. Writing the identity + // and then reserving the event sequence updated the run row twice, and the + // second UPDATE re-checked heartbeat_runs' foreign keys (KEY SHARE on those + // parents; the event row's own agent/company keys are checked either way). It waited behind the holder and then failed, so the live run + // ended as native_session_interrupted. + it.each([ + ["issue", (tx: Tx) => tx.select({ id: issues.id }).from(issues).where(eq(issues.id, issueId))], + ["predecessor run", (tx: Tx) => tx.select({ id: heartbeatRuns.id }).from(heartbeatRuns).where(eq(heartbeatRuns.id, predecessorRunId))], + ] as const)("records identity without waiting on a transaction holding the %s FOR UPDATE", async (_label, lockParent) => { + let release!: () => void; + const released = new Promise((resolve) => { release = resolve; }); + let locked!: () => void; + const holding = new Promise((resolve) => { locked = resolve; }); + const holder = other.transaction(async (tx) => { + await lockParent(tx as unknown as Tx).for("update"); + locked(); + await released; + }); + await holding; + try { + const pid = 40_000 + Math.floor(Math.random() * 1000); + const persisted = persistHeartbeatRunProcessMetadata(db, runId, { + pid, processGroupId: null, startedAt: new Date().toISOString(), targetKind: "remote", + }); + const outcome = await Promise.race([ + persisted.then((run) => ({ run }), (error: unknown) => ({ error })), + sleep(3_000).then(() => "blocked" as const), + ]); + expect(outcome).not.toBe("blocked"); + expect(outcome).not.toHaveProperty("error"); + const run = (outcome as { run: typeof heartbeatRuns.$inferSelect | null }).run; + expect(run).toMatchObject({ id: runId, processPid: pid, processHostHash: "remote" }); + // The cases share one run; select this write's event by its reserved sequence. + const [event] = await db.select().from(heartbeatRunEvents).where(and( + eq(heartbeatRunEvents.runId, runId), + eq(heartbeatRunEvents.seq, Number(run!.nextEventSeq) - 1), + )); + expect(event).toMatchObject({ eventType: PROCESS_IDENTITY_RECORDED }); + expect(event?.payload).toMatchObject({ processPid: pid, targetKind: "remote" }); + } finally { + release(); + await holder; + } + }, 30_000); +}); + +describe("transient run write retry", () => { + it("retries a deadlock victim and rethrows other failures", async () => { + let calls = 0; + await expect(retryTransientRunWrite(async () => { + calls += 1; + if (calls < 3) throw Object.assign(new Error("Failed query"), { cause: { code: "40P01" } }); + return "persisted"; + })).resolves.toBe("persisted"); + expect(calls).toBe(3); + await expect(retryTransientRunWrite(async () => { + throw Object.assign(new Error("Failed query"), { cause: { code: "23505" } }); + })).rejects.toThrow("Failed query"); + }); +}); diff --git a/server/src/services/heartbeat/run-lifecycle.ts b/server/src/services/heartbeat/run-lifecycle.ts index 1a97b33057b..650f7a504a4 100644 --- a/server/src/services/heartbeat/run-lifecycle.ts +++ b/server/src/services/heartbeat/run-lifecycle.ts @@ -576,6 +576,32 @@ function buildRunEventRuntimeProgress(input: { }; } +function isTransientRunWriteConflict(error: unknown) { + let current = error; + for (let depth = 0; current && typeof current === "object" && depth < 4; depth++) { + const code = (current as { code?: unknown }).code; + if (code === "40P01" || code === "40001") return true; + current = (current as { cause?: unknown }).cause; + } + return false; +} + +/** + * Spawn identity is the run's ownership proof. A deadlock victim or + * serialization failure rolled back cleanly, so the whole write is retried + * instead of failing a live provider session. + */ +export async function retryTransientRunWrite(write: () => Promise, attempts = 4): Promise { + for (let attempt = 1; ; attempt++) { + try { + return await write(); + } catch (error) { + if (attempt >= attempts || !isTransientRunWriteConflict(error)) throw error; + await new Promise((resolve) => setTimeout(resolve, 20 * attempt)); + } + } +} + export async function persistHeartbeatRunProcessMetadata( db: Db, runId: string, @@ -586,32 +612,36 @@ export async function persistHeartbeatRunProcessMetadata( () => null, ); const startedAt = new Date(observedStartedAt ?? meta.startedAt); + const processStartedAt = Number.isNaN(startedAt.getTime()) ? new Date() : startedAt; return db.transaction(async tx => { - const run = await tx - .update(heartbeatRuns) - .set({ + const [binding] = await tx + .select({ companyId: heartbeatRuns.companyId, agentId: heartbeatRuns.agentId }) + .from(heartbeatRuns) + .where(eq(heartbeatRuns.id, runId)) + .limit(1); + if (!binding) return null; + // One UPDATE writes the identity and reserves the event sequence. Updating + // the row twice here re-checks its foreign keys and can wait behind (or + // deadlock with) any transaction holding the agent, issue, or predecessor + // run FOR UPDATE, which failed spawn persistence and lost the run. + await appendHeartbeatRunEvent(tx as unknown as Db, { + companyId: binding.companyId, runId, agentId: binding.agentId, + eventType: PROCESS_IDENTITY_RECORDED, stream: "system", level: "info", + message: "Process identity recorded; prior stop evidence no longer applies.", + payload: { targetKind: meta.targetKind ?? "local", hostname: os.hostname(), + processPid: meta.pid, processGroupId: meta.processGroupId, + processStartedAt: processStartedAt.toISOString() }, + runUpdate: { legacyProcessSettlement: null, processPid: meta.pid, processGroupId: meta.processGroupId, // A sandbox PID never names a process in this PID space. processHostHash: meta.targetKind === "remote" ? "remote" : currentProcessSpaceHash(), - processStartedAt: Number.isNaN(startedAt.getTime()) - ? new Date() - : startedAt, - updatedAt: new Date(), - }) - .where(eq(heartbeatRuns.id, runId)) - .returning() - .then((rows) => rows[0] ?? null); - if (run) await appendHeartbeatRunEvent(tx as unknown as Db, { - companyId: run.companyId, runId, agentId: run.agentId, - eventType: PROCESS_IDENTITY_RECORDED, stream: "system", level: "info", - message: "Process identity recorded; prior stop evidence no longer applies.", - payload: { targetKind: meta.targetKind ?? "local", hostname: os.hostname(), - processPid: meta.pid, processGroupId: meta.processGroupId, - processStartedAt: run.processStartedAt?.toISOString() ?? null }, + processStartedAt, + }, }); - return run; + return tx.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, runId)) + .then((rows) => rows[0] ?? null); }); } @@ -2245,7 +2275,7 @@ export function createHeartbeatLifecycle(db: Db, dependencies: HeartbeatLifecycl runId: string, meta: { pid: number; processGroupId: number | null; startedAt: string; targetKind: "local" | "remote" }, ) { - return persistHeartbeatRunProcessMetadata(db, runId, meta); + return retryTransientRunWrite(() => persistHeartbeatRunProcessMetadata(db, runId, meta)); } async function clearDetachedRunWarning(runId: string) {