diff --git a/server/src/__tests__/heartbeat-restart-claim-lapse.test.ts b/server/src/__tests__/heartbeat-restart-claim-lapse.test.ts new file mode 100644 index 00000000000..0e8058421c5 --- /dev/null +++ b/server/src/__tests__/heartbeat-restart-claim-lapse.test.ts @@ -0,0 +1,98 @@ +import { randomUUID } from "node:crypto"; +import { and, eq, sql } from "drizzle-orm"; +import { expect, it, vi } from "vitest"; +import { authUsers, companyMemberships, heartbeatRunEvents, heartbeatRuns, nativeRunFinalizations } from "@paperclipai/db"; +import { heartbeatService } from "../services/heartbeat.js"; +import { releaseLapsedNativeRestartClaim } from "../services/native-runtime/native-execution-lease.js"; +import { instanceSettingsService } from "../services/instance-settings.js"; +import { describeEmbeddedPostgres, nativeRestartFixture } from "./helpers/native-restart.js"; +import { waitForCondition } from "./helpers/native-restart-fixture.js"; + +const preparationFault = vi.hoisted(() => ({ + current: null as null | { runId: string; phase: string; act: () => Promise }, +})); +vi.mock("../services/heartbeat/preparation-timing.js", async (importOriginal) => { + const actual = await importOriginal(); + return { + createHeartbeatPreparationTiming: (runId: string) => { + const timing = actual.createHeartbeatPreparationTiming(runId); + return (async (phase, work) => { + const fault = preparationFault.current; + if (fault?.runId === runId && fault.phase === phase) { + preparationFault.current = null; + await fault.act(); + } + return timing(phase, work); + }) as typeof timing; + }, + }; +}); + +describeEmbeddedPostgres("restart-claim preparation that outlives its claim lease", () => { + const fixture = nativeRestartFixture(); + + it.each(["live lease", "another owner", "newer generation", "other attempt", "run not running", "lapsed own claim"] as const)( + "releases a restart claim only when it is this exact lapsed claim (%s)", async (variant) => { + const f = await fixture.seedRun(`RELEASE-${variant.replace(/\W+/g, "-")}`); + const leaseOwner = `native-recovery:${randomUUID()}`; + await f.db.update(nativeRunFinalizations).set({ phase: "observed", leaseOwner, controllerGeneration: 3, attempt: 1, + leaseExpiresAt: new Date(Date.now() + (variant === "live lease" ? 60_000 : -1_000)) }) + .where(eq(nativeRunFinalizations.runId, f.runId)); + if (variant === "run not running") await f.db.update(heartbeatRuns).set({ status: "failed" }).where(eq(heartbeatRuns.id, f.runId)); + const released = await releaseLapsedNativeRestartClaim({ db: f.db, companyId: fixture.current.companyId, runId: f.runId, + leaseOwner: variant === "another owner" ? `native-recovery:${randomUUID()}` : leaseOwner, + controllerGeneration: variant === "newer generation" ? 2 : 3, providerAttempt: variant === "other attempt" ? 2 : 1 }); + const [owner] = await f.db.select().from(nativeRunFinalizations).where(eq(nativeRunFinalizations.runId, f.runId)); + if (variant === "lapsed own claim") { + expect(released).toBe(true); + expect(owner).toMatchObject({ leaseOwner: null, leaseExpiresAt: null, phase: "observed", controllerGeneration: 3 }); + return; + } + expect(released).toBe(false); + expect(owner).toMatchObject({ leaseOwner, phase: "observed", controllerGeneration: 3, attempt: 1 }); + }, 60_000); + + it("hands the run straight back to restart recovery instead of leaving it running unowned", async () => { + const f = await fixture.seedRun("LAPSED-RESTART-CLAIM"); + await instanceSettingsService(f.db).updateExperimental({ enableNativeRunner: true }); + const userId = randomUUID(); + await f.db.insert(authUsers).values({ id: userId, name: "Operator", email: `${userId}@example.invalid`, + createdAt: new Date(), updatedAt: new Date() }); + await f.db.insert(companyMemberships).values({ companyId: fixture.current.companyId, principalType: "user", + principalId: userId, membershipRole: "owner" }); + await f.db.update(heartbeatRuns).set({ responsibleUserId: userId, completionContractId: null, + completionContractSha256: null }).where(eq(heartbeatRuns.id, f.runId)); + const transitions = async () => (await f.db.select({ id: heartbeatRunEvents.id }).from(heartbeatRunEvents) + .where(and(eq(heartbeatRunEvents.runId, f.runId), eq(heartbeatRunEvents.eventType, "native.recovery.transition")))).length; + + preparationFault.current = { runId: f.runId, phase: "run_output", act: async () => { + await f.db.update(heartbeatRuns).set({ runnerProfileJson: sql`coalesce(${heartbeatRuns.runnerProfileJson}, '{}'::jsonb) + || jsonb_build_object('nativeWorkspaceSync', jsonb_build_object('schema', 'paperclip.native-workspace-sync/v3', + 'state', 'prepared', 'descriptorSha256', ${"d".repeat(64)}::text, 'baselineSha256', ${"b".repeat(64)}::text, + 'finalHostSha256', null, 'workspaceId', ${randomUUID()}::text, 'leaseId', ${randomUUID()}::text, + 'providerLeaseId', ${randomUUID()}::text, 'remoteCwd', '/work', 'resourceDisposition', 'destroy'))` }) + .where(eq(heartbeatRuns.id, f.runId)); + await f.db.update(nativeRunFinalizations).set({ leaseExpiresAt: new Date(Date.now() - 1_000) }) + .where(eq(nativeRunFinalizations.runId, f.runId)); + throw new Error("Object storage is unavailable"); + } }; + const heartbeat = heartbeatService(f.db, { + nativeSessionBackendFactory: vi.fn(() => { throw new Error("A lapsed recovery must not open a provider session"); }), + }); + let before = 0; + try { + const { claims } = await heartbeat.recoverNativeRunsAfterRestart(new Date(), [f.runId]); + expect(claims).toEqual([expect.objectContaining({ runId: f.runId })]); + before = await transitions(); + await waitForCondition("restart recovery reclaims or blocks the lapsed run", async () => { + const [owner] = await f.db.select().from(nativeRunFinalizations).where(eq(nativeRunFinalizations.runId, f.runId)); + return await transitions() > before && (owner.controllerGeneration !== claims[0]!.controllerGeneration + || owner.phase !== "observed" || owner.recoveryState === "blocked"); + }, 15_000); + } finally { + preparationFault.current = null; + await heartbeat.drainActiveRunExecutions(); + } + expect(await transitions()).toBeGreaterThan(before); + }, 90_000); +}); diff --git a/server/src/services/heartbeat.ts b/server/src/services/heartbeat.ts index 9c4365abde7..16b9897657d 100644 --- a/server/src/services/heartbeat.ts +++ b/server/src/services/heartbeat.ts @@ -53,7 +53,7 @@ import { prepareHeartbeatTaskContext } from "./heartbeat/task-context.js"; import { selectHeartbeatEnvironment } from "./heartbeat/environment-selection.js"; import { reconcileMuseNativeTerminalRun } from "./native-runtime/muse-terminal-reconciliation.js"; import { prepareHeartbeatWorkspace } from "./heartbeat/workspace-preparation.js"; -import { nativeRecoveryClaimRelinquished } from "./native-runtime/native-execution-lease.js"; +import { nativeRecoveryClaimRelinquished, releaseLapsedNativeRestartClaim } from "./native-runtime/native-execution-lease.js"; import { executeHeartbeatRuntime, NativeSessionResumeScheduledError, NativeWorkspaceFinalizeScheduledError } from "./heartbeat/runtime-execution.js"; import { preserveNativeFailedWorkspace } from "./native-runtime/native-failed-workspace-recovery.js"; import { nativeRunHasNoWorkspaceAccess } from "./workspace-restore-recovery-state.js"; @@ -2026,6 +2026,7 @@ export function heartbeatService( let githubLauncherLocation: Parameters[0] | null = null; let nativeSessionResumeScheduled = false; + let restartClaimLapsed = false; const shutdownHandedOff = watchShutdownHandoff(run.id); let nativeOwnershipHeld = false; let nativeStartupCleanupScheduled = false; @@ -3813,7 +3814,23 @@ export function heartbeatService( return; } if (await leaveToNextController(err)) return; - if (!(await preserveNativeWorkspaceOnFailure())) return; + if (!(await preserveNativeWorkspaceOnFailure())) { + const claim = runOptions.nativeRestartRecovery; + let releaseFailed = false; + restartClaimLapsed = claim !== undefined && await releaseLapsedNativeRestartClaim({ db, companyId: run.companyId, runId: run.id, + leaseOwner: claim.leaseOwner, controllerGeneration: claim.controllerGeneration, providerAttempt: claim.providerAttempt }) + .catch(releaseError => { + logger.warn({ err: releaseError, runId: run.id }, "lapsed native restart claim could not be released; restart recovery will reclaim it after expiry"); + releaseFailed = true; + return false; + }); + logger.warn({ runId: run.id, err: identityRedactor.redact({ message: err instanceof Error ? err.message : String(err), + stack: err instanceof Error ? err.stack : undefined }), restartClaim: claim?.kind ?? null, restartClaimLapsed }, + restartClaimLapsed ? "native restart claim lapsed during a failed preparation; released it for restart recovery" + : releaseFailed ? "native run failure left to restart recovery after its claim release failed" + : "native run failure left to the workspace owner that holds the run"); + return; + } await runCompletion.failRun({ run, agent, issueId, issueRef, signal: executionControl.controller.signal, output: runOutput, @@ -4129,6 +4146,7 @@ export function heartbeatService( } finally { controllerLease.stop(); activeRunExecutions.delete(run.id); + if (restartClaimLapsed) scheduleNativeSessionResumeDispatch(run.id, new Date(), "restart"); // A failed owned Stop remains visible until this exact executor settles, // including a graceful exit result arriving after the cancellation error. // It is never retained beyond the active execution's cleanup. diff --git a/server/src/services/native-runtime/native-execution-lease.ts b/server/src/services/native-runtime/native-execution-lease.ts index 851c608574a..015d6b822c8 100644 --- a/server/src/services/native-runtime/native-execution-lease.ts +++ b/server/src/services/native-runtime/native-execution-lease.ts @@ -145,6 +145,23 @@ export async function nativeRecoveryClaimRelinquished(input: { && row.lastDisposition === "restart_suspended"; } +export async function releaseLapsedNativeRestartClaim(input: { + db: Db; companyId: string; runId: string; leaseOwner: string; controllerGeneration: number; providerAttempt: number; +}): Promise { + const released = await input.db.update(nativeRunFinalizations) + .set({ leaseOwner: null, leaseExpiresAt: null, updatedAt: new Date() }) + .where(and(eq(nativeRunFinalizations.companyId, input.companyId), eq(nativeRunFinalizations.runId, input.runId), + eq(nativeRunFinalizations.phase, "observed"), isNull(nativeRunFinalizations.resultId), + eq(nativeRunFinalizations.leaseOwner, input.leaseOwner), + eq(nativeRunFinalizations.controllerGeneration, input.controllerGeneration), + eq(nativeRunFinalizations.attempt, input.providerAttempt), + sql`${nativeRunFinalizations.leaseExpiresAt} <= clock_timestamp()`, + sql`exists (select 1 from ${heartbeatRuns} where ${heartbeatRuns.id} = ${nativeRunFinalizations.runId} + and ${heartbeatRuns.companyId} = ${input.companyId} and ${heartbeatRuns.status} = 'running')`)) + .returning({ runId: nativeRunFinalizations.runId }); + return released.length > 0; +} + /** * Preparation failed after a fresh claim and before any provider turn. Release * only this exact claim, so ordinary failure settlement can retain the source