diff --git a/server/src/services/heartbeat/recovery.test.ts b/server/src/services/heartbeat/recovery.test.ts index a925edfaa7..02b3dd8397 100644 --- a/server/src/services/heartbeat/recovery.test.ts +++ b/server/src/services/heartbeat/recovery.test.ts @@ -1,9 +1,11 @@ import { randomUUID } from "node:crypto"; import type { Db } from "@paperclipai/db"; -import { agents, companies, createDb, EMBEDDED_POSTGRES_TEST_TIMEOUT_MS, heartbeatRuns, issues, nativeRunFinalizations } from "@paperclipai/db"; +import { agents, companies, createDb, EMBEDDED_POSTGRES_TEST_TIMEOUT_MS, heartbeatRuns, issues, nativeRunFinalizations, environments, environmentLeases } from "@paperclipai/db"; import { eq } from "drizzle-orm"; import { afterAll, beforeAll, beforeEach, describe, expect, it, vi } from "vitest"; import { getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase } from "../../__tests__/helpers/embedded-postgres.js"; +import { buildNativeExecutionInput } from "../native-runtime/native-execution-input.js"; +import { nativeRuntimeContextFixture } from "../native-runtime/runtime-context.test-fixture.js"; import { environmentRuntimeService } from "../environment-runtime.js"; import { environmentService } from "../environments.js"; import { createHeartbeatRecovery, type HeartbeatRecoveryDependencies } from "./recovery.js"; @@ -198,6 +200,118 @@ describePostgres("heartbeat recovery module ownership", () => { await db.delete(heartbeatRuns).where(eq(heartbeatRuns.id, run.id)); }); + async function seedComputerRun(mode: "warm" | "per_turn" = "warm") { + const run = await seedRun(); + const [environment] = await db.insert(environments).values({ name: `Computer ${run.id}`, driver: "computer" }).returning(); + const [lease] = await db.insert(environmentLeases).values({ companyId: run.companyId, environmentId: environment.id, + heartbeatRunId: run.id, issueId: run.issueId, provider: "boat", providerLeaseId: randomUUID(), + leasePolicy: "reuse_by_environment", metadata: { driver: "computer" } }).returning(); + const execution = buildNativeExecutionInput({ + companyId: run.companyId, runId: run.id, agentId: run.agentId, + issue: { id: run.issueId!, identifier: null, title: "Computer lifecycle", description: null, workMode: "standard" }, + taskPrompt: "Complete the turn", workspace: { id: randomUUID(), cwd: "/workspace", repoUrl: null, repoRef: null, branchName: null }, + normalizedSessionId: null, lifecyclePolicy: mode === "warm" ? { mode, idleTimeoutMs: 300_000 } : { mode, idleTimeoutMs: null }, + completionContract: { id: randomUUID(), sha256: `sha256:${"a".repeat(64)}`, schemaVersion: "paperclip.run-result.v1", + contract: { revision: "1", objective: "Complete the turn", criteria: [{ id: "output", requirement: "Complete the turn" }] } }, + runtimeContext: nativeRuntimeContextFixture(), + }); + const [settled] = await db.update(heartbeatRuns).set({ runtimeMode: "native", status: "succeeded", runnerProfileJson: { + nativeExecutionInput: execution, nativeComputerWorkspace: { kind: "remote-persistent", leaseId: lease.id, + providerLeaseId: lease.providerLeaseId, computerId: randomUUID(), ownerGeneration: 1, listenerPort: 47000, remoteCwd: "/workspace" }, + } }).where(eq(heartbeatRuns.id, run.id)).returning(); + // Current agent settings deliberately disagree with the frozen run policy. + await db.update(agents).set({ adapterConfig: { lifecycleMode: mode === "warm" ? "per_turn" : "warm" } }).where(eq(agents.id, run.agentId)); + return { run: settled, lease }; + } + + function settleThroughRecovery(deps: ReturnType, run: Run, startup = false) { + native.finalizations.mockImplementationOnce(async (_db, _ids, options) => { + await (options as { onWorkspaceSettled: (input: { runId: string; companyId: string; agentId: string; succeeded: boolean }) => Promise }) + .onWorkspaceSettled({ runId: run.id, companyId: run.companyId, agentId: run.agentId, succeeded: true }); + }); + const recovery = createHeartbeatRecovery(db, deps); + return startup ? recovery.recoverNativeRunsAfterRestart() : recovery.reapOrphanedRuns({ staleThresholdMs: 60_000 }); + } + + it("leaves an active computer executor's warm lease return to that executor", async () => { + const { run, lease } = await seedComputerRun(); + const deps = callbacks(db); + deps.getRun.mockResolvedValue(run); + deps.activeRunExecutions.add(run.id); + // Hold the executor's final lease return until the recovered-finalization + // callback has had its opportunity to incorrectly request physical Stop. + const stoppingRelease = vi.fn(async () => undefined); + deps.releaseEnvironmentLeasesForRun.mockImplementation(stoppingRelease); + await settleThroughRecovery(deps, run); + expect(stoppingRelease).not.toHaveBeenCalled(); + expect(deps.finalizeAgentStatus).not.toHaveBeenCalled(); + await deps.environmentRuntime.releaseRunLeases(run.id, "released", undefined, "keep_running"); + deps.activeRunExecutions.delete(run.id); + expect(await environmentService(db).getLeaseById(lease.id)).toMatchObject({ status: "retained", releasedAt: null }); + // A later recovery observation also uses the frozen warm policy, and cannot + // shorten the owner's deadline by issuing an unrelated Stop. + await settleThroughRecovery(deps, run); + expect(stoppingRelease).toHaveBeenCalledWith(expect.objectContaining({ providerResourceDisposition: "keep_running" })); + }); + + it.each([ + ["warm", false], ["per_turn", false], ["warm", true], ["per_turn", true], + ] as const)("recovers inactive computer cleanup using its frozen %s policy (startup %s)", async (mode, startup) => { + const { run } = await seedComputerRun(mode); + const deps = callbacks(db); + deps.getRun.mockResolvedValue(run); + await settleThroughRecovery(deps, run, startup); + expect(deps.releaseEnvironmentLeasesForRun).toHaveBeenCalledWith(expect.objectContaining({ + providerResourceDisposition: mode === "warm" ? "keep_running" : "stop_and_retain", + })); + expect(deps.getAgent).not.toHaveBeenCalled(); + }); + + it("rechecks computer executor ownership after the awaited run read", async () => { + const { run } = await seedComputerRun(); + const deps = callbacks(db); + deps.getRun.mockImplementation(async () => { deps.activeRunExecutions.add(run.id); return run; }); + await settleThroughRecovery(deps, run); + expect(deps.releaseEnvironmentLeasesForRun).not.toHaveBeenCalled(); + expect(deps.finalizeAgentStatus).not.toHaveBeenCalled(); + }); + + it.each(["missing", "malformed", "foreign-run", "foreign-company", "foreign-agent", "foreign-task", "noncomputer", "failed", "cancelled", "export-retry"])( + "does not earn recovered computer warmth from %s policy evidence", async variant => { + const { run } = await seedComputerRun(); + const profile = structuredClone(run.runnerProfileJson!); + const execution = profile.nativeExecutionInput as { binding: Record; session: { lifecyclePolicy: unknown } }; + if (variant === "missing") delete profile.nativeExecutionInput; + if (variant === "malformed") execution.session.lifecyclePolicy = { mode: "warm", idleTimeoutMs: -1 }; + if (variant.startsWith("foreign-")) execution.binding[({ "foreign-run": "runId", "foreign-company": "companyId", "foreign-agent": "agentId", "foreign-task": "issueId" } as Record)[variant]] = randomUUID(); + if (variant === "noncomputer") delete profile.nativeComputerWorkspace; + const candidate = { ...run, runnerProfileJson: profile, + ...(variant === "failed" || variant === "cancelled" ? { status: variant } : {}), + ...(variant === "export-retry" ? { resultJson: { workspaceExportRetry: true } } : {}), + }; + const deps = callbacks(db); + deps.getRun.mockResolvedValue(candidate); + await settleThroughRecovery(deps, candidate); + expect(deps.releaseEnvironmentLeasesForRun).toHaveBeenCalledWith(expect.objectContaining({ providerResourceDisposition: "stop_and_retain" })); + }, + ); + + it("preserves an active sandbox's delegated workspace-finalization cleanup", async () => { + const { run } = await seedComputerRun(); + const profile: Record = { ...run.runnerProfileJson, nativeWorkspaceSync: { + schema: "paperclip.native-workspace-sync/v2", state: "finalized", descriptorSha256: "a".repeat(64), baselineSha256: "b".repeat(64), + finalHostSha256: "c".repeat(64), workspaceId: randomUUID(), leaseId: randomUUID(), providerLeaseId: "sandbox", + remoteCwd: "/workspace", resourceDisposition: "destroy", + } }; + delete profile.nativeComputerWorkspace; + const deps = callbacks(db); + deps.getRun.mockResolvedValue({ ...run, runnerProfileJson: profile }); + deps.activeRunExecutions.add(run.id); + await settleThroughRecovery(deps, run); + expect(deps.releaseEnvironmentLeasesForRun).toHaveBeenCalledWith(expect.objectContaining({ providerResourceDisposition: "destroy" })); + expect(deps.finalizeAgentStatus).toHaveBeenCalledWith(run.agentId, "succeeded", null, expect.any(Object)); + }); + it("settles persisted native results before claiming restart authority and tracks the exact execution", async () => { const deps = callbacks(db); const claim: NativeRestartRecoveryClaim = { kind: "bootstrap_incomplete", runId: randomUUID(), leaseOwner: "claimed-owner", controllerGeneration: 3, providerAttempt: 2, restartKind: "hard", recoveryRequestId: null }; diff --git a/server/src/services/heartbeat/recovery.ts b/server/src/services/heartbeat/recovery.ts index f6a84bfba6..6b9abb2eb5 100644 --- a/server/src/services/heartbeat/recovery.ts +++ b/server/src/services/heartbeat/recovery.ts @@ -51,7 +51,8 @@ import { reconcileRetainedNativeSessionCleanup, reconcileRetainedNativeSessionCleanups, } from "../native-runtime/index.js"; -import { drainRetainedRunnerdMaintenanceOperations } from "../../vendor/paperclip-runner/index.js"; +import { readNativeComputerWorkspaceReference } from "../native-runtime/native-workspace-sync.js"; +import { drainRetainedRunnerdMaintenanceOperations, parseNativeExecutionInput } from "../../vendor/paperclip-runner/index.js"; import { runningProcesses } from "../../adapters/index.js"; import { parseObject } from "../../adapters/utils.js"; import { @@ -1838,9 +1839,30 @@ export function createHeartbeatRecovery(db: Db, dependencies: HeartbeatRecoveryD succeeded: boolean; }) { const settledRun = await getRun(input.runId); - const workspaceSyncReference = readNativeWorkspaceSyncReference( - parseObject(settledRun?.runnerProfileJson).nativeWorkspaceSync, - ); + const profile = parseObject(settledRun?.runnerProfileJson); + const workspaceSyncReference = readNativeWorkspaceSyncReference(profile.nativeWorkspaceSync); + let computerLifecycle: { mode: "warm" | "per_turn" } | null = null; + if (settledRun?.runtimeMode === "native" && + settledRun.companyId === input.companyId && settledRun.agentId === input.agentId && + !profile.nativeWorkspaceSync && !parseObject(settledRun.resultJson).workspaceExportRetry && + readNativeComputerWorkspaceReference(profile.nativeComputerWorkspace)) { + try { + const execution = parseNativeExecutionInput(profile.nativeExecutionInput); + if (execution.binding.runId === input.runId && execution.binding.companyId === input.companyId && + execution.binding.agentId === input.agentId && + execution.binding.issueId === (settledRun.nativeIssueId ?? settledRun.issueId)) { + computerLifecycle = execution.session.lifecyclePolicy; + } + } catch { /* Unreadable persisted policy cannot earn warm retention. */ } + } + // Persistent computer workspaces never hand off copy-back. The local + // executor owns this suffix until its finally releases the lease; recovery + // must not stop its warm session between terminalization and that release. + // Sandbox copy-back callbacks still own their explicitly delegated suffix. + if (computerLifecycle && activeRunExecutions.has(input.runId)) return; + const resourceDisposition = computerLifecycle?.mode === "warm" && settledRun?.status === "succeeded" + ? "keep_running" + : workspaceSyncReference?.resourceDisposition ?? "stop_and_retain"; await releaseEnvironmentLeasesForRun({ runId: input.runId, companyId: input.companyId, @@ -1848,7 +1870,7 @@ export function createHeartbeatRecovery(db: Db, dependencies: HeartbeatRecoveryD status: settledRun?.status, failureReason: settledRun?.error ?? undefined, providerResourceDisposition: input.succeeded && (!parseObject(settledRun?.resultJson).workspaceExportRetry || workspaceSyncReference?.resourceDisposition === "destroy") - ? (workspaceSyncReference?.resourceDisposition ?? "stop_and_retain") + ? resourceDisposition : "stop_and_retain", }); await releaseRuntimeServicesForRun(input.runId).catch(() => undefined);