diff --git a/server/src/__tests__/heartbeat-native-restart-suspension.test.ts b/server/src/__tests__/heartbeat-native-restart-suspension.test.ts index 49a7dd3385f..9a5b1cf6209 100644 --- a/server/src/__tests__/heartbeat-native-restart-suspension.test.ts +++ b/server/src/__tests__/heartbeat-native-restart-suspension.test.ts @@ -35,6 +35,30 @@ import { createDeliveryWorkCoordinator } from "../services/delivery-work-coordin import { reconcileNativeFinalizations, registerOwnershipReconciliationWork } from "../services/native-runtime/native-finalization-reconciler.js"; import { OWNABLE_KINDS } from "../services/recovery/ownership-registry.js"; import { currentProcessProbe, proveOwnerGone } from "../services/owner-proof.js"; +import { createRunLogWriterAuthority } from "../services/heartbeat/run-log-authority.js"; +import { resolveHotRestartIntentPath, writeHotRestartIntent } from "../services/hot-restart.js"; +import fs from "node:fs/promises"; +import { detachNativeSessionsForRestart } from "../services/native-runtime/native-session-executor.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("heartbeat native recovery interrupted by graceful shutdown", () => { @@ -251,6 +275,79 @@ describeEmbeddedPostgres("heartbeat native recovery interrupted by graceful shut expect(resumed).toHaveBeenCalledOnce(); }, 90_000); + it.each([ + { mode: "drain", stage: "run_output" }, + { mode: "drain", stage: "task_context" }, + { mode: "hot_restart", stage: "run_output" }, + { mode: "lease_expired", stage: "run_output" }, + { mode: "restart_detach_only", stage: "run_output" }, + { mode: "begin_failed", stage: "run_output" }, + ] as const)("leaves a preparing recovery that lost run-log writer authority to the next controller only after a drain ($mode, $stage)", async ({ mode, stage }) => { + const f = await seedRecoverableRun(`PREPARING-AUTHORITY-${mode}-${stage}`); + const dying = vi.fn(() => { throw new Error("A preparing recovery must not open a provider session"); }); + let drained: unknown; + const loseAuthority = async () => { + if (mode === "drain") drained = await heartbeat.drainRunningRunsForShutdown("SIGTERM", new Date(), [f.runId]); + else if (mode === "hot_restart") { + await writeHotRestartIntent({ previousServerPid: process.pid, previousServerVersion: "previous" }); + try { + expect(await heartbeat.prepareHotRestartShutdown("SIGTERM")).toMatchObject({ mode: "hot_restart", skipDrain: true }); + } finally { await fs.rm(resolveHotRestartIntentPath(), { force: true }); } + await f.db.update(nativeRunFinalizations).set({ leaseExpiresAt: new Date(0) }) + .where(eq(nativeRunFinalizations.runId, f.runId)); + } else if (mode === "restart_detach_only") { + await detachNativeSessionsForRestart([f.runId]); + await f.db.update(nativeRunFinalizations).set({ leaseExpiresAt: new Date(0) }) + .where(eq(nativeRunFinalizations.runId, f.runId)); + } else await f.db.update(nativeRunFinalizations).set({ leaseExpiresAt: new Date(0) }) + .where(eq(nativeRunFinalizations.runId, f.runId)); + const [admitted] = await f.db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, f.runId)); + await createRunLogWriterAuthority(f.db, admitted)(); + }; + if (stage === "task_context") preparationFault.current = { runId: f.runId, phase: stage, act: loseAuthority }; + if (mode === "begin_failed") preparationFault.current = { runId: f.runId, phase: "run_output", + act: async () => { throw new Error("Object storage is unavailable"); } }; + const heartbeat = heartbeatService(f.db, { + nativeSessionBackendFactory: dying, + ...(stage === "run_output" && mode !== "begin_failed" ? { beforeNativeRuntimeSelection: loseAuthority } : {}), + }); + try { + const { claims } = await heartbeat.recoverNativeRunsAfterRestart(new Date(), [f.runId]); + expect(claims).toEqual([expect.objectContaining({ runId: f.runId })]); + } finally { + await heartbeat.drainActiveRunExecutions(); + preparationFault.current = null; + } + 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 handoffEvents = await f.db.select().from(heartbeatRunEvents).where(and(eq(heartbeatRunEvents.runId, f.runId), + eq(heartbeatRunEvents.eventType, "lifecycle"), sql`${heartbeatRunEvents.payload}->>'disposition' = 'left_to_next_controller'`)); + if (mode !== "drain" && mode !== "hot_restart") { + expect(run.status).toBe("failed"); + expect(handoffEvents).toEqual([]); + return; + } + expect({ status: run.status, errorCode: run.errorCode, finishedAt: run.finishedAt }).toEqual({ + status: "running", errorCode: null, finishedAt: null }); + expect(handoffEvents).toHaveLength(1); + if (mode === "hot_restart") return; + expect(drained).toMatchObject({ restartSuspendedRunIds: [f.runId] }); + expect(owner).toMatchObject({ phase: "observed", resultId: null, recoveryState: "awaiting_runner_reattach" }); + expect(handoffEvents).toEqual([expect.objectContaining({ payload: expect.objectContaining({ + errorCode: "run_log_writer_authority_lost", processPid: process.pid }) })]); + + 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(); + }, 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"); }); diff --git a/server/src/services/heartbeat.ts b/server/src/services/heartbeat.ts index dbb1836d81b..9c4365abde7 100644 --- a/server/src/services/heartbeat.ts +++ b/server/src/services/heartbeat.ts @@ -44,6 +44,7 @@ import { createHeartbeatInstructionCopyRelease, createHeartbeatInstructionWorkin import { prepareHeartbeatRuntimeServiceIntents, startHeartbeatRuntimeServices, recordHeartbeatAdapterRuntimeServices } from "./heartbeat/runtime-services.js"; import { bindHeartbeatRunnerHomeMode, prepareHeartbeatExecutionConfiguration } from "./heartbeat/execution-configuration.js"; import { createRunLogWriterAuthority } from "./heartbeat/run-log-authority.js"; +import { watchShutdownHandoff } from "./heartbeat/shutdown-handoff.js"; import { createHeartbeatRunOutput } from "./heartbeat/run-output.js"; import { readRecoverableRunLog } from "./native-runtime/native-run-log-recovery.js"; import { observabilitySecretValues } from "@paperclipai/adapter-utils/agent-observability"; @@ -2025,6 +2026,7 @@ export function heartbeatService( let githubLauncherLocation: Parameters[0] | null = null; let nativeSessionResumeScheduled = false; + const shutdownHandedOff = watchShutdownHandoff(run.id); let nativeOwnershipHeld = false; let nativeStartupCleanupScheduled = false; let nativeStartupCleanupError: NativeStartupCleanupPendingError | undefined; @@ -2041,6 +2043,21 @@ export function heartbeatService( }; let instructionWorkingCopy: ReturnType | null = null; let nativeDispatchStarted = false; + const leaveToNextController = async (error: unknown) => { + if (!shutdownHandedOff() || legacyAdapterEntered || nativeDispatchStarted) return false; + nativeSessionResumeScheduled = true; + const reason = error instanceof Error ? error.message : String(error); + const errorCode = /^[a-z][a-z0-9_]{0,80}$/.test(reason) ? reason : null; + logger.info({ runId: run.id, err: reason }, "run preparation stopped after a graceful shutdown handed the run to the next controller"); + await appendRunEvent(run, { + eventType: "lifecycle", + stream: "system", + level: "info", + message: "Run preparation stopped; a graceful shutdown handed the run to the next controller", + payload: { disposition: "left_to_next_controller", errorCode, processPid: process.pid }, + }).catch(() => undefined); + return true; + }; let nativeWorkspaceFinalizeScheduled = false; let nativeWorkspaceSync: Awaited< ReturnType @@ -3795,6 +3812,7 @@ export function heartbeatService( } return; } + if (await leaveToNextController(err)) return; if (!(await preserveNativeWorkspaceOnFailure())) return; await runCompletion.failRun({ run, agent, issueId, issueRef, signal: executionControl.controller.signal, @@ -3852,7 +3870,7 @@ export function heartbeatService( ); }, ); - } else { + } else if (!(await leaveToNextController(outerErr))) { await preserveNativeWorkspaceOnFailure(); if (!nativeOwnershipHeld && !nativeStartupCleanupScheduled && !nativeStartupCleanupError) { await runCompletion.failRunSetup({ diff --git a/server/src/services/heartbeat/recovery.ts b/server/src/services/heartbeat/recovery.ts index a4dedbf0471..18798f49eeb 100644 --- a/server/src/services/heartbeat/recovery.ts +++ b/server/src/services/heartbeat/recovery.ts @@ -1,4 +1,5 @@ import { settleOwnableKind } from "../recovery/ownership-invariant-report.js"; +import { markRunHandedOffForShutdown } from "./shutdown-handoff.js"; import { ORPHANED_RUN_LEASE_FAILURE_REASON } from "../environment-run-lease-settlement.js"; import { ENVIRONMENT_LEASE_CONTROLLER_KEY } from "../environment-lease-controller.js"; import { readNativeStartupCleanup, readNativeStartupContext } from "../native-runtime/native-startup-cleanup.js"; @@ -597,6 +598,7 @@ export function createHeartbeatRecovery(db: Db, dependencies: HeartbeatRecoveryD isNativeSessionId(run.nativeSessionId), ) .map(({ run }) => run.id); + for (const runId of nativeRunIds) markRunHandedOffForShutdown(runId); const detachedNativeSessions = await detachNativeSessionsForRestart(nativeRunIds); @@ -1190,6 +1192,7 @@ export function createHeartbeatRecovery(db: Db, dependencies: HeartbeatRecoveryD // hot restart. Leaving the old event consumer attached lets its // finalizer interrupt/suspend Claude while the next server is adopting // the same turn. + markRunHandedOffForShutdown(run.id); await detachNativeSessionsForRestart([run.id]); const recoveryHistoryEntry = JSON.stringify({ at: now.toISOString(), diff --git a/server/src/services/heartbeat/shutdown-handoff.ts b/server/src/services/heartbeat/shutdown-handoff.ts new file mode 100644 index 00000000000..57df43b2cef --- /dev/null +++ b/server/src/services/heartbeat/shutdown-handoff.ts @@ -0,0 +1,13 @@ +const runsHandedOffForShutdown = new Map(); + +export function markRunHandedOffForShutdown(runId: string) { + runsHandedOffForShutdown.set(runId, Symbol("shutdown handoff")); +} + +export function watchShutdownHandoff(runId: string): () => boolean { + const atStart = runsHandedOffForShutdown.get(runId); + return () => { + const current = runsHandedOffForShutdown.get(runId); + return current !== undefined && current !== atStart; + }; +} diff --git a/server/src/services/native-runtime/native-session-executor.ts b/server/src/services/native-runtime/native-session-executor.ts index aad0b18f0aa..cccbcee85df 100644 --- a/server/src/services/native-runtime/native-session-executor.ts +++ b/server/src/services/native-runtime/native-session-executor.ts @@ -382,6 +382,7 @@ export class NativeControllerDetachedForRestartError extends Error { const activeNativeSessions = new Map(); // Shutdown can race provider startup before onSession publishes its handle. const nativeRunsDetachingForRestart = new Map(); + /** Flushes a run's pending restart-recovery budget check when the controller * detaches it for another restart (its trace never finishes). */ const nativeRunRecoveryFlushers = new Map void>();