diff --git a/server/src/__tests__/heartbeat-native-restart-suspension.test.ts b/server/src/__tests__/heartbeat-native-restart-suspension.test.ts index cdf63886b62..892cb063735 100644 --- a/server/src/__tests__/heartbeat-native-restart-suspension.test.ts +++ b/server/src/__tests__/heartbeat-native-restart-suspension.test.ts @@ -1,9 +1,12 @@ import { randomUUID } from "node:crypto"; import { and, eq } from "drizzle-orm"; -import { authUsers, budgetReservations, companyMemberships, heartbeatRuns, issues, nativeRunFinalizations, workspaceOperations } from "@paperclipai/db"; +import { authUsers, budgetReservations, companyMemberships, completionContracts, heartbeatRunEvents, heartbeatRuns, issues, nativeRunFinalizations, nativeRunResults, workspaceOperations } from "@paperclipai/db"; import { expect, it, vi } from "vitest"; import { heartbeatService } from "../services/heartbeat.js"; import { instanceSettingsService } from "../services/instance-settings.js"; +import { currentNativeControllerIdentity } from "../services/native-runtime/native-restart-recovery.js"; +import { PaperclipControlPlanePort } from "../services/native-runtime/paperclip-control-plane-port.js"; +import { CONTROL_PLANE_CONFORMANCE_OPEN, CONTROL_PLANE_CONFORMANCE_RESULT, CONTROL_PLANE_CONFORMANCE_TERMINAL } from "../vendor/paperclip-runner/testing.js"; import { describeEmbeddedPostgres, nativeRestartFixture } from "./helpers/native-restart.js"; @@ -92,4 +95,99 @@ describeEmbeddedPostgres("heartbeat native recovery interrupted by graceful shut const [run] = await f.db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, f.runId)); expect(run.status).toBe("failed"); }, 60_000); + + + async function shutdownDispositions(db: Awaited>["db"], runId: string) { + const events = await db.select().from(heartbeatRunEvents).where(and(eq(heartbeatRunEvents.runId, runId), + eq(heartbeatRunEvents.eventType, "native.recovery.transition"))); + return events.map((event) => (event.payload as Record | null)?.runnerDisposition); + } + + /** A run whose provider turn ended and whose result this boot's controller accepted. */ + async function seedAcceptedResult(label: string, leaseOwner: (runnerInstanceId: string) => string) { + const f = await seedRecoverableRun(label); + const [contract] = await f.db.select().from(completionContracts).where(eq(completionContracts.issueId, f.issueId)); + await f.db.update(heartbeatRuns).set({ completionContractId: contract.id, + completionContractSha256: contract.canonicalSha256 }).where(eq(heartbeatRuns.id, f.runId)); + const controller = await currentNativeControllerIdentity(); + const leaseExpiresAt = new Date(Date.now() + 20 * 60_000); + await f.db.update(nativeRunFinalizations).set({ + leaseOwner: leaseOwner(f.native.runnerInstanceId), + leaseExpiresAt, + controllerBootId: controller.bootId, + controllerPid: controller.pid, + controllerProcessStartedAt: controller.processStartedAt, + controllerGeneration: 1, + }).where(eq(nativeRunFinalizations.runId, f.runId)); + const [prepared] = await f.db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, f.runId)); + const identity = { companyId: fixture.current.companyId, issueId: f.issueId, runId: f.runId, + agentId: fixture.current.agentId, sessionId: prepared.nativeSessionId! }; + const port = new PaperclipControlPlanePort(f.db, { + ...identity, + completionContractId: prepared.completionContractId!, + completionContractSha256: prepared.completionContractSha256!, + sourceInstanceId: prepared.runnerInstanceId!, + controlPlaneSourceInstanceId: `${label}-control`, + }); + await port.openRun({ ...CONTROL_PLANE_CONFORMANCE_OPEN, identity, sourceInstanceId: prepared.runnerInstanceId! }); + const { revision, criterionIds } = f.native.completionContract; + await port.completeRun({ + result: { ...CONTROL_PLANE_CONFORMANCE_RESULT, completionClaim: { + ...CONTROL_PLANE_CONFORMANCE_RESULT.completionClaim, contractRevision: revision, + criteria: criterionIds.map((criterionId) => ({ criterionId, status: "satisfied" as const, evidenceRefs: ["event:2"] })), + } }, + terminal: CONTROL_PLANE_CONFORMANCE_TERMINAL, + callerResultId: `${label}-result`, + }); + // The detached controller had begun local workspace finalization. + await f.db.insert(workspaceOperations).values({ companyId: fixture.current.companyId, issueId: f.issueId, + heartbeatRunId: f.runId, phase: "workspace_finalize", status: "running", cwd: fixture.current.runtimeRoot }); + const [accepted] = await f.db.select().from(nativeRunFinalizations).where(eq(nativeRunFinalizations.runId, f.runId)); + expect(accepted).toMatchObject({ phase: "workspace_finalizing", resultId: expect.any(String) }); + return { ...f, leaseExpiresAt }; + } + + it("finalizes an accepted terminal result whose controller was suspended by a graceful shutdown", async () => { + // RC-30: SIGTERM lands after run.result.accepted/run.terminal but before the + // controller finalizes. The run must not stay `running` behind the old + // controller's execution lease once the next boot starts recovery. + const f = await seedAcceptedResult("ACCEPTED-THEN-SIGTERM", (runner) => `${runner}:${randomUUID()}`); + const dying = heartbeatService(f.db); + const drained = await dying.drainRunningRunsForShutdown("SIGTERM", new Date(), [f.runId]); + expect(drained).toMatchObject({ restartSuspendedRunIds: [f.runId] }); + const [released] = await f.db.select().from(nativeRunFinalizations).where(eq(nativeRunFinalizations.runId, f.runId)); + // Finalization, not a runner reattach, is what remains. + expect(released.recoveryState).toBeNull(); + expect(await shutdownDispositions(f.db, f.runId)).toEqual(["awaiting_result_finalization"]); + + const providerDispatch = vi.fn(() => { throw new Error("an accepted result must not open another provider turn"); }); + const next = heartbeatService(f.db, { nativeSessionBackendFactory: providerDispatch }); + try { + await next.recoverNativeRunsAfterRestart(new Date()); + } finally { await next.drainActiveRunExecutions(); } + expect(providerDispatch).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, phase: owner.phase }).toEqual({ status: "succeeded", phase: "committed" }); + expect(task.status).toBe("done"); + expect(await f.db.select({ id: nativeRunResults.id }).from(nativeRunResults) + .where(eq(nativeRunResults.runId, f.runId))).toHaveLength(1); + expect(await f.db.select({ id: heartbeatRuns.id }).from(heartbeatRuns) + .where(eq(heartbeatRuns.nativeIssueId, f.issueId))).toHaveLength(1); + }, 90_000); + + it("keeps a finalizer lease already settling the accepted result through a graceful shutdown", async () => { + // The shutdown drain waits for in-process finalizers; releasing their + // lease would let the next controller finalize the same result concurrently. + const f = await seedAcceptedResult("FINALIZER-THEN-SIGTERM", () => `native-finalizer:${randomUUID()}`); + const dying = heartbeatService(f.db); + await dying.drainRunningRunsForShutdown("SIGTERM", new Date(), [f.runId]); + const [owner] = await f.db.select().from(nativeRunFinalizations).where(eq(nativeRunFinalizations.runId, f.runId)); + expect(owner).toMatchObject({ phase: "workspace_finalizing", leaseExpiresAt: f.leaseExpiresAt, recoveryState: null }); + expect(await shutdownDispositions(f.db, f.runId)).toEqual(["finalizer_in_progress"]); + const [run] = await f.db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, f.runId)); + expect(run.status).toBe("running"); + }, 60_000); }); diff --git a/server/src/services/heartbeat/recovery.ts b/server/src/services/heartbeat/recovery.ts index c3501633098..3d5695339bc 100644 --- a/server/src/services/heartbeat/recovery.ts +++ b/server/src/services/heartbeat/recovery.ts @@ -1188,11 +1188,18 @@ export function createHeartbeatRecovery(db: Db, dependencies: HeartbeatRecoveryD processPid: run.processPid, processStartedAt: run.processStartedAt?.toISOString() ?? null, }); - await db.transaction(async tx => { + // An accepted result needs finalization, not a runner reattach. Its + // execution lease must still be released: the detached executor will + // not finalize, and the next controller's reconciler skips leased rows. + // A finalizer lease belongs to finalization already in progress here. + const released = await db.transaction(async tx => { await notifyLiveRunRecovery(tx, { companyId: run.companyId, runId: run.id }); - await tx.update(nativeRunFinalizations).set({ - leaseExpiresAt: now, - recoveryState: "awaiting_runner_reattach", + return tx.update(nativeRunFinalizations).set({ + leaseExpiresAt: sql`case when ${nativeRunFinalizations.resultId} is not null + and ${nativeRunFinalizations.leaseOwner} like 'native-finalizer:%' + then ${nativeRunFinalizations.leaseExpiresAt} else ${now.toISOString()}::timestamptz end`, + recoveryState: sql`case when ${nativeRunFinalizations.resultId} is null + then 'awaiting_runner_reattach' else ${nativeRunFinalizations.recoveryState} end`, recoveryHistory: sql`( select coalesce(jsonb_agg(item order by ordinal), '[]'::jsonb) from jsonb_array_elements( @@ -1214,20 +1221,26 @@ export function createHeartbeatRecovery(db: Db, dependencies: HeartbeatRecoveryD eq(nativeRunFinalizations.runId, run.id), owner?.bootId ? eq(nativeRunFinalizations.controllerBootId, owner.bootId) : isNull(nativeRunFinalizations.controllerBootId), - isNull(nativeRunFinalizations.resultId), ), - ); + ) + .returning({ resultId: nativeRunFinalizations.resultId, leaseOwner: nativeRunFinalizations.leaseOwner }); }); + const resultAccepted = Boolean(released[0]?.resultId); + const finalizerHeld = resultAccepted && Boolean(released[0]?.leaseOwner?.startsWith("native-finalizer:")); await appendRunEvent(run, { eventType: "native.recovery.transition", stream: "system", level: "info", - message: - "Server shutdown suspended native controller ownership without cancelling provider work", + message: finalizerHeld + ? "Server shutdown kept the in-progress finalizer lease of the accepted native result" + : resultAccepted + ? "Server shutdown released the accepted native result for finalization by the next controller" + : "Server shutdown suspended native controller ownership without cancelling provider work", payload: { restartKind: "graceful", signal, - runnerDisposition: "awaiting_runner_reattach", + runnerDisposition: finalizerHeld ? "finalizer_in_progress" + : resultAccepted ? "awaiting_result_finalization" : "awaiting_runner_reattach", processPid: run.processPid, processGroupId: run.processGroupId, processStartedAt: run.processStartedAt?.toISOString() ?? null,