diff --git a/doc/execution-semantics.md b/doc/execution-semantics.md index eda4b250cac..eb7f8cff4d3 100644 --- a/doc/execution-semantics.md +++ b/doc/execution-semantics.md @@ -3285,6 +3285,10 @@ waiting for provider startup. An abrupt controller exit during session opening can then recover through the same exact process-identity checks as an active turn. Both graceful and hot restarts detach the old controller from native sessions. +A graceful shutdown can also suspend a restart-recovery claim whose execution +has not yet reached provider dispatch. That execution has lost its authority: +it leaves the run running, its budget hold held and its workspace unmarked, and +the next controller adopts the same run. It does not settle the run as failed. Remote Runner processes use their authenticated connection lease as the bound on disconnected lifetime. The controller's shorter connection-recovery budget limits its own attempt; it must not terminate the remote owner while another diff --git a/server/src/__tests__/heartbeat-native-restart-suspension.test.ts b/server/src/__tests__/heartbeat-native-restart-suspension.test.ts new file mode 100644 index 00000000000..cdf63886b62 --- /dev/null +++ b/server/src/__tests__/heartbeat-native-restart-suspension.test.ts @@ -0,0 +1,95 @@ +import { randomUUID } from "node:crypto"; +import { and, eq } from "drizzle-orm"; +import { authUsers, budgetReservations, companyMemberships, heartbeatRuns, issues, nativeRunFinalizations, workspaceOperations } from "@paperclipai/db"; +import { expect, it, vi } from "vitest"; +import { heartbeatService } from "../services/heartbeat.js"; +import { instanceSettingsService } from "../services/instance-settings.js"; +import { describeEmbeddedPostgres, nativeRestartFixture } from "./helpers/native-restart.js"; + + +describeEmbeddedPostgres("heartbeat native recovery interrupted by graceful shutdown", () => { + const fixture = nativeRestartFixture(); + + async function seedRecoverableRun(label: string) { + const f = await fixture.seedRun(label); + 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)); + return f; + } + + it.each(["before_budget_reservation", "before_dispatch_gate"] as const)( + "leaves a recovery execution suspended %s for the next controller", + async (point) => { + const f = await seedRecoverableRun(`SUSPENDED-${point}`); + if (point === "before_dispatch_gate") await f.db.update(heartbeatRuns) + .set({ scheduledRetryReason: "native_provider_overloaded" }).where(eq(heartbeatRuns.id, f.runId)); + const dying = vi.fn(() => { throw new Error("A suspended recovery must not open a provider session"); }); + let drained: unknown; + const drain = async () => { + if (drained) return; + drained = await heartbeat.drainRunningRunsForShutdown("SIGTERM", new Date(), [f.runId]); + }; + const heartbeat = heartbeatService(f.db, { + nativeSessionBackendFactory: dying, + ...(point === "before_budget_reservation" + ? { beforeNativeRuntimeSelection: drain } + : { afterResolvedInteractionContinuationDispatchCheck: drain }), + }); + try { + const { claims } = await heartbeat.recoverNativeRunsAfterRestart(new Date(), [f.runId]); + expect(claims).toEqual([expect.objectContaining({ runId: f.runId })]); + } finally { await heartbeat.drainActiveRunExecutions(); } + expect(drained).toMatchObject({ restartSuspendedRunIds: [f.runId] }); + expect(dying).not.toHaveBeenCalled(); + + const [run] = await f.db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, f.runId)); + const [owner] = await f.db.select().from(nativeRunFinalizations).where(eq(nativeRunFinalizations.runId, f.runId)); + const [task] = await f.db.select().from(issues).where(eq(issues.id, f.issueId)); + expect({ status: run.status, errorCode: run.errorCode, finishedAt: run.finishedAt }).toEqual({ + status: "running", errorCode: null, finishedAt: null }); + expect(owner).toMatchObject({ phase: "observed", resultId: null, recoveryState: "awaiting_runner_reattach" }); + expect(task.executionRunId).toBe(f.runId); + expect(await f.db.select().from(workspaceOperations).where(and( + eq(workspaceOperations.heartbeatRunId, f.runId), eq(workspaceOperations.status, "failed")))).toEqual([]); + expect(await f.db.select().from(budgetReservations).where(and(eq(budgetReservations.runId, f.runId), + eq(budgetReservations.state, "released")))).toEqual([]); + await f.db.update(nativeRunFinalizations).set({ controllerPid: 2_000_000_001 }) + .where(eq(nativeRunFinalizations.runId, f.runId)); + const resumed = vi.fn(() => { throw new Error("next boot reached provider dispatch"); }); + const next = heartbeatService(f.db, { nativeSessionBackendFactory: resumed }); + try { + const { claims } = await next.recoverNativeRunsAfterRestart(new Date(), [f.runId]); + expect(claims).toEqual([expect.objectContaining({ runId: f.runId })]); + } finally { await next.drainActiveRunExecutions(); } + expect(resumed).toHaveBeenCalledOnce(); + expect(await f.db.select({ id: heartbeatRuns.id }).from(heartbeatRuns) + .where(eq(heartbeatRuns.nativeIssueId, f.issueId))).toHaveLength(1); + }, + 90_000, + ); + + it("still settles a recovery whose claim merely expired without a shutdown suspension", async () => { + const f = await seedRecoverableRun("EXPIRED-NOT-SUSPENDED"); + const dying = vi.fn(() => { throw new Error("An expired claim must not open a provider session"); }); + const heartbeat = heartbeatService(f.db, { + nativeSessionBackendFactory: dying, + beforeNativeRuntimeSelection: async () => { + await f.db.update(nativeRunFinalizations).set({ leaseExpiresAt: new Date(0) }) + .where(eq(nativeRunFinalizations.runId, f.runId)); + }, + }); + try { + const { claims } = await heartbeat.recoverNativeRunsAfterRestart(new Date(), [f.runId]); + expect(claims).toEqual([expect.objectContaining({ runId: f.runId })]); + } finally { await heartbeat.drainActiveRunExecutions(); } + expect(dying).not.toHaveBeenCalled(); + const [run] = await f.db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, f.runId)); + expect(run.status).toBe("failed"); + }, 60_000); +}); diff --git a/server/src/services/heartbeat.ts b/server/src/services/heartbeat.ts index 4f54e6d2915..574d08ee974 100644 --- a/server/src/services/heartbeat.ts +++ b/server/src/services/heartbeat.ts @@ -43,6 +43,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 { 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"; @@ -2924,19 +2925,28 @@ export function heartbeatService( const dispatchResolvedInteractionContinuationWithAtomicGate = async ( dispatch: (markDispatchStarted: () => void) => Promise, ) => { + const relinquished = async () => runOptions.nativeLeaseOwner !== undefined && await nativeRecoveryClaimRelinquished({ + db, companyId: run.companyId, runId: run.id, leaseOwner: runOptions.nativeLeaseOwner }); // Admission can wait for accounting locks. Finish that wait before the // final ownership gate, whose callback must enter the adapter directly. const reservation = await measureRunPhase("execution.budget_reservation", () => reserveRunBudget(db, run.companyId, run.id, - readNonEmptyString(runLedgerScope.projectId), runLedgerScope, runOptions.nativeLeaseOwner)); - let entered = false; + readNonEmptyString(runLedgerScope.projectId), runLedgerScope, runOptions.nativeLeaseOwner)) + .catch(async (error: unknown) => { throw await relinquished() ? new NativeControllerDetachedForRestartError() : error; }); + let entered = false, handedOff = false; try { return await dispatchResolvedInteractionContinuationAfterAdmission((markDispatchStarted) => { entered = true; return dispatch(markDispatchStarted); }, reservation); + } catch (error) { + if (!entered && await relinquished()) { + handedOff = true; + throw new NativeControllerDetachedForRestartError(); + } + throw error; } finally { - if (!entered && !reservation.reused) { + if (!entered && !reservation.reused && !handedOff) { await db.update(heartbeatRuns).set({ costAccountingPending: true, usageJson: sql`coalesce(${heartbeatRuns.usageJson}, '{}'::jsonb) || '{"accountingProviderWorkStarted":false}'::jsonb`, }).where(and(eq(heartbeatRuns.id, run.id), isNull(heartbeatRuns.costAccountedAt))); diff --git a/server/src/services/native-runtime/native-execution-lease.ts b/server/src/services/native-runtime/native-execution-lease.ts index 4e17b04eccf..98c95766ef5 100644 --- a/server/src/services/native-runtime/native-execution-lease.ts +++ b/server/src/services/native-runtime/native-execution-lease.ts @@ -30,6 +30,23 @@ export async function assertNativeSessionPreparationAuthority(input: { if (!owned) throw new Error("native_session_lease_lost"); } +export async function nativeRecoveryClaimRelinquished(input: { + db: Db; companyId: string; runId: string; leaseOwner: string; +}): Promise { + const [row] = await input.db.select({ + status: heartbeatRuns.status, resultId: nativeRunFinalizations.resultId, + leaseOwner: nativeRunFinalizations.leaseOwner, recoveryState: nativeRunFinalizations.recoveryState, + leaseLive: sql`coalesce(${nativeRunFinalizations.leaseExpiresAt} > clock_timestamp(), false)`, + lastDisposition: sql`${nativeRunFinalizations.recoveryHistory} -> -1 ->> 'disposition'`, + }).from(nativeRunFinalizations).innerJoin(heartbeatRuns, eq(heartbeatRuns.id, nativeRunFinalizations.runId)) + .where(and(eq(nativeRunFinalizations.companyId, input.companyId), eq(nativeRunFinalizations.runId, input.runId), + eq(heartbeatRuns.companyId, input.companyId))).limit(1); + if (!row || row.status !== "running" || row.resultId) return false; + if (row.leaseOwner !== null && row.leaseOwner !== input.leaseOwner) return true; + return row.leaseOwner === input.leaseOwner && !row.leaseLive && row.recoveryState === "awaiting_runner_reattach" + && row.lastDisposition === "restart_suspended"; +} + /** * Preparation failed after a fresh claim and before any provider turn. Release * only this exact claim, so ordinary failure settlement can retain the source