Stream large remote provider archives with bounded parallel uploads

Co-Authored-By: Paperclip <noreply@paperclip.ing>

server/src/services/native-runtime/native-session-executor.test.ts

@@ -47,7 +47,7 @@ import {
   type NativeExecutionInput,
   type PrpEvent,
 } from "@paperclipai/paperclip-runner";
-import { createHash, generateKeyPairSync, randomUUID } from "node:crypto";
+import { createHash, generateKeyPairSync, randomBytes, randomUUID } from "node:crypto";
 import { DatabaseSync } from "node:sqlite";
 import { nativeSha256 } from "./canonical.js";
 import { agentDirectoryWorkingCopyService } from "../agent-directory-working-copies.js";
@@ -2463,7 +2463,7 @@ describe("remote provider checkpoint snapshots", () => {
 });
 
 describe("remote provider checkpoint restores", () => {
-  it("stages a provider pack through a command-only computer runner", async () => {
+  it.each([{ sizeMiB: 0, corrupt: false }, { sizeMiB: 65, corrupt: false }, { sizeMiB: 0, corrupt: true }])("stages a $sizeMiB MiB provider pack through a command-only computer runner (corrupt=$corrupt)", async ({ sizeMiB, corrupt }) => {
     const root = await mkdtemp(join(tmpdir(), "paperclip-computer-pack-"));
     const sourcePath = join(root, "source");
     const targetPath = join(root, "remote pack's files");
@@ -2472,27 +2472,50 @@ describe("remote provider checkpoint restores", () => {
       await mkdir(join(sourcePath, "bin"), { recursive: true });
       await writeFile(join(sourcePath, "provider-pack.json"), manifest);
       await writeFile(join(sourcePath, "bin", "provider"), "#!/bin/sh\nexit 0\n", { mode: 0o755 });
-      const execute = vi.fn(async (request: { command: string; args: string[]; stdin: string }) => {
-        execFileSync(request.command, request.args, { input: request.stdin });
+      const chunk = randomBytes(1024 * 1024);
+      const expectedDigest = createHash("sha256");
+      const payload = await open(join(sourcePath, "payload.bin"), "w");
+      try {
+        for (let index = 0; index < sizeMiB; index++) {
+          await payload.write(chunk);
+          expectedDigest.update(chunk);
+        }
+      } finally { await payload.close(); }
+      const execute = vi.fn(async (request: { command: string; args: string[]; stdin?: string }) => {
+        expect(request.stdin?.length ?? 0).toBeLessThanOrEqual(Math.ceil(4 * 1024 * 1024 / 3) * 4);
+        const stdin = corrupt && request.stdin
+          ? Buffer.from(request.stdin, "base64").subarray(0, 32).toString("base64")
+          : request.stdin;
+        execFileSync(request.command, request.args, { input: stdin });
         return { exitCode: 0, timedOut: false, stdout: "", stderr: "" };
       });
       const runner = { execute };
-      await stageRemoteRunnerDirectory({
+      const staging = stageRemoteRunnerDirectory({
         target: { kind: "remote", transport: "computer", remoteCwd: root, runner } as never,
         runner: runner as never,
         sourcePath,
         targetPath,
         mode: 0o700,
       });
+      if (corrupt) {
+        await expect(staging).rejects.toThrow();
+        await expect(access(targetPath)).rejects.toThrow();
+        expect((await readdir(root)).some((name) => name.includes(".upload-"))).toBe(false);
+        return;
+      }
+      await staging;
       expect(await readFile(join(targetPath, "provider-pack.json"), "utf8")).toBe(manifest);
       expect((await lstat(join(targetPath, "bin", "provider"))).mode & 0o100).toBe(0o100);
       expect((await lstat(targetPath)).mode & 0o777).toBe(0o700);
-      expect(execute).toHaveBeenCalledOnce();
+      const actualDigest = execFileSync("shasum", ["-a", "256", join(targetPath, "payload.bin")], { encoding: "utf8" }).split(" ")[0];
+      expect(actualDigest).toBe(expectedDigest.digest("hex"));
+      expect(execute.mock.calls.filter(([request]) => request.stdin)).toHaveLength(sizeMiB ? 17 : 1);
       expect(execute.mock.calls[0]![0]).toMatchObject({ bypassSession: true });
+      expect((await readdir(root)).some((name) => name.includes(".upload-"))).toBe(false);
     } finally {
       await rm(root, { recursive: true, force: true });
     }
-  });
+  }, 60_000);
 
   it("does not upload excluded Codex scratch trees or credentials", async () => {
     const sourcePath = await mkdtemp(

server/src/services/native-runtime/native-session-executor.ts

@@ -49,9 +49,10 @@ import {
   EXECUTION_CONTROL_DEADLINE_MS,
 } from "../execution-control-deadline.js";
 import { createHash, randomUUID } from "node:crypto";
-import { execFileSync } from "node:child_process";
+import { execFile, execFileSync } from "node:child_process";
 import {
   chmodSync,
+  createReadStream,
   copyFileSync,
   closeSync,
   constants,
@@ -10348,24 +10349,55 @@ export async function stageRemoteRunnerDirectory(input: {
     }
     return;
   }
-  const archive = execFileSync(
-    "tar",
-    [...excludeArgs, "-czf", "-", "-C", input.sourcePath, "."],
-    { encoding: "buffer", maxBuffer: 64 * 1024 * 1024 },
-  );
-  const escapedTarget = input.targetPath.replaceAll("'", "'\\''");
-  const script =
-    `umask 077; mkdir -p '${escapedTarget}' && ` +
-    `base64 -d | tar -xzf - -C '${escapedTarget}' && ` +
-    `chmod ${input.mode.toString(8)} '${escapedTarget}'`;
-  const result = await input.runner.execute({
-    command: "sh",
-    args: ["-c", script],
-    stdin: archive.toString("base64"),
-    bypassSession: true,
-  });
-  if (result.exitCode !== 0 || result.timedOut) {
-    throw new Error("runner_remote_directory_staging_failed");
+  const stagingRoot = mkdtempSync(join(tmpdir(), "paperclip-runner-upload-"));
+  const archivePath = join(stagingRoot, "payload.tar.gz");
+  const remoteArchive = `${input.targetPath}.upload-${randomUUID()}.tar.gz`;
+  const quote = (value: string) => `'${value.replaceAll("'", "'\\''")}'`;
+  const run = async (script: string, stdin?: string) => {
+    const result = await input.runner.execute({
+      command: "sh", args: ["-c", script], stdin,
+      bypassSession: true, timeoutMs: 180_000,
+    });
+    if (result.exitCode !== 0 || result.timedOut) {
+      throw new Error("runner_remote_directory_staging_failed");
+    }
+  };
+  const pending: Promise<void>[] = [];
+  let uploadError: unknown;
+  try {
+    // Provider packs can exceed a gigabyte. Keep compression off the event loop
+    // and transfer from disk in bounded chunks instead of buffering the archive.
+    await new Promise<void>((resolve, reject) => {
+      execFile("tar", [...excludeArgs, "-czf", archivePath, "-C", input.sourcePath, "."],
+        { maxBuffer: 1024 * 1024 }, (error) => error ? reject(error) : resolve());
+    });
+    await run(`umask 077; mkdir -p ${quote(posix.dirname(remoteArchive))} && : > ${quote(remoteArchive)}`);
+    const digest = createHash("sha256");
+    let chunkIndex = 0;
+    const stream = createReadStream(archivePath, { highWaterMark: 4 * 1024 * 1024 });
+    for await (const bytes of stream) {
+      const chunk = bytes as Buffer;
+      digest.update(chunk);
+      // Disjoint fixed offsets permit bounded parallel upload and idempotent
+      // writes. The final digest rejects missing, duplicated, or partial bytes.
+      pending.push(run(`base64 -d | dd of=${quote(remoteArchive)} bs=4194304 seek=${chunkIndex++} conv=notrunc 2>/dev/null`,
+        chunk.toString("base64")).catch((error) => { uploadError ??= error; }));
+      if (pending.length === 4) {
+        await Promise.all(pending);
+        pending.length = 0;
+        if (uploadError) throw uploadError;
+      }
+    }
+    await Promise.all(pending);
+    if (uploadError) throw uploadError;
+    const expectedDigest = digest.digest("hex");
+    await run(`test "$(if command -v sha256sum >/dev/null 2>&1; then sha256sum ${quote(remoteArchive)}; else shasum -a 256 ${quote(remoteArchive)}; fi | cut -d ' ' -f 1)" = ${quote(expectedDigest)} && ` +
+      `umask 077 && mkdir -p ${quote(input.targetPath)} && ` +
+      `tar -xzf ${quote(remoteArchive)} -C ${quote(input.targetPath)} && chmod ${input.mode.toString(8)} ${quote(input.targetPath)}`);
+  } finally {
+    await Promise.all(pending);
+    rmSync(stagingRoot, { recursive: true, force: true });
+    await run(`rm -f ${quote(remoteArchive)}`).catch(() => undefined);
   }
 }