diff --git a/server/src/__tests__/heartbeat-orphaned-active-lease-sweep.test.ts b/server/src/__tests__/heartbeat-orphaned-active-lease-sweep.test.ts index 6ff69f16f0a..5d5b65277e5 100644 --- a/server/src/__tests__/heartbeat-orphaned-active-lease-sweep.test.ts +++ b/server/src/__tests__/heartbeat-orphaned-active-lease-sweep.test.ts @@ -5,6 +5,7 @@ import { costEvents, agents, companies, + computers, createDb, environmentLeases, environments, @@ -65,6 +66,7 @@ describeEmbeddedPostgres("heartbeat sweepOrphanedActiveLeases", () => { await db.delete(environmentLeases); await db.delete(costEvents); await db.delete(heartbeatRuns); + await db.delete(computers); await db.delete(environments); await db.delete(agents); await db.delete(companies); @@ -98,7 +100,7 @@ describeEmbeddedPostgres("heartbeat sweepOrphanedActiveLeases", () => { await db.insert(environments).values({ id: environmentId, companyId, - name: "Fake Sandbox", + name: `Fake Sandbox ${environmentId}`, driver: "sandbox", status: "active", config: { provider: "fake", image: "ubuntu:24.04" }, @@ -160,6 +162,104 @@ describeEmbeddedPostgres("heartbeat sweepOrphanedActiveLeases", () => { const oldEnough = () => new Date(Date.now() - 60 * 60 * 1000); + async function seedWarmComputerLease() { + const scope = await seedCompanyAgentAndEnvironment(); + await db.update(environments).set({ driver: "computer" }).where(eq(environments.id, scope.environmentId)); + const runId = await insertHeartbeatRun({ ...scope, status: "succeeded" }); + const computerId = randomUUID(), ownerId = randomUUID(); + const owner = { id: ownerId, generation: 1, kind: "runner", phase: "warm", agentId: scope.agentId, + runId, admittedAt: oldEnough().toISOString(), deadline: new Date(Date.now() + 300_000).toISOString(), + absoluteDeadline: null, retirementDeadline: null, process: { bootId: randomUUID(), unitName: "test", nonce: "test", launchGeneration: 1 } }; + await db.insert(computers).values({ id: computerId, companyId: scope.companyId, environmentId: scope.environmentId, + provider: "boat", providerId: randomUUID(), ledger: { status: "attached", owners: [owner], placements: {}, action: null } }); + const leaseId = await insertActiveLease({ ...scope, heartbeatRunId: runId, updatedAt: oldEnough(), + provider: "boat", providerLeaseId: ownerId }); + await db.update(environmentLeases).set({ metadata: { driver: "computer", agentId: scope.agentId, + computerOwner: { computerId, ownerId, generation: 1 } } }).where(eq(environmentLeases.id, leaseId)); + return { ...scope, runId, computerId, ownerId, leaseId, owner }; + } + + it("preserves an authoritative warm computer while successful-run instruction cleanup delays lease retention", async () => { + const fixture = await seedWarmComputerLease(); + const destroyRunLease = vi.fn(); + const heartbeat = heartbeatService(db, { environmentRuntime: { destroyRunLease } as unknown as HeartbeatEnvironmentRuntime }); + expect(await heartbeat.sweepOrphanedActiveLeases({ backoffMs: 300_000 })).toEqual({ recovered: 0 }); + await heartbeat.sweepPendingCleanupLeases({ backoffMs: 0 }); + expect(destroyRunLease).not.toHaveBeenCalled(); + expect(await leaseRow(fixture.leaseId)).toMatchObject({ status: "active", failureReason: null }); + // The normal heartbeat can finish its delayed cleanup and retain the lease. + await db.update(environmentLeases).set({ status: "retained" }).where(eq(environmentLeases.id, fixture.leaseId)); + expect(await heartbeat.sweepOrphanedActiveLeases({ backoffMs: 0 })).toEqual({ recovered: 0 }); + }); + + it.each(["expired", "retired", "foreign_run", "foreign_agent", "generation", "computer", "environment", "company", "provider_lease", "invalid_deadline", "missing_process", "invalid_process", "invalid_launch_generation"])( + "still recovers a computer lease without exact bounded warm authority: %s", async (kind) => { + const fixture = await seedWarmComputerLease(); + const owner = { ...fixture.owner }; + if (kind === "expired") owner.deadline = new Date(Date.now() - 60_000).toISOString(); + if (kind === "retired") owner.phase = "retired"; + if (kind === "foreign_run") owner.runId = randomUUID(); + if (kind === "foreign_agent") owner.agentId = randomUUID(); + if (kind === "generation") owner.generation += 1; + if (kind === "invalid_deadline") owner.deadline = "invalid"; + if (kind === "missing_process") Object.assign(owner, { process: null }); + if (kind === "invalid_process") owner.process = { ...owner.process, nonce: "" }; + if (kind === "invalid_launch_generation") owner.process = { ...owner.process, launchGeneration: 0 }; + await db.update(computers).set({ ledger: { status: "attached", owners: [owner] } }).where(eq(computers.id, fixture.computerId)); + if (kind === "computer") await db.update(environmentLeases).set({ metadata: { driver: "computer", agentId: fixture.agentId, + computerOwner: { computerId: randomUUID(), ownerId: fixture.ownerId, generation: 1 } } }).where(eq(environmentLeases.id, fixture.leaseId)); + if (kind === "provider_lease") await db.update(environmentLeases).set({ providerLeaseId: randomUUID() }).where(eq(environmentLeases.id, fixture.leaseId)); + if (kind === "environment" || kind === "company") { + const foreign = await seedCompanyAgentAndEnvironment(); + await db.update(computers).set(kind === "company" ? { companyId: foreign.companyId } : { environmentId: foreign.environmentId }) + .where(eq(computers.id, fixture.computerId)); + } + expect(await heartbeatService(db).sweepOrphanedActiveLeases({ backoffMs: 300_000 })).toEqual({ recovered: 1 }); + expect(await leaseRow(fixture.leaseId)).toMatchObject({ status: "pending_cleanup", failureReason: "orphaned_active_lease_recovered" }); + }); + + it("preserves a warm owner whose stopped process was relaunched within the same owner generation", async () => { + const fixture = await seedWarmComputerLease(); + await db.update(computers).set({ ledger: { owners: [{ ...fixture.owner, + process: { ...fixture.owner.process, launchGeneration: 2 } }] } }) + .where(eq(computers.id, fixture.computerId)); + // Process launch generation is an independent counter, not an owner epoch. + expect(await heartbeatService(db).sweepOrphanedActiveLeases({ backoffMs: 300_000 })).toEqual({ recovered: 0 }); + }); + + it("honors only the existing bounded warm shutdown grace", async () => { + const fixture = await seedWarmComputerLease(); + await db.update(computers).set({ ledger: { owners: [{ ...fixture.owner, deadline: new Date(Date.now() - 10_000).toISOString() }] } }) + .where(eq(computers.id, fixture.computerId)); + expect(await heartbeatService(db).sweepOrphanedActiveLeases({ backoffMs: 300_000 })).toEqual({ recovered: 0 }); + }); + + it("allows cleanup after the same warm reservation and shutdown grace expire", async () => { + const fixture = await seedWarmComputerLease(); + const heartbeat = heartbeatService(db); + expect(await heartbeat.sweepOrphanedActiveLeases({ backoffMs: 0 })).toEqual({ recovered: 0 }); + await db.update(computers).set({ ledger: { owners: [{ ...fixture.owner, deadline: new Date(Date.now() - 60_000).toISOString() }] } }) + .where(eq(computers.id, fixture.computerId)); + expect(await heartbeat.sweepOrphanedActiveLeases({ backoffMs: 0 })).toEqual({ recovered: 1 }); + }); + + it("does not overwrite lease retention completed after the sweep selected its stale snapshot", async () => { + const fixture = await seedWarmComputerLease(); + await db.update(computers).set({ ledger: { owners: [{ ...fixture.owner, phase: "retired" }] } }) + .where(eq(computers.id, fixture.computerId)); + const transaction = db.transaction.bind(db); + const intercept = vi.spyOn(db, "transaction").mockImplementationOnce(async (callback, config) => { + // Sweep has selected the old active row, but heartbeat finishes retention + // before the computer-row lock/cleanup claim. The status CAS must lose. + await db.update(environmentLeases).set({ status: "retained" }).where(eq(environmentLeases.id, fixture.leaseId)); + return transaction(callback, config); + }); + try { + expect(await heartbeatService(db).sweepOrphanedActiveLeases({ backoffMs: 0 })).toEqual({ recovered: 0 }); + expect(await leaseRow(fixture.leaseId)).toMatchObject({ status: "retained", failureReason: null }); + } finally { intercept.mockRestore(); } + }); + it.each([false, true])("releases an interrupted local run's bookkeeping lease (deleted environment: %s)", async (deletedEnvironment) => { const { companyId, agentId, environmentId } = await seedCompanyAgentAndEnvironment(); await db.delete(environments).where(eq(environments.driver, "local")); diff --git a/server/src/modules/computers/application/service.ts b/server/src/modules/computers/application/service.ts index ba2ab0f2514..2aa863ea292 100644 --- a/server/src/modules/computers/application/service.ts +++ b/server/src/modules/computers/application/service.ts @@ -1,6 +1,7 @@ import { randomUUID } from "node:crypto"; import type { ComputerRepository, ComputerBackend } from "./ports.js"; import { + COMPUTER_RETIREMENT_GRACE_MS as RETIREMENT_GRACE_MS, ComputerError, ComputerStopPendingError, assertAdmission, @@ -25,8 +26,6 @@ class ComputerSessionBusyError extends ComputerError { } } -// This is shutdown time only; admission and ordinary work stop at the warm deadline. -const RETIREMENT_GRACE_MS = 30_000; const FILE_OPERATION_BUDGET_MS = 120_000; const FILE_OPERATION_SETTLEMENT_MS = 5000; type Scope = { companyId: string; environmentId: string }; diff --git a/server/src/modules/computers/domain/ledger.ts b/server/src/modules/computers/domain/ledger.ts index eceef75e212..1d5f50bb86c 100644 --- a/server/src/modules/computers/domain/ledger.ts +++ b/server/src/modules/computers/domain/ledger.ts @@ -1,3 +1,6 @@ +/** Shutdown only; ordinary work remains fenced at the warm deadline. */ +export const COMPUTER_RETIREMENT_GRACE_MS = 30_000; + export type OwnerRef = { computerId: string; ownerId: string; diff --git a/server/src/services/heartbeat/recovery.ts b/server/src/services/heartbeat/recovery.ts index a727f4ad5f4..95cd90b00ea 100644 --- a/server/src/services/heartbeat/recovery.ts +++ b/server/src/services/heartbeat/recovery.ts @@ -1,3 +1,4 @@ +import { COMPUTER_RETIREMENT_GRACE_MS } from "../../modules/computers/domain/ledger.js"; import { preserveWorkspaceRestoreRecoveryMetadataSql } from "../legacy-workspace-restore-recovery.js"; import { pendingNativeWorkspaceFinalizationCondition } from "../native-runtime/native-workspace-finalization-state.js"; import { CHAT_COMPLETION_WAKE_REASON } from "../chat-completion-delivery.js"; @@ -30,6 +31,7 @@ import { import type { Db } from "@paperclipai/db"; import { agents, + computers, environmentLeases, heartbeatRuns, issues, @@ -1413,7 +1415,7 @@ export function createHeartbeatRecovery(db: Db, dependencies: HeartbeatRecoveryD const cutoff = new Date(Date.now() - opts.backoffMs); const rows = await db - .select({ lease: environmentLeases }) + .select({ lease: environmentLeases, runAgentId: heartbeatRuns.agentId }) .from(environmentLeases) .leftJoin( heartbeatRuns, @@ -1435,7 +1437,7 @@ export function createHeartbeatRecovery(db: Db, dependencies: HeartbeatRecoveryD .limit(ORPHANED_ACTIVE_LEASE_SWEEP_PAGE_SIZE); let recovered = 0; - for (const { lease } of rows) { + for (const { lease, runAgentId } of rows) { // A provider resource id names one physical sandbox. A different lease // row can still hold that same resource in a live status, so this sweep // must not tear down a sandbox that a different lease still owns. @@ -1470,19 +1472,59 @@ export function createHeartbeatRecovery(db: Db, dependencies: HeartbeatRecoveryD // pending_cleanup sweep can accept the same row in this same tick. A // fresh timestamp here would push the row inside that sweep's own // backoff window and delay the teardown by one full tick. - const flipped = await db + const sameLease = and( + eq(environmentLeases.id, lease.id), + eq(environmentLeases.status, "active"), + // A completed heartbeat may retain or replace its bookkeeping while + // this sweep waits for the computer serialization boundary. + eq(environmentLeases.companyId, lease.companyId), + lease.environmentId ? eq(environmentLeases.environmentId, lease.environmentId) : isNull(environmentLeases.environmentId), + lease.heartbeatRunId ? eq(environmentLeases.heartbeatRunId, lease.heartbeatRunId) : isNull(environmentLeases.heartbeatRunId), + lease.metadata === null ? isNull(environmentLeases.metadata) : eq(environmentLeases.metadata, lease.metadata), + lease.provider === null ? isNull(environmentLeases.provider) : eq(environmentLeases.provider, lease.provider), + lease.providerLeaseId === null ? isNull(environmentLeases.providerLeaseId) : eq(environmentLeases.providerLeaseId, lease.providerLeaseId), + ); + const claim = (query: Pick) => query .update(environmentLeases) - .set({ - status: "pending_cleanup", - failureReason: "orphaned_active_lease_recovered", - }) - .where( - and( - eq(environmentLeases.id, lease.id), - eq(environmentLeases.status, "active"), - ), - ) + .set({ status: "pending_cleanup", failureReason: "orphaned_active_lease_recovered" }) + .where(sameLease) .returning({ id: environmentLeases.id }); + const binding = parseObject(lease.metadata?.computerOwner); + const flipped = lease.provider === "boat" && lease.metadata?.driver === "computer" + ? await db.transaction(async (tx) => { + const [computer] = typeof binding.computerId === "string" && lease.environmentId + ? await tx.select().from(computers).where(and( + eq(sql`${computers.id}::text`, binding.computerId), + eq(computers.companyId, lease.companyId), + eq(computers.environmentId, lease.environmentId), + eq(computers.provider, "boat"), + )).for("update") : []; + const owner = Array.isArray(computer?.ledger.owners) + ? computer.ledger.owners.map(parseObject).find((value) => + value.id === binding.ownerId && value.id === lease.providerLeaseId && + Number.isSafeInteger(binding.generation) && Number(binding.generation) > 0 && + value.generation === binding.generation && value.kind === "runner" && + value.runId === lease.heartbeatRunId && runAgentId !== null && + value.agentId === runAgentId && value.agentId === lease.metadata?.agentId) + : undefined; + const process = parseObject(owner?.process); + const processBound = [process.bootId, process.unitName, process.nonce] + .every((value) => typeof value === "string" && value.length > 0) && + Number.isSafeInteger(process.launchGeneration) && Number(process.launchGeneration) > 0; + const deadline = typeof owner?.deadline === "string" ? Date.parse(owner.deadline) : NaN; + if (owner?.phase === "warm" && processBound && Number.isFinite(deadline) && + deadline + COMPUTER_RETIREMENT_GRACE_MS > Date.now()) { + // Native completion retains the durable owner before heartbeat's + // later instruction cleanup can mark the lease retained. A terminal + // run is not permission to kill this bounded warm process. Keep it + // out of the page until a later sweep; expired/mismatched owners + // still follow the ordinary orphan cleanup path. + await tx.update(environmentLeases).set({ updatedAt: new Date() }) + .where(sameLease); + return []; + } + return claim(tx); + }) : await claim(db); if (flipped.length > 0) recovered += 1; }