perf(workspaces): reduce snapshot stream scheduling overhead (#19900)

Use bounded 1 MiB filesystem buffers during workspace archive hashing, encryption, and restore. Preserve integrity, ownership, cancellation, and snapshot format guarantees.

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

server/src/services/remote-workspace-archive-io.ts

@@ -0,0 +1 @@
+export const WORKSPACE_ARCHIVE_IO_BUFFER_BYTES = 1024 * 1024;

server/src/services/remote-workspace-snapshot-storage.test.ts

@@ -1,4 +1,4 @@
-import { randomUUID } from "node:crypto";
+import { randomBytes, randomUUID } from "node:crypto";
 import fs from "node:fs/promises";
 import os from "node:os";
 import path from "node:path";
@@ -40,10 +40,10 @@ suite("remote workspace snapshot storage", () => {
     storage = createStorageService(createLocalDiskStorageProvider(path.join(root, "objects")));
   });
   afterEach(async () => { vi.unstubAllEnvs(); await db.delete(executionWorkspaceSnapshots); await fs.rm(root, { recursive: true, force: true }); });
-  async function fixture() {
+  async function fixture(content: string | Buffer = "acknowledged workspace content\n") {
     const companyId = randomUUID(), workspaceId = randomUUID(), operationId = randomUUID();
     const cwd = path.join(root, "source"), archivePath = path.join(root, "archive.pcws"); await fs.mkdir(cwd);
-    await fs.writeFile(path.join(cwd, "saved.txt"), "acknowledged workspace content\n");
+    await fs.writeFile(path.join(cwd, "saved.txt"), content);
     const captured = await captureRemoteWorkspaceDirectory({ cwd, archivePath });
     const row = await beginWorkspaceSnapshot(db, { companyId, workspaceId, operationId, generation: 1, epoch: 3 });
     const upload = { db, storage, snapshotId: row.id, companyId, workspaceId, archivePath, manifestDigest: captured.manifestDigest };
@@ -63,6 +63,18 @@ suite("remote workspace snapshot storage", () => {
     const [row] = await db.select().from(executionWorkspaceSnapshots).where(eq(executionWorkspaceSnapshots.id, ref.id));
     expect(row?.status).toBe("uploaded"); // The owning service still controls atomic publication.
   });
+  it("preserves multi-chunk archive bytes through encryption, authenticated readback, and cold restore", async () => {
+    const content = randomBytes(3 * 1024 * 1024 + 137), f = await fixture(content);
+    const verifiedArchivePath = path.join(root, "verified-large.pcws");
+    const ref = await uploadWorkspaceSnapshot({ ...f.upload, verifiedArchivePath });
+    expect((await fs.readFile(verifiedArchivePath)).equals(await fs.readFile(f.archivePath))).toBe(true);
+    await fs.rm(f.cwd, { recursive: true }); await fs.rm(f.archivePath);
+    const restoredArchive = path.join(root, "cold-large.pcws"), targetDir = path.join(root, "cold-large");
+    await loadWorkspaceSnapshotArchive({ storage, reference: ref, targetPath: restoredArchive });
+    expect((await fs.readFile(restoredArchive)).equals(await fs.readFile(verifiedArchivePath))).toBe(true);
+    await restoreRemoteWorkspaceArchive({ archivePath: restoredArchive, targetDir, expectedManifestDigest: ref.manifestDigest });
+    expect((await fs.readFile(path.join(targetDir, "saved.txt"))).equals(content)).toBe(true);
+  });
   it.each(["fresh", "uploaded", "lost-reply"])("reuses one authenticated %s readback for a complete cold restore", async state => {
     const f = await fixture(), put = storage.putFile.bind(storage);
     if (state === "uploaded") await uploadWorkspaceSnapshot(f.upload);
@@ -332,15 +344,16 @@ suite("remote workspace snapshot storage", () => {
     await expect(fs.access(targetPath)).rejects.toMatchObject({ code: "ENOENT" });
   });
   it("cancels an active object body and removes private decryption scratch before rejecting", async () => {
-    const f = await fixture(), ref = await uploadWorkspaceSnapshot(f.upload);
+    const f = await fixture(randomBytes(3 * 1024 * 1024 + 137)), ref = await uploadWorkspaceSnapshot(f.upload);
     const ciphertext = await fs.readFile(path.join(root, "objects", ref.objectKey));
     const controller = new AbortController();
-    let produced = 0, closed = false;
+    let produced = 0, reads = 0, closed = false;
     const body = new Readable({
       highWaterMark: 1,
       read() {
-        if (produced === 2) { controller.abort(new Error("client disconnected")); return; }
-        this.push(ciphertext.subarray(produced, ++produced));
+        if (reads++ === 2) { controller.abort(new Error("client disconnected")); return; }
+        const start = produced; produced += 1024 * 1024 + 17;
+        this.push(ciphertext.subarray(start, produced));
       },
       destroy(error, callback) { closed = true; callback(error); },
     });

server/src/services/remote-workspace-snapshot-storage.ts

@@ -12,6 +12,7 @@ import { localEncryptedProvider } from "../secrets/local-encrypted-provider.js";
 import type { StoredSecretVersionMaterial } from "../secrets/types.js";
 import type { StorageService } from "../storage/types.js";
 import { notifyRemoteWorkspaceSnapshots, REMOTE_WORKSPACE_SNAPSHOT_GRACE_MS } from "./remote-workspace-snapshot-retention.js";
+import { WORKSPACE_ARCHIVE_IO_BUFFER_BYTES } from "./remote-workspace-archive-io.js";
 
 const SCHEMA = "paperclip.workspace-snapshot.v1";
 const KEY_SCHEMA = "paperclip.workspace-snapshot-key.v1";
@@ -55,7 +56,7 @@ function validateReference(ref: WorkspaceSnapshotRef) {
 }
 async function digestFile(file: string, signal?: AbortSignal) {
   const digest = createHash("sha256"); let byteSize = 0;
-  for await (const part of createReadStream(file, { signal })) { digest.update(part); byteSize += part.length; }
+  for await (const part of createReadStream(file, { signal, highWaterMark: WORKSPACE_ARCHIVE_IO_BUFFER_BYTES })) { digest.update(part); byteSize += part.length; }
   return { sha256: digest.digest("hex"), byteSize };
 }
 async function verifyObject(storage: StorageService, ref: WorkspaceSnapshotRef, signal?: AbortSignal) {
@@ -167,7 +168,7 @@ export async function uploadWorkspaceSnapshot(input: WorkspaceSnapshotUpload): P
       const before = await handle.stat(), digest = createHash("sha256"); let bytes = 0;
       if (!before.isFile() || (before.mode & 0o777) !== 0o600 || before.nlink !== 1
         || typeof process.getuid === "function" && before.uid !== process.getuid()) failure("readback_file_changed");
-      for await (const chunk of createReadStream(file, { fd: handle.fd, autoClose: false, signal: input.signal })) {
+      for await (const chunk of createReadStream(file, { fd: handle.fd, autoClose: false, signal: input.signal, highWaterMark: WORKSPACE_ARCHIVE_IO_BUFFER_BYTES })) {
         digest.update(chunk); bytes += chunk.length;
       }
       const after = await handle.stat(), named = await lstat(file);
@@ -241,7 +242,7 @@ async function uploadVerifiedWorkspaceSnapshot(input: WorkspaceSnapshotUpload, v
     try { await header.writeFile(Buffer.concat([MAGIC, iv])); } finally { await header.close(); }
     const incomingDigest = createHash("sha256"); let incomingBytes = 0;
     const observed = new Transform({ transform(bytes, _encoding, callback) { incomingDigest.update(bytes); incomingBytes += bytes.length; callback(null, bytes); } });
-    await pipeline(createReadStream(input.archivePath, { signal: input.signal }), observed, cipher, createWriteStream(encrypted, { flags: "a", mode: 0o600 }), { signal: input.signal });
+    await pipeline(createReadStream(input.archivePath, { signal: input.signal, highWaterMark: WORKSPACE_ARCHIVE_IO_BUFFER_BYTES }), observed, cipher, createWriteStream(encrypted, { flags: "a", mode: 0o600, highWaterMark: WORKSPACE_ARCHIVE_IO_BUFFER_BYTES }), { signal: input.signal });
     if (incomingBytes !== plaintext.byteSize || incomingDigest.digest("hex") !== plaintext.sha256) failure("archive_changed");
     const footer = await open(encrypted, "a"); try { await footer.writeFile(cipher.getAuthTag()); } finally { await footer.close(); }
     dataKey.fill(0);
@@ -275,7 +276,7 @@ async function uploadVerifiedWorkspaceSnapshot(input: WorkspaceSnapshotUpload, v
     // The object key is unique per attempt, so a late losing uploader cannot
     // overwrite an object already used by a committed generation.
     await input.storage.putFile({ companyId: identity.companyId, namespace: "workspace-snapshots", objectKey: attempt.objectKey,
-      originalFilename: null, contentType: "application/octet-stream", body: createReadStream(encrypted),
+      originalFilename: null, contentType: "application/octet-stream", body: createReadStream(encrypted, { highWaterMark: WORKSPACE_ARCHIVE_IO_BUFFER_BYTES }),
       byteSize: attempt.byteSize, sha256: attempt.sha256, signal: input.signal });
     await verify({ schema: SCHEMA, ...identity, ...attempt } as WorkspaceSnapshotRef);
     if (recapture) {
@@ -318,7 +319,7 @@ export async function loadWorkspaceSnapshotArchive(input: {
         received += bytes.length; digest.update(bytes);
         callback(received > ref.byteSize ? new Error("remote_workspace_snapshot_object_size_mismatch") : null, bytes);
       } });
-      try { await pipeline(object.stream, bounded, createWriteStream(ciphertext, { flags: "wx", mode: 0o600 }), { signal }); }
+      try { await pipeline(object.stream, bounded, createWriteStream(ciphertext, { flags: "wx", mode: 0o600, highWaterMark: WORKSPACE_ARCHIVE_IO_BUFFER_BYTES }), { signal }); }
       finally { object.stream.destroy(); }
     } catch (error) {
       if (signal.aborted && signal.reason === stalled) throw stalled;
@@ -332,8 +333,8 @@ export async function loadWorkspaceSnapshotArchive(input: {
     const decipher = createDecipheriv("aes-256-gcm", key, header.subarray(MAGIC.length));
     decipher.setAAD(aad(ref)); decipher.setAuthTag(tag);
     try {
-      await pipeline(createReadStream(ciphertext, { start: header.length, end: ref.byteSize - tag.length - 1, signal: input.signal }),
-        decipher, createWriteStream(plaintext, { flags: "wx", mode: 0o600 }), { signal: input.signal });
+      await pipeline(createReadStream(ciphertext, { start: header.length, end: ref.byteSize - tag.length - 1, signal: input.signal, highWaterMark: WORKSPACE_ARCHIVE_IO_BUFFER_BYTES }),
+        decipher, createWriteStream(plaintext, { flags: "wx", mode: 0o600, highWaterMark: WORKSPACE_ARCHIVE_IO_BUFFER_BYTES }), { signal: input.signal });
     } finally { key.fill(0); }
     const actual = await digestFile(plaintext, input.signal);
     if (actual.sha256 !== wrapped.plaintextSha256 || actual.byteSize !== wrapped.plaintextBytes) failure("plaintext_hash_mismatch");
@@ -341,7 +342,7 @@ export async function loadWorkspaceSnapshotArchive(input: {
     // failed/cancelled relay cannot leave a plaintext cache that looks complete.
     const staged = path.join(path.dirname(input.targetPath), `.workspace-snapshot-${randomUUID()}.tmp`);
     try {
-      await pipeline(createReadStream(plaintext, { signal: input.signal }), createWriteStream(staged, { flags: "wx", mode: 0o600 }), { signal: input.signal });
+      await pipeline(createReadStream(plaintext, { signal: input.signal, highWaterMark: WORKSPACE_ARCHIVE_IO_BUFFER_BYTES }), createWriteStream(staged, { flags: "wx", mode: 0o600, highWaterMark: WORKSPACE_ARCHIVE_IO_BUFFER_BYTES }), { signal: input.signal });
       input.signal?.throwIfAborted(); await link(staged, input.targetPath);
     } finally { await rm(staged, { force: true }); }
   } finally { await rm(scratch, { recursive: true, force: true }); }

server/src/services/remote-workspace-transport.test.ts

@@ -1,5 +1,5 @@
 import { execFile } from "node:child_process";
-import { createHash, randomUUID } from "node:crypto";
+import { createHash, randomBytes, randomUUID } from "node:crypto";
 import fs from "node:fs/promises";
 import os from "node:os";
 import path from "node:path";
@@ -162,6 +162,17 @@ describe("remote workspace transport safety with real sandbox commands", () => {
     expect(await fs.lstat(root).catch(() => null)).toBeNull();
   });
 
+  it("verifies a multi-chunk downloaded archive and restores every byte", async () => {
+    const content = randomBytes(3 * 1024 * 1024 + 137);
+    await fs.writeFile(path.join(workspace, "large.bin"), content);
+    const captured = await captureRemoteWorkspaceArchive({ target: target().remote, remoteCwd: workspace, operationId: randomUUID() });
+    try {
+      const targetDir = path.join(scratch, "large-restored");
+      await restoreRemoteWorkspaceArchive({ archivePath: captured.archivePath, targetDir, expectedManifestDigest: captured.manifestDigest });
+      expect((await fs.readFile(path.join(targetDir, "large.bin"))).equals(content)).toBe(true);
+    } finally { await captured.cleanup(); }
+  });
+
   it("adopts an atomic restored folder after a lost reply without overwriting later edits", async () => {
     await fs.writeFile(path.join(workspace, "saved.txt"), "saved synthetic files\n");
     const { remote, execute } = target();

server/src/services/remote-workspace-transport.ts

@@ -11,6 +11,7 @@ import { removeTaskDirectoryTree } from "./isolated-task-directory-snapshot.js";
 import type { GitRemoteAuthProvider } from "./git-credentials.js";
 import { classifyWorkspaceRestoreFailure } from "@paperclipai/adapter-utils/workspace-restore-merge";
 import { MAX_SANDBOX_SYNC_TIMEOUT_MS, type SandboxSyncOperation } from "@paperclipai/adapter-utils/sandbox-managed-runtime";
+import { WORKSPACE_ARCHIVE_IO_BUFFER_BYTES } from "./remote-workspace-archive-io.js";
 
 export type RemoteWorkspaceTarget = Extract<AdapterExecutionTarget, { transport: "sandbox" }>;
 
@@ -580,7 +581,7 @@ async function captureSandboxWorkspaceArchive(input: {
     timing.transferMs = elapsedSince(transferStarted);
     const verifyStarted = performance.now();
     const hash = createHash("sha256"); let bytes = 0;
-    for await (const chunk of createReadStream(archivePath)) { hash.update(chunk); bytes += chunk.length; }
+    for await (const chunk of createReadStream(archivePath, { highWaterMark: WORKSPACE_ARCHIVE_IO_BUFFER_BYTES })) { hash.update(chunk); bytes += chunk.length; }
     timing.verifyMs = elapsedSince(verifyStarted);
     if (hash.digest("hex") !== receipt.sha256 || bytes !== receipt.byteSize) throw new Error("remote_workspace_capture_transfer_mismatch");
     return { ...receipt, archivePath, cleanup: () => rm(root, { recursive: true, force: true }) };