From 8198c9a571ea24aedc32040d0ebae90535694421 Mon Sep 17 00:00:00 2001 From: Dotta Date: Sun, 11 Oct 2026 02:59:57 -0500 Subject: [PATCH 1/2] fix: share Boat upload capacity before command admission --- .../src/command-managed-runtime.ts | 10 ++++ .../modules/computers/adapters/boat.test.ts | 28 +++++++++ server/src/modules/computers/adapters/boat.ts | 46 ++++++++++++++- .../computers/application/service.test.ts | 23 ++++++++ .../modules/computers/application/service.ts | 1 + .../native-session-executor.test.ts | 57 +++++++++++++++++++ .../native-runtime/native-session-executor.ts | 12 ++-- .../remote-provider-pack-cache.test.ts | 13 +++++ .../remote-provider-pack-cache.ts | 1 + 9 files changed, 184 insertions(+), 7 deletions(-) diff --git a/packages/adapter-utils/src/command-managed-runtime.ts b/packages/adapter-utils/src/command-managed-runtime.ts index b46b2476f1..ea013b89a2 100644 --- a/packages/adapter-utils/src/command-managed-runtime.ts +++ b/packages/adapter-utils/src/command-managed-runtime.ts @@ -58,6 +58,16 @@ export interface CommandManagedDuplexChannel { } export interface CommandManagedRuntimeRunner { + /** + * Optional transport-shared backpressure for stdin artifact chunks. Acquire + * before execute starts its command deadline; this grants no execution + * authority. The callback must still use the owned execute wrapper. Queue + * waiting is bounded separately, and signal cancels only a queued operation. + */ + withStdinUploadSlot?( + operation: () => Promise, + options: { timeoutMs: number; signal?: AbortSignal }, + ): Promise; /** * True when the provider verified the concurrent-sync opt-in. A native runner * carries the value from the effective capability snapshot diff --git a/server/src/modules/computers/adapters/boat.test.ts b/server/src/modules/computers/adapters/boat.test.ts index 725bf128b7..365a6d813b 100644 --- a/server/src/modules/computers/adapters/boat.test.ts +++ b/server/src/modules/computers/adapters/boat.test.ts @@ -891,6 +891,34 @@ describe("Boat SSH admission", () => { await Promise.all([...jobs, queued]); } finally { vi.useRealTimers(); } }); + it("bounds and cancels queued upload slots without starting their callbacks", async () => { + vi.useFakeTimers(); + try { + const admission = createBoatTransportAdmission(); + const completions: Array<() => void> = []; + const occupied = Array.from({ length: 4 }, () => admission.withStdinUploadSlot( + () => new Promise(resolve => completions.push(resolve)), { timeoutMs: 180_000 })); + await vi.advanceTimersByTimeAsync(0); + const aborted = new AbortController(); + const cancelledOperation = vi.fn(async () => undefined); + const cancelled = admission.withStdinUploadSlot(cancelledOperation, { timeoutMs: 180_000, signal: aborted.signal }); + const cancelledAssertion = expect(cancelled).rejects.toThrow("cancelled"); + aborted.abort(); + await cancelledAssertion; + const expiredOperation = vi.fn(async () => undefined); + const expired = admission.withStdinUploadSlot(expiredOperation, { timeoutMs: 999_999 }); + const expiredAssertion = expect(expired).rejects.toThrow("timed out waiting for capacity"); + await vi.advanceTimersByTimeAsync(180_000); + await expiredAssertion; + expect(cancelledOperation).not.toHaveBeenCalled(); + expect(expiredOperation).not.toHaveBeenCalled(); + completions.splice(0).forEach(resolve => resolve()); + await Promise.all(occupied); + await expect(admission.withStdinUploadSlot(async () => { throw new Error("failed upload"); }, { timeoutMs: 1000 })).rejects.toThrow("failed upload"); + await expect(admission.withStdinUploadSlot(async () => "next", { timeoutMs: 1000 })).resolves.toBe("next"); + } finally { vi.useRealTimers(); } + }); + it("releases all capacity after failed or cancelled transport commands without replay", async () => { const admission = createBoatTransportAdmission(); const runner = { execute: vi.fn(async () => { throw new Error("transport cancelled"); }) }; diff --git a/server/src/modules/computers/adapters/boat.ts b/server/src/modules/computers/adapters/boat.ts index 1c2f077fd8..b2bfc01e66 100644 --- a/server/src/modules/computers/adapters/boat.ts +++ b/server/src/modules/computers/adapters/boat.ts @@ -17,6 +17,8 @@ const apiOrigin = "https://boat.dev/api/v1"; const transportCache = new Map>(); type CommandInput = Parameters[0]; +const BOAT_BULK_COMMAND_LIMIT = 4; + /** Reserve SSH admission for lifecycle/identity calls while bulk and legacy commands run. */ export function createBoatTransportAdmission() { let active = 0; @@ -26,14 +28,14 @@ export function createBoatTransportAdmission() { function drain() { while (active < 8) { const eligible = (entry: (typeof queue)[number]) => - entry.control || (ordinary < 6 && (!entry.bulk || bulk < 4)); + entry.control || (ordinary < 6 && (!entry.bulk || bulk < BOAT_BULK_COMMAND_LIMIT)); let index = queue.findIndex((entry) => entry.control && eligible(entry)); if (index < 0) index = queue.findIndex(eligible); if (index < 0) break; queue.splice(index, 1)[0]!.start(); } } - return async function execute( + async function execute( runner: CommandManagedRuntimeRunner, input: CommandInput, control: boolean, @@ -72,7 +74,44 @@ export function createBoatTransportAdmission() { queue.push(entry); drain(); }); + } + + // All runner views for a machine share these slots. Wait outside execute so + // one upload cannot consume another upload's SSH command deadline in queue. + let uploads = 0; + const uploadQueue: Array<{ start: () => void }> = []; + function drainUploads() { + while (uploads < BOAT_BULK_COMMAND_LIMIT && uploadQueue.length) uploadQueue.shift()!.start(); + } + const withStdinUploadSlot: NonNullable = (operation, options) => { + if (uploadQueue.length >= 128) return Promise.reject(new ComputerError("conflict", "Computer upload admission is busy")); + return new Promise((resolve, reject) => { + const cleanup = () => { clearTimeout(timer); options.signal?.removeEventListener("abort", aborted); }; + const remove = (message: string) => { + const index = uploadQueue.indexOf(entry); + if (index < 0) return; + uploadQueue.splice(index, 1); + cleanup(); + reject(new ComputerError("provider_error", message)); + }; + const aborted = () => remove("Computer upload admission was cancelled"); + const entry = { start: () => { + cleanup(); + uploads++; + void Promise.resolve().then(operation).then(resolve, reject).finally(() => { + uploads--; + drainUploads(); + }); + } }; + const timer = setTimeout(() => remove("Computer upload timed out waiting for capacity"), + Math.max(1, Math.min(options.timeoutMs, 180_000))); + uploadQueue.push(entry); + options.signal?.addEventListener("abort", aborted, { once: true }); + if (options.signal?.aborted) aborted(); + else drainUploads(); + }); }; + return Object.assign(execute, { withStdinUploadSlot }); } const transportAdmission = new Map>(); const desktopCache = new Map< @@ -358,7 +397,8 @@ export function boatBackend( args: ["--signal=KILL", `${input.timeoutMs! / 1000}s`, input.command, ...(input.args ?? [])], }), } : raw; - return { ...raw, execute: (input) => execute(bounded, input, options?.control === true) }; + return { ...raw, withStdinUploadSlot: admission.withStdinUploadSlot, + execute: (input) => execute(bounded, input, options?.control === true) }; } async function fileTransport( record: ComputerRecord, diff --git a/server/src/modules/computers/application/service.test.ts b/server/src/modules/computers/application/service.test.ts index c47c30e5e7..a261e54d81 100644 --- a/server/src/modules/computers/application/service.test.ts +++ b/server/src/modules/computers/application/service.test.ts @@ -1,3 +1,4 @@ +import { createBoatTransportAdmission } from "../adapters/boat.js"; import { describe, it, expect, vi } from "vitest"; import { existsSync, @@ -115,6 +116,28 @@ function fixture() { }; } describe("computer ownership", () => { + it("rechecks owner authority after a shared upload slot wait before SSH dispatch", async () => { + const f = fixture(); + const admission = createBoatTransportAdmission(); + const execute = vi.fn(); + vi.mocked(f.backend.runner).mockResolvedValue({ execute, withStdinUploadSlot: admission.withStdinUploadSlot }); + await f.attach(); + const binding = await f.admit(); + expect(binding.runner.withStdinUploadSlot).toBe(binding.process.runner.withStdinUploadSlot); + const completions: Array<() => void> = []; + const occupied = Array.from({ length: 4 }, () => admission.withStdinUploadSlot( + () => new Promise(resolve => completions.push(resolve)), { timeoutMs: 180_000 })); + await Promise.resolve(); + const queued = binding.process.runner.withStdinUploadSlot!( + () => binding.process.runner.execute({ command: "must-not-launch", stdin: "payload", timeoutMs: 180_000 }), + { timeoutMs: 180_000 }); + const assertion = expect(queued).rejects.toThrow(/expired|retired/); + await f.service.retire({ ...f.scope, owner: binding.owner }); + completions.forEach(resolve => resolve()); + await Promise.all(occupied); + await assertion; + expect(execute).not.toHaveBeenCalled(); + }); it("allocates outside guest ephemeral and occupied ports while reserving live owners", async () => { const f = fixture(); await f.attach(); diff --git a/server/src/modules/computers/application/service.ts b/server/src/modules/computers/application/service.ts index 3f5e0471b6..ba2ab0f251 100644 --- a/server/src/modules/computers/application/service.ts +++ b/server/src/modules/computers/application/service.ts @@ -228,6 +228,7 @@ export function createComputerService( const runnerFor = ( processScoped: boolean, ): CommandManagedRuntimeRunner => ({ + withStdinUploadSlot: (processScoped ? control : raw).withStdinUploadSlot, execute: async (input) => { const current = await (processScoped ? processScope() diff --git a/server/src/services/native-runtime/native-session-executor.test.ts b/server/src/services/native-runtime/native-session-executor.test.ts index cc80747010..e05a130661 100644 --- a/server/src/services/native-runtime/native-session-executor.test.ts +++ b/server/src/services/native-runtime/native-session-executor.test.ts @@ -1,3 +1,4 @@ +import { createBoatTransportAdmission } from "../../modules/computers/adapters/boat.js"; import { closeWarmNativeSession, compareAndSealRetiredComputerCheckpoint, retiredComputerCheckpointIsSettled, retiredComputerCheckpointUnsettledReason, retireUnlaunchedComputerAttempt, nativeSessionBootstrapTimeoutMs } from "./native-session-executor.js"; import { prepareHeartbeatGitHubLaunchers } from "../heartbeat-github-launchers.js"; import { RemoteProviderPackVerificationError } from "./remote-provider-pack-cache.js"; @@ -2717,6 +2718,62 @@ describe("remote provider checkpoint restores", () => { } finally { await rm(root, { recursive: true, force: true }); } }); + it("shares Boat upload capacity across concurrent packs before command deadlines begin", async () => { + const root = await mkdtemp(join(tmpdir(), "paperclip-boat-upload-pressure-")); + const sourcePath = join(root, "source"); + const bytes = randomBytes(24 * 1024 * 1024); + await writeFile(sourcePath, bytes); + const admission = createBoatTransportAdmission(); + const result = { exitCode: 0, signal: null, timedOut: false, stdout: "", stderr: "", pid: null, startedAt: "" }; + let active = 0; + let peak = 0; + let entered!: () => void; + const firstBatch = new Promise(resolve => { entered = resolve; }); + let wake: (() => void) | undefined; + const budgets: number[] = []; + const raw = { execute: vi.fn(async (request: { command: string; args?: string[]; stdin?: string; timeoutMs?: number }) => { + if (!request.stdin) { + execFileSync(request.command, request.args ?? []); + return result; + } + active++; + peak = Math.max(peak, active); + if (active === 4) entered(); + wake?.(); wake = undefined; + budgets.push(request.timeoutMs!); + try { + await new Promise(resolve => setTimeout(resolve, Math.min(100_000, request.timeoutMs!))); + if (request.timeoutMs! < 100_000) return { ...result, exitCode: null, timedOut: true }; + execFileSync(request.command, request.args ?? [], { input: request.stdin }); + return result; + } finally { active--; } + }) }; + const runner = { + withStdinUploadSlot: admission.withStdinUploadSlot, + execute: (request: Parameters[1]) => admission(raw, request, true), + }; + vi.useFakeTimers({ toFake: ["Date", "setTimeout", "clearTimeout"] }); + try { + let settled = false; + const transfer = Promise.allSettled(["one", "two"].map(name => stageRemoteRunnerFile({ + target: { kind: "remote", transport: "computer" } as never, + runner: runner as never, sourcePath, targetPath: join(root, name), mode: 0o700, + }))).then(results => { settled = true; wake?.(); wake = undefined; return results; }); + await firstBatch; + await expect(admission(raw, { command: "true", timeoutMs: 1000 }, true)).resolves.toMatchObject({ exitCode: 0 }); + while (!settled) { + if (!active) await new Promise(resolve => { wake = resolve; }); + if (!settled) await vi.advanceTimersByTimeAsync(100_000); + } + const results = await transfer; + expect(results.map(result => result.status)).toEqual(["fulfilled", "fulfilled"]); + expect(peak).toBe(4); + expect(budgets).toHaveLength(12); + expect(budgets.every(budget => budget > 99_999)).toBe(true); + for (const name of ["one", "two"]) expect((await readFile(join(root, name))).equals(bytes)).toBe(true); + } finally { vi.useRealTimers(); await rm(root, { recursive: true, force: true }); } + }); + it("shares eight upload slots across concurrent artifact transfers", async () => { const root = await mkdtemp(join(tmpdir(), "paperclip-runner-upload-slots-")); const sourcePath = join(root, "source"); diff --git a/server/src/services/native-runtime/native-session-executor.ts b/server/src/services/native-runtime/native-session-executor.ts index a7ba645a48..5d82abd6e3 100644 --- a/server/src/services/native-runtime/native-session-executor.ts +++ b/server/src/services/native-runtime/native-session-executor.ts @@ -10545,7 +10545,7 @@ export async function stageRemoteRunnerFile(input: { } const remoteUpload = `${input.targetPath}.upload-${randomUUID()}`; try { - await uploadRemoteRunnerFileChunks({ sourcePath: input.sourcePath, remotePath: remoteUpload, run }); + await uploadRemoteRunnerFileChunks({ sourcePath: input.sourcePath, remotePath: remoteUpload, run, runner }); await run(`chmod ${input.mode.toString(8)} ${quote(remoteUpload)} && mv -f ${quote(remoteUpload)} ${quote(input.targetPath)}`); } finally { await run(`rm -f ${quote(remoteUpload)}`).catch(() => undefined); @@ -10572,6 +10572,7 @@ async function withRemoteArtifactUploadSlot(operation: () => Promise): Pro } async function uploadRemoteRunnerFileChunks(input: { + runner: CommandManagedRuntimeRunner; sourcePath: string; remotePath: string; run: (script: string, stdin?: string) => Promise; @@ -10597,9 +10598,12 @@ async function uploadRemoteRunnerFileChunks(input: { // Disjoint fixed offsets permit bounded parallel upload and idempotent // writes. The final digest rejects missing, duplicated, or partial bytes. const offset = chunkIndex++; - pending.push(withRemoteArtifactUploadSlot(() => run( + const upload = () => run( `base64 -d | dd of=${quote(remotePath)} bs=4194304 seek=${offset} conv=notrunc 2>/dev/null`, - chunk.toString("base64"))).catch((error) => { uploadError ??= error; })); + chunk.toString("base64")); + pending.push(withRemoteArtifactUploadSlot(() => input.runner.withStdinUploadSlot + ? input.runner.withStdinUploadSlot(upload, { timeoutMs: 180_000 }) + : upload()).catch((error) => { uploadError ??= error; })); if (pending.length === REMOTE_ARTIFACT_UPLOAD_CONCURRENCY) { await Promise.all(pending); pending.length = 0; @@ -10708,7 +10712,7 @@ export async function stageRemoteRunnerDirectory(input: { (error) => error ? reject(error) : resolve()); }); await uploadRemoteRunnerFileChunks({ - sourcePath: archivePath, remotePath: remoteArchive, run, onProgress: input.onProgress, + sourcePath: archivePath, remotePath: remoteArchive, run, runner: input.runner, onProgress: input.onProgress, }); await run(`umask 077 && mkdir -p ${quote(input.targetPath)} && ` + `tar -xzf ${quote(remoteArchive)} -C ${quote(input.targetPath)} && chmod ${input.mode.toString(8)} ${quote(input.targetPath)}`); diff --git a/server/src/services/native-runtime/remote-provider-pack-cache.test.ts b/server/src/services/native-runtime/remote-provider-pack-cache.test.ts index 3b7ad01276..d24528a30a 100644 --- a/server/src/services/native-runtime/remote-provider-pack-cache.test.ts +++ b/server/src/services/native-runtime/remote-provider-pack-cache.test.ts @@ -120,6 +120,19 @@ describe("persistent computer provider pack cache", () => { expect(() => computerProviderPackCachePath("/agents/../a", digest)).toThrow("identity_invalid"); }); + it("preserves transport-shared upload backpressure through the ownership guard", async () => { + const input = await fixture(); + const withStdinUploadSlot = vi.fn(async (operation: () => Promise) => operation()); + await expect(prepareComputerProviderPackCache({ ...input, + runner: { ...input.runner, withStdinUploadSlot } as CommandManagedRuntimeRunner, + stage: async (root, guarded) => { + expect(guarded.withStdinUploadSlot).toBe(withStdinUploadSlot); + await guarded.withStdinUploadSlot!(() => input.stage(root, guarded), { timeoutMs: 180_000 }); + }, + })).resolves.toBe("published"); + expect(withStdinUploadSlot).toHaveBeenCalledOnce(); + }); + it("verifies a shared hit before the compression/upload callback and excludes its ownership marker", async () => { const input = await fixture(); const stage = vi.fn(input.stage); const verify = vi.fn(input.verify); diff --git a/server/src/services/native-runtime/remote-provider-pack-cache.ts b/server/src/services/native-runtime/remote-provider-pack-cache.ts index c7ae3c89dc..550db61c3c 100644 --- a/server/src/services/native-runtime/remote-provider-pack-cache.ts +++ b/server/src/services/native-runtime/remote-provider-pack-cache.ts @@ -298,6 +298,7 @@ export async function prepareComputerProviderPackCache(input: { } const stagingRoot = posix.join(posix.dirname(input.cacheRoot), ".uploads", session, posix.basename(input.cacheRoot), identity.attemptId, "pack"); const guardedRunner: CommandManagedRuntimeRunner = { + withStdinUploadSlot: input.runner.withStdinUploadSlot, execute: (command) => input.runner.execute({ ...command, command: "python3", args: [...args("guard"), command.command, ...(command.args ?? [])] }), }; try { From d14578ad13e681cd5dde0025061c8b1493b0aed1 Mon Sep 17 00:00:00 2001 From: Dotta Date: Sun, 11 Oct 2026 03:00:55 -0500 Subject: [PATCH 2/2] fix: preserve authoritative warm computers during orphan lease sweep --- ...rtbeat-orphaned-active-lease-sweep.test.ts | 102 +++++++++++++++++- .../modules/computers/application/service.ts | 3 +- server/src/modules/computers/domain/ledger.ts | 3 + server/src/services/heartbeat/recovery.ts | 68 +++++++++--- 4 files changed, 160 insertions(+), 16 deletions(-) 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 6ff69f16f0..5d5b65277e 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 ba2ab0f251..2aa863ea29 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 eceef75e21..1d5f50bb86 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 a727f4ad5f..95cd90b00e 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; }