diff --git a/packages/adapter-utils/src/command-managed-runtime.ts b/packages/adapter-utils/src/command-managed-runtime.ts index b46b2476f1c..ea013b89a20 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 725bf128b76..365a6d813bb 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 1c2f077fd88..b2bfc01e66b 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 c47c30e5e72..a261e54d81d 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 3f5e0471b6f..ba2ab0f2514 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 cc807470100..e05a1306614 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 a7ba645a48b..5d82abd6e36 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 3b7ad012768..d24528a30a9 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 c7ae3c89dcb..550db61c3cc 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 {