diff --git a/doc/sandbox-work-folders.md b/doc/sandbox-work-folders.md index a13875808a..49b7bef3dd 100644 --- a/doc/sandbox-work-folders.md +++ b/doc/sandbox-work-folders.md @@ -209,6 +209,29 @@ transfer uses the downloaded version's size, hash, and executable bit. Its sync baseline records those same bytes, so an unchanged copy cannot overwrite a later shared edit. +Providers with native file synchronization or explicit streaming-stdin support +hydrate files through bounded stdin batches instead of one remote command per +small chunk. A batch carries at most 4 MiB of file bytes and 256 operations; +each file still uses confined paths, SHA-256 validation, and atomic publication. +Before publishing an incoming batch or applying incoming deletions, the host +persists the intended versions in Postgres. After a failed or lost response, +the next run reconciles those intents against observed disk contents before +saving outgoing changes. Imported bytes therefore cannot overwrite a newer +shared version by being mistaken for an agent edit; actual subsequent edits +still synchronize normally. Failed explicit refreshes retain the same intent. +Other providers keep the small-argument transport. Incoming storage responses +are prefetched four at a time and closed if transfer fails. Repository restores +use the same path, then recreate confined repository links. + +Repository checkpoints transfer at most four distinct content-addressed blobs +concurrently, avoiding duplicate uploads for identical files. Small-file reads +are grouped into at most 1 MiB and 64 files per remote command, with at most +four read batches cached per checkpoint. Larger files stream independently. +Retries bypass that cache and reopen the actual file. All active transfers +must settle, and a second filesystem scan must match, before the +complete checkpoint reference can advance. Scoped-file retry receipts and +last-write-wins publication remain ordered. + Outgoing checkpoints run every **180 seconds**, with at most one in flight, and a final flush when execution stops. File signatures include content and executable state. Unchanged stale working copies do not overwrite newer shared diff --git a/server/src/__tests__/helpers/work-folder-runner.ts b/server/src/__tests__/helpers/work-folder-runner.ts index 9bba68071c..fe0384c259 100644 --- a/server/src/__tests__/helpers/work-folder-runner.ts +++ b/server/src/__tests__/helpers/work-folder-runner.ts @@ -5,8 +5,10 @@ const exec = promisify(execFile); export const localTestWorkFolderRunner: CommandManagedRuntimeRunner = { async execute(input) { try { - const { stdout, stderr } = await exec(input.command, input.args ?? [], { cwd: input.cwd, + const execution = exec(input.command, input.args ?? [], { cwd: input.cwd, env: { ...process.env, ...input.env }, timeout: input.timeoutMs, maxBuffer: 32 * 1024 * 1024 }); + execution.child.stdin?.end(input.stdin); + const { stdout, stderr } = await execution; return { stdout, stderr, exitCode: 0, signal: null, timedOut: false }; } catch (error) { const value = error as Error & { stdout?: string; stderr?: string }; diff --git a/server/src/__tests__/sandbox-work-folders.test.ts b/server/src/__tests__/sandbox-work-folders.test.ts index 62ec64a620..52b4183eb4 100644 --- a/server/src/__tests__/sandbox-work-folders.test.ts +++ b/server/src/__tests__/sandbox-work-folders.test.ts @@ -17,6 +17,7 @@ import * as activityLog from "../services/activity-log.js"; import { workFolderService } from "../services/work-folders.js"; import * as workFolderServices from "../services/work-folders.js"; import { collectWorkFolderGarbage } from "../services/work-folder-garbage.js"; +import type { CommandManagedRuntimeRunner } from "@paperclipai/adapter-utils/command-managed-runtime"; import { localTestWorkFolderRunner } from "./helpers/work-folder-runner.js"; const exec = promisify(execFile); @@ -146,7 +147,7 @@ describe("shared sandbox work-folder lifecycle", () => { await expect(bindWarmSandboxWorkspace(db, input)).rejects.toThrow("active task run"); }); async function prepare(home: string, leaseId: string, physicalId = leaseId, responsibleUserId: string | null = null, - options: { taskId?: string; branchName?: string; agentId?: string } = {}) { + options: { taskId?: string; branchName?: string; agentId?: string; bulkStdin?: boolean } = {}) { await fs.mkdir(home, { recursive: true }); const runId = randomUUID(); const boundAgentId = options.agentId ?? agentId; @@ -157,7 +158,7 @@ describe("shared sandbox work-folder lifecycle", () => { const run = await prepareSandboxWorkFolders({ db, companyId, agentId: boundAgentId, projectId, taskId: options.taskId ?? taskId, runId, primaryWorkspaceId: primary!.id, primaryBranchName: options.branchName, responsibleUserId, storage, sandboxKey: workFolderSandboxKey(lease), target: { kind: "remote", transport: "sandbox", leaseId, remoteCwd: home, - runner: { execute: (input) => localTestWorkFolderRunner.execute({ ...input, env: { ...input.env, HOME: home } }) } } }); + runner: { supportsSingleStreamStdinProgress: options.bulkStdin, execute: (input) => localTestWorkFolderRunner.execute({ ...input, env: { ...input.env, HOME: home } }) } } }); active.push(run); return run; } it("uses downloaded metadata when a shared file changes after the startup listing", async () => { @@ -248,11 +249,11 @@ describe("shared sandbox work-folder lifecycle", () => { active.splice(active.indexOf(run), 1); }, 120_000); - it("reuses clones and restores saved unpushed work, staged changes, and task files after losing the sandbox", async () => { + it.each([false, true])("reuses clones and restores saved unpushed work, staged changes, and task files after losing the sandbox (bulk stdin: %s)", async (bulkStdin) => { const repositoryTaskId = randomUUID(); await db.insert(issues).values({ id: repositoryTaskId, companyId, projectId, title: "Repository recovery", assigneeAgentId: agentId }); - const task = { taskId: repositoryTaskId }; - const home = path.join(root, "sandbox"); + const task = { taskId: repositoryTaskId, bulkStdin }; + const home = path.join(root, `sandbox-${bulkStdin}`); const leaseId = randomUUID(); const first = await prepare(home, leaseId, leaseId, null, task); expect(first.home).toBe(home); @@ -278,7 +279,7 @@ describe("shared sandbox work-folder lifecycle", () => { await warm.stop(); active.splice(active.indexOf(warm), 1); await fs.rm(home, { recursive: true }); const replacementId = randomUUID(); - const restored = await prepare(path.join(root, "replacement"), replacementId, replacementId, null, task); + const restored = await prepare(path.join(root, `replacement-${bulkStdin}`), replacementId, replacementId, null, task); expect(await fs.readFile(path.join(restored.primaryRepo, ".setup-count"), "utf8")).toBe("initialized\ninitialized\n"); expect(await fs.readFile(path.join(restored.primaryRepo, "node_modules/acceptance/installed"), "utf8")).toBe("ready"); await expect(fs.stat(path.join(restored.primaryRepo, "node_modules/acceptance/warm-cache"))).rejects.toMatchObject({ code: "ENOENT" }); @@ -290,6 +291,120 @@ describe("shared sandbox work-folder lifecycle", () => { expect(await fs.readlink(path.join(restored.primaryRepo, "link"))).toBe("tracked"); await restored.stop(); active.splice(active.indexOf(restored), 1); }, 120_000); + it.each([ + { bulk: true, refresh: false, edit: false, apply: true }, + { bulk: true, refresh: false, edit: false, apply: false }, + { bulk: false, refresh: false, edit: false, apply: true }, + { bulk: true, refresh: false, edit: true, apply: true }, + { bulk: true, refresh: true, edit: false, apply: true }, + ])("reconciles imported files after a lost response (bulk=$bulk refresh=$refresh edit=$edit apply=$apply)", async ({ bulk, refresh, edit, apply }) => { + const scopedAgent = randomUUID(), sandboxKey = randomUUID(); + await db.insert(agents).values({ id: scopedAgent, companyId, name: "Inbound recovery" }); + const svc = workFolderService(db, storage); + const folder = await svc.ensure({ companyId, scope: "agent", ownerId: scopedAgent }); + const home = path.join(root, `incoming-${sandboxKey}`); + await fs.mkdir(home); + let loseResponse = false; + let failedRunId = ""; + const runner: CommandManagedRuntimeRunner = { + supportsSingleStreamStdinProgress: bulk, + execute: async (input) => { + const request = input.command === "node" ? JSON.parse(Buffer.from(input.args!.at(-1)!, "base64").toString()) : {}; + const fail = loseResponse && request.root === path.join(home, "agent") + && (request.operation === "publish" || request.operation === "batch"); + if (fail && !apply) { loseResponse = false; throw new Error("Injected lost incoming response"); } + const result = await localTestWorkFolderRunner.execute({ ...input, env: { ...input.env, HOME: home } }); + if (fail) { + expect(result.exitCode).toBe(0); + loseResponse = false; + throw new Error("Injected lost incoming response"); + } + return result; + }, + }; + async function start() { + const runId = randomUUID(); failedRunId = runId; + await db.insert(heartbeatRuns).values({ id: runId, companyId, agentId: scopedAgent, status: "running" }); + const run = await prepareSandboxWorkFolders({ db, companyId, agentId: scopedAgent, responsibleUserId: null, + projectId: null, taskId: null, runId, storage, sandboxKey, + target: { kind: "remote", transport: "sandbox", leaseId: randomUUID(), remoteCwd: home, runner } }); + active.push(run); return run; + } + await svc.write(folder, { path: "nested/shared.txt", body: Buffer.from("v0"), operationId: randomUUID() }); + const first = await start(); + if (!refresh) { await first.stop(); active.splice(active.indexOf(first), 1); } + await svc.write(folder, { path: "nested/shared.txt", body: Buffer.from("v1-imported"), executable: true, operationId: randomUUID() }); + loseResponse = true; + if (refresh) { + await db.update(workFolderRuns).set({ refreshRequested: true }).where(eq(workFolderRuns.runId, first.manifest.runId)); + await expect(first.stop()).rejects.toThrow("Injected lost incoming response"); + active.splice(active.indexOf(first), 1); + } else { + await expect(start()).rejects.toThrow("Injected lost incoming response"); + } + expect(await fs.readFile(path.join(home, "agent/nested/shared.txt"), "utf8")).toBe(apply ? "v1-imported" : "v0"); + const [failed] = await db.select().from(workFolderRuns).where(eq(workFolderRuns.runId, failedRunId)); + expect(failed?.state).toBe("failed"); + expect(failed?.baselines["incoming:agent"]).toEqual(expect.arrayContaining([ + expect.objectContaining({ path: "nested/shared.txt", executable: true }), + expect.objectContaining({ path: "nested", kind: "directory" }), + ])); + await svc.write(folder, { path: "nested/shared.txt", body: Buffer.from("v2-newer-shared"), operationId: randomUUID() }); + if (edit) await fs.writeFile(path.join(home, "agent/nested/shared.txt"), "v3-actual-local-edit"); + const retry = await start(); + const expected = edit ? "v3-actual-local-edit" : "v2-newer-shared"; + expect(await fs.readFile(path.join(home, "agent/nested/shared.txt"), "utf8")).toBe(expected); + const content = await svc.content(folder, "nested/shared.txt"); + expect(await content.stream.toArray().then((chunks) => Buffer.concat(chunks).toString())).toBe(expected); + const [recovered] = await db.select().from(workFolderRuns).where(eq(workFolderRuns.runId, retry.manifest.runId)); + expect(recovered?.baselines["incoming:agent"]).toBeUndefined(); + expect(recovered?.baselines["incomingRemoved:agent"]).toBeUndefined(); + await retry.stop(); active.splice(active.indexOf(retry), 1); + }, 60_000); + + it("does not repeat an imported deletion against a newer shared file after a lost remove response", async () => { + const scopedAgent = randomUUID(), sandboxKey = randomUUID(); + await db.insert(agents).values({ id: scopedAgent, companyId, name: "Deletion recovery" }); + const svc = workFolderService(db, storage); + const folder = await svc.ensure({ companyId, scope: "agent", ownerId: scopedAgent }); + const home = path.join(root, `incoming-delete-${sandboxKey}`); + await fs.mkdir(home); + let loseResponse = false, failedRunId = ""; + const runner: CommandManagedRuntimeRunner = { + supportsSingleStreamStdinProgress: true, + execute: async (input) => { + const result = await localTestWorkFolderRunner.execute({ ...input, env: { ...input.env, HOME: home } }); + const request = input.command === "node" ? JSON.parse(Buffer.from(input.args!.at(-1)!, "base64").toString()) : {}; + if (loseResponse && request.root === path.join(home, "agent") && request.operation === "remove") { + expect(result.exitCode).toBe(0); loseResponse = false; + throw new Error("Injected lost removal response"); + } + return result; + }, + }; + async function start() { + const runId = randomUUID(); failedRunId = runId; + await db.insert(heartbeatRuns).values({ id: runId, companyId, agentId: scopedAgent, status: "running" }); + const run = await prepareSandboxWorkFolders({ db, companyId, agentId: scopedAgent, responsibleUserId: null, + projectId: null, taskId: null, runId, storage, sandboxKey, + target: { kind: "remote", transport: "sandbox", leaseId: randomUUID(), remoteCwd: home, runner } }); + active.push(run); return run; + } + await svc.write(folder, { path: "nested/shared.txt", body: Buffer.from("old"), operationId: randomUUID() }); + const first = await start(); await first.stop(); active.splice(active.indexOf(first), 1); + await svc.remove(folder, "nested", randomUUID()); + loseResponse = true; + await expect(start()).rejects.toThrow("Injected lost removal response"); + await expect(fs.stat(path.join(home, "agent/nested/shared.txt"))).rejects.toMatchObject({ code: "ENOENT" }); + const [failed] = await db.select().from(workFolderRuns).where(eq(workFolderRuns.runId, failedRunId)); + expect(failed?.baselines["incomingRemoved:agent"].map((entry) => entry.path)).toEqual(["nested/shared.txt", "nested"]); + await svc.write(folder, { path: "nested/shared.txt", body: Buffer.from("new shared after deletion"), operationId: randomUUID() }); + const retry = await start(); + expect(await fs.readFile(path.join(home, "agent/nested/shared.txt"), "utf8")).toBe("new shared after deletion"); + const content = await svc.content(folder, "nested/shared.txt"); + expect(await content.stream.toArray().then((chunks) => Buffer.concat(chunks).toString())).toBe("new shared after deletion"); + await retry.stop(); active.splice(active.indexOf(retry), 1); + }, 60_000); it("does not let an unchanged stale shared file overwrite a newer durable value", async () => { const svc = workFolderService(db, storage); const folder = await svc.ensure({ companyId, scope: "project", ownerId: projectId }); @@ -348,9 +463,12 @@ describe("shared sandbox work-folder lifecycle", () => { expect((await svc.list(folder)).files.map((file) => file.path)).toContain("private"); }, 120_000); - it("reopens repository blobs after transient PUT failures and retains the previous checkpoint when retries exhaust", async () => { + it.each([false, true])("reopens repository blobs after transient PUT failures and retains the previous checkpoint when retries exhaust (bulk stdin: %s)", async (bulkStdin) => { const leaseId = randomUUID(); - const run = await prepare(path.join(root, "replayed-repository-upload"), leaseId); + const repositoryTaskId = randomUUID(); + await db.insert(issues).values({ id: repositoryTaskId, companyId, projectId, title: "Repository retry", assigneeAgentId: agentId }); + const run = await prepare(path.join(root, `replayed-repository-upload-${bulkStdin}`), leaseId, leaseId, null, + { bulkStdin, taskId: repositoryTaskId }); await run.flush(); const filename = path.join(run.primaryRepo, "retry-upload"); const content = "replay the entire repository file"; diff --git a/server/src/__tests__/work-folder-read-cache.test.ts b/server/src/__tests__/work-folder-read-cache.test.ts new file mode 100644 index 0000000000..bc5a5b4074 --- /dev/null +++ b/server/src/__tests__/work-folder-read-cache.test.ts @@ -0,0 +1,121 @@ +import { Readable } from "node:stream"; +import { describe, expect, it, vi } from "vitest"; +import { createWorkFolderReadCache } from "../services/work-folder-read-cache.js"; +import type { WorkTreeEntry } from "../services/work-folder-transport.js"; + +function entry(path: string, byteSize = 1): WorkTreeEntry { + return { path, kind: "file", byteSize, sha256: "a".repeat(64), executable: false }; +} +async function consume(stream: Readable) { + return Buffer.concat(await stream.toArray()); +} +const bytes = (entries: WorkTreeEntry[]) => entries.map((entry) => Buffer.alloc(entry.byteSize, "x")); + +describe("checkpoint small-file read cache", () => { + it("loads lazily, coalesces concurrent paths and always reopens a repeated path", async () => { + const entries = [entry("a"), entry("b"), entry("c")]; + const load = vi.fn(async (group: WorkTreeEntry[]) => bytes(group)); + const fallback = vi.fn(() => Readable.from([Buffer.from("fresh")])); + const cache = createWorkFolderReadCache(entries, load, fallback); + const a = cache.read(entries[0]!), b = cache.read(entries[1]!); + expect(load).not.toHaveBeenCalled(); + expect((await Promise.all([consume(a), consume(b)])).map((value) => value.toString())).toEqual(["x", "x"]); + expect(await consume(cache.read(entries[2]!))).toEqual(Buffer.from("x")); + expect(load).toHaveBeenCalledTimes(1); + expect(await consume(cache.read(entries[0]!))).toEqual(Buffer.from("fresh")); + expect(fallback).toHaveBeenCalledTimes(1); + cache.clear(); + }); + + it("bounds groups by byte size and count, including empty files, and streams larger files", async () => { + const entries = [entry("large", 1024 * 1024 + 1), entry("full", 1024 * 1024), entry("next", 1), + ...Array.from({ length: 129 }, (_, i) => entry(`empty-${i}`, 0))]; + const load = vi.fn(async (group: WorkTreeEntry[]) => bytes(group)); + const fallback = vi.fn(() => Readable.from([Buffer.from("large-stream")])); + const cache = createWorkFolderReadCache(entries, load, fallback); + for (const value of entries) await consume(cache.read(value)); + expect(fallback).toHaveBeenCalledTimes(1); + expect(fallback).toHaveBeenCalledWith(entries[0]); + expect(load).toHaveBeenCalledTimes(4); + for (const [group] of load.mock.calls) { + expect(group.length).toBeLessThanOrEqual(64); + expect(group.reduce((sum, value) => sum + value.byteSize, 0)).toBeLessThanOrEqual(1024 * 1024); + } + expect(load.mock.calls.flatMap(([group]) => group.map((value) => value.path))).toEqual(entries.slice(1).map((value) => value.path)); + cache.clear(); + }); + + it("retains only four least-recently-used groups even when most grouped files are never requested", async () => { + const entries = Array.from({ length: 10 }, (_, i) => entry(`file-${i}`, 512 * 1024)); + const load = vi.fn(async (group: WorkTreeEntry[]) => bytes(group)); + const cache = createWorkFolderReadCache(entries, load, () => { throw new Error("unexpected fallback"); }); + for (const index of [0, 2, 4, 6]) await consume(cache.read(entries[index]!)); + // Touch the first group, then create a fifth: group two is now the LRU. + await consume(cache.read(entries[1]!)); + await consume(cache.read(entries[8]!)); + expect(load).toHaveBeenCalledTimes(5); + await consume(cache.read(entries[5]!)); + expect(load).toHaveBeenCalledTimes(5); + await consume(cache.read(entries[3]!)); + expect(load).toHaveBeenCalledTimes(6); + expect(load.mock.calls.at(-1)![0].map((value) => value.path)).toEqual(["file-2", "file-3"]); + cache.clear(); + }); + + it("uses a fresh source after a rejected load and allows untouched paths to reload their failed group", async () => { + const entries = [entry("a"), entry("b"), entry("c")]; + const load = vi.fn(async (group: WorkTreeEntry[]) => bytes(group)).mockRejectedValueOnce(new Error("read failed")); + const fallback = vi.fn(() => Readable.from([Buffer.from("fresh")])); + const cache = createWorkFolderReadCache(entries, load, fallback); + const results = await Promise.allSettled([consume(cache.read(entries[0]!)), consume(cache.read(entries[1]!))]); + expect(results.map((result) => result.status)).toEqual(["rejected", "rejected"]); + expect(load).toHaveBeenCalledTimes(1); + expect(await consume(cache.read(entries[0]!))).toEqual(Buffer.from("fresh")); + expect(await consume(cache.read(entries[2]!))).toEqual(Buffer.from("x")); + expect(load).toHaveBeenCalledTimes(2); + expect(fallback).toHaveBeenCalledTimes(1); + cache.clear(); + }); + + it("clear releases cached groups and sends existing lazy and future readers to fresh streams", async () => { + const entries = [entry("a"), entry("b"), entry("c")]; + const load = vi.fn(async (group: WorkTreeEntry[]) => bytes(group)); + const sources: Readable[] = []; + const fallback = vi.fn(() => { + const source = Readable.from([Buffer.from("fresh")]); sources.push(source); return source; + }); + const cache = createWorkFolderReadCache(entries, load, fallback); + await consume(cache.read(entries[0]!)); + const lazy = cache.read(entries[1]!); + cache.clear(); + expect(await consume(lazy)).toEqual(Buffer.from("fresh")); + expect(await consume(cache.read(entries[2]!))).toEqual(Buffer.from("fresh")); + expect(load).toHaveBeenCalledTimes(1); + expect(fallback).toHaveBeenCalledTimes(2); + expect(sources.every((source) => source.destroyed)).toBe(true); + }); + + it("does not repopulate the cache after clearing an in-flight load", async () => { + const entries = [entry("a"), entry("b")]; + let release!: (value: Buffer[]) => void; + const load = vi.fn(() => new Promise((resolve) => { release = resolve; })); + const fallback = vi.fn(() => Readable.from([Buffer.from("fresh")])); + const cache = createWorkFolderReadCache(entries, load, fallback); + const active = consume(cache.read(entries[0]!)); + await vi.waitFor(() => expect(load).toHaveBeenCalledTimes(1)); + cache.clear(); release(bytes(entries)); + expect(await active).toEqual(Buffer.from("x")); + expect(await consume(cache.read(entries[1]!))).toEqual(Buffer.from("fresh")); + expect(load).toHaveBeenCalledTimes(1); + }); + + it.each([{ buffers: [] }, { buffers: [Buffer.alloc(2)] }])("rejects malformed batch contents and reopens the next attempt", async ({ buffers }) => { + const item = entry("a"); + const load = vi.fn(async () => buffers); + const fallback = vi.fn(() => Readable.from([Buffer.from("fresh")])); + const cache = createWorkFolderReadCache([item], load, fallback); + await expect(consume(cache.read(item))).rejects.toThrow("does not match its entries"); + expect(await consume(cache.read(item))).toEqual(Buffer.from("fresh")); + cache.clear(); + }); +}); diff --git a/server/src/__tests__/work-folder-repositories.test.ts b/server/src/__tests__/work-folder-repositories.test.ts new file mode 100644 index 0000000000..2237c5b3ec --- /dev/null +++ b/server/src/__tests__/work-folder-repositories.test.ts @@ -0,0 +1,221 @@ +import { createHash, randomUUID } from "node:crypto"; +import { Readable } from "node:stream"; +import { setImmediate } from "node:timers/promises"; +import { afterAll, afterEach, beforeAll, describe, expect, it, vi } from "vitest"; +import { eq } from "drizzle-orm"; +import { companies, createDb, issues, startEmbeddedPostgresTestDatabase, + taskRepositoryBindings, workFolderObjects, type Db } from "@paperclipai/db"; +import * as garbage from "../services/work-folder-garbage.js"; +import { workFolderRepositoryService } from "../services/work-folder-repositories.js"; +import type { WorkFolderTransport, WorkTreeEntry } from "../services/work-folder-transport.js"; +import type { PutObjectInput, StorageProvider } from "../storage/types.js"; + +function gate() { + let release!: () => void; + const promise = new Promise((resolve) => { release = resolve; }); + return { promise, release }; +} + +describe("bounded repository checkpoint transfers", () => { + let database: Awaited>; + let db: Db; + const companyId = randomUUID(); + beforeAll(async () => { + database = await startEmbeddedPostgresTestDatabase("paperclip-repository-pool-"); + db = createDb(database.connectionString); + await db.insert(companies).values({ id: companyId, name: "Checkpoint pool" }); + }, 60_000); + afterEach(() => vi.restoreAllMocks()); + afterAll(async () => { await database?.cleanup(); }); + + async function fixture() { + const taskId = randomUUID(); + await db.insert(issues).values({ id: taskId, companyId, title: "Checkpoint" }); + const [binding] = await db.insert(taskRepositoryBindings).values({ companyId, taskId, + workspaceId: randomUUID(), name: "repository" }).returning(); + const contents = new Map(); + const objects = new Map(); + const sources: Readable[] = []; + let entries: WorkTreeEntry[] = []; + const scan = vi.fn(async () => entries); + const transport: WorkFolderTransport = { + home: async () => "/home/runner", scan, readBatch: undefined, + read: vi.fn((_root, filePath) => { + const source = Readable.from([contents.get(filePath)!]); + sources.push(source); + return source; + }), + write: vi.fn(async () => {}), + writeMany: vi.fn(async (_root, _stagingRoot, transfers) => { + for await (const { body } of transfers) { + if (body) for await (const chunk of body) { void chunk; } + } + }), + moveRoot: vi.fn(async () => {}), + symlink: vi.fn(async () => {}), mkdirRoot: vi.fn(async () => {}), + mkdir: vi.fn(async () => {}), remove: vi.fn(async () => {}), + }; + const putObject = vi.fn(async (input: PutObjectInput) => { + const chunks: Buffer[] = []; + if (Buffer.isBuffer(input.body)) chunks.push(input.body); + else for await (const chunk of input.body) chunks.push(Buffer.from(chunk)); + objects.set(input.objectKey, Buffer.concat(chunks)); + }); + const headObject = vi.fn(async ({ objectKey }: { objectKey: string }) => ({ exists: objects.has(objectKey) })); + const storage: StorageProvider = { + id: "local_disk", putObject, headObject, + getObject: async ({ objectKey }) => ({ stream: Readable.from([objects.get(objectKey)!]) }), + deleteObject: async ({ objectKey }) => { objects.delete(objectKey); }, + }; + function file(filePath: string, text = filePath): WorkTreeEntry { + const bytes = Buffer.from(text); + contents.set(filePath, bytes); + return { path: filePath, kind: "file", byteSize: bytes.length, + sha256: createHash("sha256").update(bytes).digest("hex"), executable: false }; + } + const service = workFolderRepositoryService(db, storage, transport); + return { binding: binding!, file, objects, sources, scan, transport, storage, putObject, headObject, + setEntries: (next: WorkTreeEntry[]) => { entries = next; }, + save: () => service.checkpoint(binding!, "/repository"), + current: async () => (await db.select().from(taskRepositoryBindings).where(eq(taskRepositoryBindings.id, binding!.id)))[0]!, + }; + } + + it("bounds HEADs and streaming PUTs to four, deduplicates content, and preserves manifest order", async () => { + const f = await fixture(); + const entries = [f.file("z"), f.file("b"), f.file("same-z", "z"), + f.file("c"), f.file("a"), f.file("fifth"), f.file("empty", "")]; + entries[2]!.executable = true; + f.setEntries(entries); + const registered = vi.spyOn(garbage, "registerWorkFolderObject"); + const heads = gate(), puts = gate(); + let activeHeads = 0, maximumHeads = 0, activePuts = 0, maximumPuts = 0; + f.headObject.mockImplementation(async () => { + activeHeads++; + maximumHeads = Math.max(maximumHeads, activeHeads); + try { await heads.promise; return { exists: false }; } finally { activeHeads--; } + }); + const put = f.putObject.getMockImplementation()!; + f.putObject.mockImplementation(async (input) => { + if (!input.objectKey.includes("/blobs/")) return put(input); + activePuts++; + maximumPuts = Math.max(maximumPuts, activePuts); + try { await puts.promise; await put(input); } finally { activePuts--; } + }); + const saving = f.save(); + try { + await vi.waitFor(() => expect(activeHeads).toBe(4)); + expect(f.headObject).toHaveBeenCalledTimes(4); + expect(registered).toHaveBeenCalledTimes(4); + heads.release(); + await vi.waitFor(() => expect(activePuts).toBe(4)); + expect(f.headObject).toHaveBeenCalledTimes(4); + expect(registered).toHaveBeenCalledTimes(4); + expect((await f.current()).checkpointKey).toBeNull(); + } finally { heads.release(); puts.release(); } + await saving; + expect(maximumHeads).toBe(4); + expect(maximumPuts).toBe(4); + expect(f.headObject).toHaveBeenCalledTimes(6); + const blobPuts = f.putObject.mock.calls.filter(([input]) => input.objectKey.includes("/blobs/")); + expect(blobPuts).toHaveLength(6); + expect(new Set(blobPuts.map(([input]) => input.objectKey)).size).toBe(6); + expect(registered.mock.calls.filter(([, , input]) => input.objectKey.includes("/blobs/"))).toHaveLength(6); + const manifest = JSON.parse(f.objects.get(f.binding.checkpointKey!)!.toString("utf8")); + expect(manifest.files.map(({ objectKey: _key, ...entry }: WorkTreeEntry & { objectKey: string }) => entry)).toEqual(entries); + expect(manifest.files[0].objectKey).toBe(manifest.files[2].objectKey); + expect(f.sources.every((source) => source.destroyed)).toBe(true); + }); + + it("drains in-flight PUTs after failure without scheduling more blobs or replacing the protected checkpoint", async () => { + const f = await fixture(); + const original = f.file("saved"); + f.setEntries([original]); + await f.save(); + const previous = await f.current(); + const failed = f.file("fail"), queued = f.file("must-not-start"); + f.setEntries([original, failed, f.file("held-a"), f.file("held-b"), f.file("held-c"), queued]); + const failedKey = `${companyId}/task-repositories/${f.binding.id}/blobs/${failed.sha256}`; + const queuedKey = `${companyId}/task-repositories/${f.binding.id}/blobs/${queued.sha256}`; + const fail = gate(), held = gate(); + const error = new Error("Injected permanent PUT failure"); + const put = f.putObject.getMockImplementation()!; + let active = 0, settled = false; + f.putObject.mockClear(); + f.headObject.mockClear(); + f.scan.mockClear(); + f.putObject.mockImplementation(async (input) => { + active++; + try { + if (input.objectKey === failedKey) { await fail.promise; throw error; } + await held.promise; + await put(input); + } finally { active--; } + }); + const outcome = f.save().then(() => ({ error: null }), (failure: unknown) => ({ error: failure })) + .finally(() => { settled = true; }); + try { + await vi.waitFor(() => expect(active).toBe(4)); + fail.release(); + await vi.waitFor(() => expect(active).toBe(3)); + await setImmediate(); + expect(settled).toBe(false); + expect((await f.current()).checkpointKey).toBe(previous.checkpointKey); + expect(f.headObject.mock.calls.some(([input]) => input.objectKey === queuedKey)).toBe(false); + expect(f.putObject.mock.calls.some(([input]) => input.objectKey.includes("/checkpoints/"))).toBe(false); + } finally { fail.release(); held.release(); } + expect((await outcome).error).toBe(error); + expect(active).toBe(0); + expect(f.sources.every((source) => source.destroyed)).toBe(true); + expect(f.scan).toHaveBeenCalledTimes(1); + expect((await f.current()).checkpointKey).toBe(previous.checkpointKey); + const tracked = await db.select().from(workFolderObjects).where(eq(workFolderObjects.repositoryBindingId, f.binding.id)); + expect(tracked.find((object) => object.objectKey === previous.checkpointKey)!.deleteAfter).toBeNull(); + expect(tracked.find((object) => object.objectKey.endsWith(`/blobs/${original.sha256}`))!.deleteAfter).toBeNull(); + expect(tracked.filter((object) => object.deleteAfter !== null)).toHaveLength(4); + }); + + it("drains failed concurrent HEADs without opening more file streams", async () => { + const f = await fixture(); + const first = f.file("fail-head"); + f.setEntries([first, f.file("two"), f.file("three"), f.file("four"), f.file("queued")]); + const fail = gate(), held = gate(); + const error = new Error("HEAD unavailable"); + let active = 0, settled = false; + f.headObject.mockImplementation(async ({ objectKey }) => { + active++; + try { + if (objectKey.endsWith(`/blobs/${first.sha256}`)) { await fail.promise; throw error; } + await held.promise; + return { exists: false }; + } finally { active--; } + }); + const outcome = f.save().then(() => ({ error: null }), (failure: unknown) => ({ error: failure })) + .finally(() => { settled = true; }); + try { + await vi.waitFor(() => expect(active).toBe(4)); + fail.release(); + await vi.waitFor(() => expect(active).toBe(3)); + await setImmediate(); + expect(settled).toBe(false); + expect(f.sources).toHaveLength(0); + } finally { fail.release(); held.release(); } + expect((await outcome).error).toBe(error); + expect(active).toBe(0); + expect(f.headObject).toHaveBeenCalledTimes(4); + expect(f.sources).toHaveLength(0); + expect(f.putObject).not.toHaveBeenCalled(); + expect((await f.current()).checkpointKey).toBeNull(); + }); + + it("rejects a changed second scan even after every concurrent blob transfer succeeds", async () => { + const f = await fixture(); + const entries = [f.file("one"), f.file("two"), f.file("three"), f.file("four"), f.file("five")]; + f.scan.mockResolvedValueOnce(entries).mockResolvedValueOnce([...entries, f.file("changed")]); + await expect(f.save()).rejects.toThrow("Repository changed during checkpoint"); + expect(f.putObject).toHaveBeenCalledTimes(5); + expect(f.putObject.mock.calls.every(([input]) => input.objectKey.includes("/blobs/"))).toBe(true); + expect((await f.current()).checkpointKey).toBeNull(); + expect(f.sources.every((source) => source.destroyed)).toBe(true); + }); +}); diff --git a/server/src/__tests__/work-folder-transfer.test.ts b/server/src/__tests__/work-folder-transfer.test.ts new file mode 100644 index 0000000000..8975b1d191 --- /dev/null +++ b/server/src/__tests__/work-folder-transfer.test.ts @@ -0,0 +1,51 @@ +import { Readable } from "node:stream"; +import { describe, expect, it } from "vitest"; +import { prefetchWorkFiles } from "../services/work-folder-transfer.js"; + +const entry = { path: "file", kind: "file" as const, byteSize: 0, sha256: null, executable: false }; +describe("bounded work file prefetch", () => { + it("preserves order and keeps at most four streams open", async () => { + let opened = 0, peak = 0; + const seen: string[] = []; + for await (const transfer of prefetchWorkFiles(Array.from({ length: 30 }, (_, i) => i), async (i) => { + opened++; peak = Math.max(peak, opened); + const body = new Readable({ read() { this.push(null); }, destroy(_error, done) { opened--; done(); } }); + return { entry: { ...entry, path: String(i) }, body }; + })) seen.push(transfer.entry.path); + expect(seen).toEqual(Array.from({ length: 30 }, (_, i) => String(i))); + expect(peak).toBe(4); + expect(opened).toBe(0); + }); + it("handles a prefetched stream failure before the consumer reaches it", async () => { + const source = new Readable({ read() {} }); + await expect((async () => { + for await (const transfer of prefetchWorkFiles([0, 1], async (i) => ({ + entry, body: i === 0 ? Readable.from([]) : source, + }))) { + if (transfer.body === source) for await (const _chunk of source) { /* consume */ } + else { + source.destroy(new Error("queued response failed")); + await new Promise((resolve) => setImmediate(resolve)); + } + } + })()).rejects.toThrow("queued response failed"); + expect(source.destroyed).toBe(true); + }); + it("closes pending responses after consumer cancellation and asynchronous open failures", async () => { + const streams: Readable[] = []; + const open = async (i: number) => { + await new Promise((resolve) => setTimeout(resolve, i * 2)); + if (i === 2) throw new Error("storage unavailable"); + const body = Readable.from([]); streams.push(body); + return { entry, body }; + }; + for await (const _transfer of prefetchWorkFiles([0, 1, 2, 3, 4], open)) break; + expect(streams).toHaveLength(3); + expect(streams.every((stream) => stream.destroyed)).toBe(true); + streams.length = 0; + await expect((async () => { + for await (const _transfer of prefetchWorkFiles([0, 1, 2, 3, 4], open)) { /* consume */ } + })()).rejects.toThrow("storage unavailable"); + expect(streams.every((stream) => stream.destroyed)).toBe(true); + }); +}); diff --git a/server/src/__tests__/work-folder-transport.test.ts b/server/src/__tests__/work-folder-transport.test.ts index 8f80a9f6d1..8eff3fc9a6 100644 --- a/server/src/__tests__/work-folder-transport.test.ts +++ b/server/src/__tests__/work-folder-transport.test.ts @@ -68,6 +68,108 @@ describe("sandbox work folder transport with real Node and Git", () => { const files = await transport.scan(dir); expect(files.find((file) => file.path === entry.path)).toEqual(entry); }); + it("batches many files and large files with bounded stdin, preserving empty files and executable modes", async () => { + const dir = await root(), staging = await root(); + const execute = vi.fn(async (input: Parameters[0]) => { + expect(Buffer.byteLength(input.stdin ?? "")).toBeLessThanOrEqual(8 * 1024 * 1024); + expect(Buffer.byteLength((input.args ?? []).join(" "))).toBeLessThan(120 * 1024); + return localTestWorkFolderRunner.execute(input); + }); + const fast = workFolderTransport({ execute, supportsSingleStreamStdinProgress: true }); + const data = Buffer.alloc(5 * 1024 * 1024 + 17, "x"); + const emptyHash = createHash("sha256").digest("hex"); + await fast.writeMany(dir, staging, (async function* () { + yield { entry: { path: "empty-dir", kind: "directory" as const, byteSize: 0, sha256: null, executable: false } }; + for (let i = 0; i < 169; i++) yield { entry: { path: `nested/empty-${i}`, kind: "file" as const, + byteSize: 0, sha256: emptyHash, executable: false }, body: Readable.from([]) }; + yield { entry: { path: "large", kind: "file" as const, byteSize: data.length, + sha256: createHash("sha256").update(data).digest("hex"), executable: true }, body: Readable.from([data]) }; + })()); + expect(execute).toHaveBeenCalledTimes(3); + expect(await readFile(path.join(dir, "large"))).toEqual(data); + const listed = await fast.scan(dir); + expect(listed.filter((entry) => entry.kind === "file")).toHaveLength(170); + expect(listed.find((entry) => entry.path === "large")?.executable).toBe(true); + expect(listed.find((entry) => entry.path === "empty-dir")?.kind).toBe("directory"); + execute.mockClear(); + const chunks = []; + for await (const chunk of fast.read(dir, "large", data.length)) chunks.push(chunk); + expect(Buffer.concat(chunks)).toEqual(data); + expect(execute).toHaveBeenCalledTimes(6); + }, 60_000); + it("reads small files in one confined batch and rejects short or oversized batches", async () => { + const dir = await root(), outside = await root(); + const execute = vi.fn(localTestWorkFolderRunner.execute); + const fast = workFolderTransport({ execute, supportsSingleStreamStdinProgress: true }); + const entries = []; + for (let i = 0; i < 32; i++) { + const data = Buffer.alloc(i === 0 ? 0 : 1024, i); + await writeFile(path.join(dir, String(i)), data); + entries.push({ path: String(i), kind: "file" as const, byteSize: data.length, + sha256: createHash("sha256").update(data).digest("hex"), executable: false }); + } + const buffers = await fast.readBatch!(dir, entries); + expect(execute).toHaveBeenCalledTimes(1); + expect(buffers.map((buffer) => createHash("sha256").update(buffer).digest("hex"))) + .toEqual(entries.map((entry) => entry.sha256)); + await writeFile(path.join(dir, "1"), "short"); + await expect(fast.readBatch!(dir, entries)).rejects.toThrow("changed during transfer"); + await writeFile(path.join(outside, "private"), "secret"); + await symlink(path.join(outside, "private"), path.join(dir, "link")); + await expect(fast.readBatch!(dir, [{ ...entries[0]!, path: "link", byteSize: 6 }])).rejects.toThrow("symlink_not_allowed"); + execute.mockClear(); + await expect(fast.readBatch!(dir, [{ ...entries[0]!, byteSize: 1024 * 1024 + 1 }])).rejects.toThrow("byte limit"); + await expect(fast.readBatch!(dir, Array.from({ length: 65 }, () => entries[0]!))).rejects.toThrow("entry limit"); + expect(execute).not.toHaveBeenCalled(); + expect(transport.readBatch).toBeUndefined(); + }); + it("persists publication intent before dispatch and aborts if that persistence fails", async () => { + const execute = vi.fn(); + const fast = workFolderTransport({ execute, supportsSingleStreamStdinProgress: true }); + const entry = { path: "file", kind: "file" as const, byteSize: 3, + sha256: createHash("sha256").update("new").digest("hex"), executable: false }; + const beforePublish = vi.fn(async (entries) => { + expect(entries).toEqual([entry]); + expect(execute).not.toHaveBeenCalled(); + throw new Error("provenance unavailable"); + }); + const body = Readable.from(["new"]); + await expect(fast.writeMany("/task", "/staging", (async function* () { yield { entry, body }; })(), beforePublish)) + .rejects.toThrow("provenance unavailable"); + expect(beforePublish).toHaveBeenCalledTimes(1); + expect(execute).not.toHaveBeenCalled(); + expect(body.destroyed).toBe(true); + }); + it("keeps previous files on bad hashes and never follows a symlink during bulk publication", async () => { + const dir = await root(), staging = await root(), outside = await root(); + const fast = workFolderTransport({ ...localTestWorkFolderRunner, supportsSingleStreamStdinProgress: true }); + await writeFile(path.join(dir, "existing"), "old"); + const entry = { path: "existing", kind: "file" as const, byteSize: 3, + sha256: createHash("sha256").update("different").digest("hex"), executable: false }; + await expect(fast.write(dir, staging, entry, Readable.from(["new"]))).rejects.toThrow("content_changed_during_transfer"); + expect(await readFile(path.join(dir, "existing"), "utf8")).toBe("old"); + await writeFile(path.join(outside, "private"), "secret"); + await symlink(outside, path.join(dir, "escape")); + const body = Readable.from(["new"]); + await expect(fast.write(dir, staging, { ...entry, path: "escape/private", + sha256: createHash("sha256").update("new").digest("hex") }, body)).rejects.toThrow("symlink_not_allowed"); + expect(await readFile(path.join(outside, "private"), "utf8")).toBe("secret"); + expect(body.destroyed).toBe(true); + const invalid = Readable.from(["new"]); + await expect(fast.write(dir, staging, { ...entry, path: "../private" }, invalid)).rejects.toThrow(); + expect(invalid.destroyed).toBe(true); + }); + it("does not retry a lost batch response or publish an incomplete source", async () => { + const execute = vi.fn().mockRejectedValue(new Error("socket hang up")); + const fast = workFolderTransport({ execute, supportsSingleStreamStdinProgress: true }); + const entry = { path: "file", kind: "file" as const, byteSize: 3, + sha256: createHash("sha256").update("new").digest("hex"), executable: false }; + await expect(fast.write("/task", "/staging", entry, Readable.from(["new"]))).rejects.toThrow("socket hang up"); + expect(execute).toHaveBeenCalledTimes(1); + execute.mockClear(); + await expect(fast.write("/task", "/staging", entry, Readable.from(["n"]))).rejects.toThrow("size changed"); + expect(execute).not.toHaveBeenCalled(); + }); it("rejects links out of a work folder on scan and download", async () => { const dir = await root(); const outside = await root(); diff --git a/server/src/services/sandbox-work-folders.ts b/server/src/services/sandbox-work-folders.ts index dbab5cea0f..ab2ca0b34f 100644 --- a/server/src/services/sandbox-work-folders.ts +++ b/server/src/services/sandbox-work-folders.ts @@ -1,3 +1,4 @@ +import { prefetchWorkFiles } from "./work-folder-transfer.js"; import { createHash, randomUUID } from "node:crypto"; import path from "node:path"; import { managedAgentFiles } from "./work-folder-agent-import.js"; @@ -134,10 +135,31 @@ export async function prepareSandboxWorkFolders(input: { } await db.update(workFolders).set({ importedAt: new Date() }).where(eq(workFolders.id, folder.id)); } + async function reconcileIncoming(scope: WorkFolderScope, current: WorkTreeEntry[]) { + const targets = baselines[`incoming:${scope}`]; + const removed = baselines[`incomingRemoved:${scope}`]; + if (!targets && !removed) return; + const before = new Map((baselines[scope] ?? []).map((entry) => [entry.path, entry])); + const observed = new Map(current.map((entry) => [entry.path, entry])); + // A failed command may have published some imports or removed some files. + // Only adopt effects actually observed on disk; different bytes remain + // genuine local edits and must still pass through outgoing synchronization. + for (const entry of removed ?? []) if (!observed.has(entry.path)) before.delete(entry.path); + for (const entry of targets ?? []) { + if (signature(observed.get(entry.path)) === signature(entry)) before.set(entry.path, entry); + } + baselines[scope] = [...before.values()]; + delete baselines[`incoming:${scope}`]; + delete baselines[`incomingRemoved:${scope}`]; + // Persist the reconciled baseline and remove its provenance atomically, + // before any outgoing write can be accepted by the shared collection. + await saveState("saving"); + } async function outgoing(scope: WorkFolderScope) { const folder = folders[scope]; if (!folder) return; const current = await transport.scan(paths[scope]!); + await reconcileIncoming(scope, current); const before = new Map((baselines[scope] ?? []).map((entry) => [entry.path, entry])); const after = new Map(current.map((entry) => [entry.path, entry])); async function operation(filePath: string, nextSignature: string, apply: (id: string) => Promise, accept: () => void) { @@ -186,25 +208,48 @@ export async function prepareSandboxWorkFolders(input: { cursor = page.nextCursor ?? undefined; } while (cursor); const desired = new Map(saved.map((entry) => [entry.path, entry])); - // Remove stale children before replacing their parent directory with a file. - for (const entry of [...current.values()].sort((a, b) => b.path.length - a.path.length)) { - if (!desired.has(entry.path) || desired.get(entry.path)!.kind !== entry.kind) { - await transport.remove(paths[scope]!, entry.path); - current.delete(entry.path); - } + // Persist all deletion intents before the first removal, including children + // that can disappear when a directory is replaced. A lost response must not + // turn an imported deletion into a new delete against a newer shared file. + const removed = [...current.values()].filter((entry) => !desired.has(entry.path) || desired.get(entry.path)!.kind !== entry.kind) + .sort((a, b) => b.path.length - a.path.length); + if (removed.length) { + baselines[`incomingRemoved:${scope}`] = removed; + await saveState("starting"); } - for (const entry of saved.sort((a, b) => a.path.length - b.path.length)) { - if (signature(current.get(entry.path)) === signature(entry)) continue; - if (entry.kind === "directory") await transport.mkdir(paths[scope]!, entry.path); - else { - const result = await svc.content(folder, entry.path); - // A shared file can change after listing. Validate and baseline the - // version opened by content(), whose metadata and stream belong together. - Object.assign(entry, { byteSize: result.file.byteSize, sha256: result.file.sha256, executable: result.file.executable }); - try { await transport.write(paths[scope]!, staging, entry, result.stream); } finally { result.stream.destroy(); } - } + for (const entry of removed) { + await transport.remove(paths[scope]!, entry.path); + current.delete(entry.path); } + const changed = saved.sort((a, b) => a.path.length - b.path.length) + .filter((entry) => signature(current.get(entry.path)) !== signature(entry)); + await transport.writeMany(paths[scope]!, staging, prefetchWorkFiles(changed, async (entry) => { + if (entry.kind === "directory") return { entry }; + const result = await svc.content(folder, entry.path); + // A shared file can change after listing. Validate and baseline the + // version opened by content(), whose metadata and stream belong together. + Object.assign(entry, { byteSize: result.file.byteSize, sha256: result.file.sha256, executable: result.file.executable }); + return { entry, body: result.stream }; + }), async (entries) => { + const targets = new Map((baselines[`incoming:${scope}`] ?? []).map((entry) => [entry.path, entry])); + for (const entry of entries) { + targets.set(entry.path, entry); + // Publishing nested files can create parents absent from the listing. + // Record those directory imports too, so they cannot be mistaken for + // agent-created directories after an interrupted batch. + let parent = path.posix.dirname(entry.path); + while (parent !== ".") { + if (!targets.has(parent)) targets.set(parent, { path: parent, kind: "directory", byteSize: 0, sha256: null, executable: false }); + parent = path.posix.dirname(parent); + } + } + baselines[`incoming:${scope}`] = [...targets.values()]; + await saveState("starting"); + }); baselines[scope] = saved; + delete baselines[`incoming:${scope}`]; + delete baselines[`incomingRemoved:${scope}`]; + await saveState("starting"); } const bindings: Array<{ binding: typeof taskRepositoryBindings.$inferSelect; root: string }> = []; @@ -340,10 +385,15 @@ export async function prepareSandboxWorkFolders(input: { .where(eq(workFolderRuns.runId, input.runId)); if (run?.refreshRequested) { // The successful final flush protects edits before incoming refresh. - await assertBindings(); - for (const scope of WORK_FOLDER_SCOPES) await incoming(scope); - await db.update(workFolderRuns).set({ refreshRequested: false, baselines, updatedAt: new Date() }) - .where(eq(workFolderRuns.runId, input.runId)); + try { + await assertBindings(); + for (const scope of WORK_FOLDER_SCOPES) await incoming(scope); + await db.update(workFolderRuns).set({ refreshRequested: false, baselines, updatedAt: new Date() }) + .where(eq(workFolderRuns.runId, input.runId)); + } catch (error) { + await saveState("failed", "Work folder refresh failed; existing files were retained"); + throw error; + } } // Publish resume identity after data is durable, before releasing a turn. await beforeCompletion?.(); diff --git a/server/src/services/scripts/work-folder-io.mjs b/server/src/services/scripts/work-folder-io.mjs index f4f3615415..f2a5209d0f 100644 --- a/server/src/services/scripts/work-folder-io.mjs +++ b/server/src/services/scripts/work-folder-io.mjs @@ -6,7 +6,7 @@ import os from "node:os"; import { createHash } from "node:crypto"; import { execFileSync } from "node:child_process"; -const input = JSON.parse(Buffer.from(process.argv[1], "base64").toString("utf8")); +let input = JSON.parse(Buffer.from(process.argv[1], "base64").toString("utf8")); const MAX_CHUNK = 256 * 1024; const MAX_ENTRIES = 100_000; function safeRelative(value) { @@ -129,6 +129,8 @@ function scan() { return results.sort((a, b) => a.path.localeCompare(b.path)); } +function execute(request) { +input = request; let result; if (input.operation === "home") { result = { home: os.homedir() }; @@ -147,7 +149,9 @@ if (input.operation === "home") { else if (input.operation === "read") { const fd = checked(full(input.path)); try { - const buffer = Buffer.alloc(MAX_CHUNK); + const length = input.length ?? MAX_CHUNK; + if (!Number.isSafeInteger(length) || length < 1 || length > 1024 * 1024) throw new Error("invalid_read_length"); + const buffer = Buffer.alloc(length); const count = fs.readSync(fd, buffer, 0, buffer.length, input.offset); result = { data: buffer.subarray(0, count).toString("base64") }; } finally { fs.closeSync(fd); } @@ -200,4 +204,43 @@ if (input.operation === "home") { result = {}; } else throw new Error("unknown_operation"); } -process.stdout.write(JSON.stringify(result)); +return result; +} + +const request = input; +if (request.operation === "read-batch") { + if (!Array.isArray(request.entries) || request.entries.length > 64) throw new Error("invalid_read_batch"); + let bytes = 0; + for (const entry of request.entries) { + safeRelative(entry.path); + if (!Number.isSafeInteger(entry.byteSize) || entry.byteSize < 0) throw new Error("invalid_read_length"); + bytes += entry.byteSize; + if (bytes > 1024 * 1024) throw new Error("read_batch_too_large"); + } + const results = request.entries.map((entry) => execute({ operation: "read", root: request.root, + path: entry.path, offset: 0, length: Math.max(1, entry.byteSize) })); + process.stdout.write(JSON.stringify(results)); +} else if (request.operation === "batch") { + // The provider transports stdin as one bounded upload. The roots stay in the + // host-authored argv envelope, so batch contents cannot redirect operations. + const chunks = []; + let size = 0; + const buffer = Buffer.alloc(64 * 1024); + let count; + while ((count = fs.readSync(0, buffer, 0, buffer.length, null)) > 0) { + size += count; + if (size > 8 * 1024 * 1024) throw new Error("batch_too_large"); + chunks.push(Buffer.from(buffer.subarray(0, count))); + } + const operations = JSON.parse(Buffer.concat(chunks).toString("utf8")); + if (!Array.isArray(operations) || operations.length > 512) throw new Error("invalid_batch"); + for (const operation of operations) { + if (!operation || !["write", "publish", "mkdir"].includes(operation.operation)) throw new Error("invalid_batch_operation"); + // Never accept roots, source roots or filesystem commands from the body. + execute({ ...operation, root: operation.operation === "write" ? request.stagingRoot : request.root, + stagingRoot: request.stagingRoot }); + } + process.stdout.write(JSON.stringify({ completed: operations.length })); +} else { + process.stdout.write(JSON.stringify(execute(request))); +} diff --git a/server/src/services/work-folder-read-cache.ts b/server/src/services/work-folder-read-cache.ts new file mode 100644 index 0000000000..6b1ea222ab --- /dev/null +++ b/server/src/services/work-folder-read-cache.ts @@ -0,0 +1,78 @@ +import { Readable } from "node:stream"; +import type { WorkTreeEntry } from "./work-folder-transport.js"; + +const MAX_BATCH_BYTES = 1024 * 1024; +const MAX_BATCH_ENTRIES = 64; +const MAX_CACHED_BATCHES = 4; + +/** + * A checkpoint-local cache, not a snapshot or a retry source. The caller limits + * concurrent readers to four. At most four 1 MiB batches stay cached; evicted + * batches held by those active readers can add at most another 4 MiB. Larger + * files use the separately bounded streaming fallback. Metadata is O(entries). + */ +export function createWorkFolderReadCache( + entries: WorkTreeEntry[], + load: (entries: WorkTreeEntry[]) => Promise, + fallback: (entry: WorkTreeEntry) => Readable, +) { + const locations = new Map(); + let batch: WorkTreeEntry[] = []; + let batchBytes = 0; + for (const entry of entries) { + if (entry.kind !== "file" || entry.linkTarget || entry.byteSize > MAX_BATCH_BYTES) continue; + if (batch.length >= MAX_BATCH_ENTRIES || batchBytes + entry.byteSize > MAX_BATCH_BYTES) { + batch = []; batchBytes = 0; + } + locations.set(entry.path, { batch, index: batch.length }); + batch.push(entry); batchBytes += entry.byteSize; + } + const cached = new Map>(); + const readPaths = new Set(); + let cleared = false; + + function getBatch(group: WorkTreeEntry[]) { + let result = cached.get(group); + if (result) { + cached.delete(group); cached.set(group, result); + return result; + } + result = Promise.resolve().then(() => load(group)).then((buffers) => { + if (buffers.length !== group.length || buffers.some((buffer, index) => + !Buffer.isBuffer(buffer) || buffer.length !== group[index]!.byteSize)) { + throw new Error("Work folder batch content does not match its entries"); + } + return buffers; + }).catch((error: unknown) => { + if (cached.get(group) === result) cached.delete(group); + throw error; + }); + cached.set(group, result); + while (cached.size > MAX_CACHED_BATCHES) cached.delete(cached.keys().next().value!); + return result; + } + + function read(entry: WorkTreeEntry): Readable { + const repeated = readPaths.has(entry.path); + readPaths.add(entry.path); + const location = locations.get(entry.path); + // Mark each source request, even if its stream is never consumed. A retry + // must reopen the physical file instead of replaying possibly stale bytes. + return Readable.from((async function* () { + if (cleared || repeated || !location) { + const source = fallback(entry); + try { yield* source; } finally { source.destroy(); } + return; + } + const buffers = await getBatch(location.batch); + yield buffers[location.index]!; + })()); + } + function clear() { + cleared = true; + cached.clear(); + locations.clear(); + readPaths.clear(); + } + return { read, clear }; +} diff --git a/server/src/services/work-folder-repositories.ts b/server/src/services/work-folder-repositories.ts index 1822f6e7d8..2122f02152 100644 --- a/server/src/services/work-folder-repositories.ts +++ b/server/src/services/work-folder-repositories.ts @@ -1,3 +1,5 @@ +import { createWorkFolderReadCache } from "./work-folder-read-cache.js"; +import { prefetchWorkFiles } from "./work-folder-transfer.js"; import { createHash, randomUUID } from "node:crypto"; import { and, eq, inArray, isNull } from "drizzle-orm"; import { taskRepositoryBindings, workFolderObjects, type Db } from "@paperclipai/db"; @@ -35,17 +37,47 @@ export function workFolderRepositoryService(db: Db, storage: StorageProvider, tr const prefix = `${binding.companyId}/task-repositories/${binding.id}/`; if (!knownByBinding.has(binding.id) && binding.checkpointKey) await loadManifest(binding); const known = knownByBinding.get(binding.id) ?? new Set(); - const files: Array = []; - for (const entry of entries) { - const objectKey = entry.kind === "file" && !entry.linkTarget ? `${prefix}blobs/${entry.sha256}` : null; - if (objectKey && !known.has(objectKey)) await registerWorkFolderObject(db, storage, { objectKey, companyId: binding.companyId, repositoryBindingId: binding.id }); - if (objectKey && !known.has(objectKey) && !(await storage.headObject({ objectKey })).exists) { - await uploadWorkFolderObject(storage, { objectKey, contentType: "application/octet-stream", - contentLength: entry.byteSize, sha256: entry.sha256!, - createSource: () => transport.read(root, entry.path, entry.byteSize) }); + // Keep manifest order independent of transfer completion, and send each + // content-addressed blob only once even when multiple paths share bytes. + const files = entries.map((entry) => ({ ...entry, + objectKey: entry.kind === "file" && !entry.linkTarget ? `${prefix}blobs/${entry.sha256}` : null, + })); + const unknown = new Map(); + for (const entry of files) { + if (entry.objectKey && !known.has(entry.objectKey) && !unknown.has(entry.objectKey)) { + unknown.set(entry.objectKey, entry); } - files.push({ ...entry, objectKey }); } + const objects = [...unknown]; + const readCache = transport.readBatch ? createWorkFolderReadCache([...unknown.values()], + (entries) => transport.readBatch!(root, entries), + (entry) => transport.read(root, entry.path, entry.byteSize)) : undefined; + let next = 0; + let failure: { error: unknown } | undefined; + try { + await Promise.all(Array.from({ length: Math.min(4, objects.length) }, async () => { + while (!failure) { + const object = objects[next++]; + if (!object) return; + const [objectKey, entry] = object; + try { + await registerWorkFolderObject(db, storage, { objectKey, companyId: binding.companyId, repositoryBindingId: binding.id }); + if (failure) return; + const { exists } = await storage.headObject({ objectKey }); + if (!exists && !failure) { + await uploadWorkFolderObject(storage, { objectKey, contentType: "application/octet-stream", + contentLength: entry.byteSize, sha256: entry.sha256!, + createSource: () => readCache ? readCache.read(entry) : transport.read(root, entry.path, entry.byteSize) }); + } + } catch (error) { + // Stop scheduling after the first error, but drain the other workers + // before returning. Their streaming PUTs must not outlive this save. + failure ??= { error }; + } + } + })); + } finally { readCache?.clear(); } + if (failure) throw failure.error; // Never publish a torn Git index/worktree snapshot as a completed save. if (signature(await transport.scan(root, true)) !== digest) throw new Error("Repository changed during checkpoint; retry required"); const checkpointKey = `${prefix}checkpoints/${randomUUID()}.json`; @@ -109,14 +141,17 @@ export function workFolderRepositoryService(db: Db, storage: StorageProvider, tr if (!binding.checkpointKey) return false; const manifest = await loadManifest(binding); await transport.mkdirRoot(root); - for (const entry of manifest.files) { - if (entry.kind === "directory") await transport.mkdir(root, entry.path); - else if (entry.linkTarget) await transport.symlink(root, stagingRoot, entry); - else { + await transport.writeMany(root, stagingRoot, prefetchWorkFiles( + manifest.files.filter((entry) => !entry.linkTarget), async (entry) => { + if (entry.kind === "directory") return { entry }; if (!entry.objectKey) throw new Error("Repository checkpoint file is missing"); const result = await storage.getObject({ objectKey: entry.objectKey }); - try { await transport.write(root, stagingRoot, entry, result.stream); } finally { result.stream.destroy(); } - } + return { entry, body: result.stream }; + })); + // Restore links only after ordinary files. No transfer follows a link as + // a parent, and symlink() still confines its target to this repository. + for (const entry of manifest.files) { + if (entry.linkTarget) await transport.symlink(root, stagingRoot, entry); } return true; } diff --git a/server/src/services/work-folder-transfer.ts b/server/src/services/work-folder-transfer.ts new file mode 100644 index 0000000000..6312c8abf9 --- /dev/null +++ b/server/src/services/work-folder-transfer.ts @@ -0,0 +1,36 @@ +import type { WorkFileTransfer } from "./work-folder-transport.js"; + +// Open a few storage responses ahead without buffering their bodies. Drain all +// pending opens on failure so abandoned HTTP response streams are closed too. +export async function* prefetchWorkFiles( + entries: Iterable, + open: (entry: T) => Promise, +): AsyncGenerator { + type Result = { value: WorkFileTransfer } | { error: unknown }; + const iterator = entries[Symbol.iterator](); + const pending: Array> = []; + function enqueue() { + const next = iterator.next(); + if (!next.done) pending.push(Promise.resolve().then(() => open(next.value)) + .then((value): Result => { + // A response can fail while queued, before its async iterator exists. + // Keep that error handled; consuming the stream still throws it. + value.body?.on("error", () => {}); + return { value }; + }, (error): Result => ({ error }))); + } + try { + for (let i = 0; i < 4; i++) enqueue(); + while (pending.length) { + const result = await pending.shift()!; + if ("error" in result) throw result.error; + try { yield result.value; } finally { result.value.body?.destroy(); } + enqueue(); + } + } finally { + for (const result of await Promise.all(pending)) { + if ("value" in result) result.value.body?.destroy(); + } + iterator.return?.(); + } +} diff --git a/server/src/services/work-folder-transport.ts b/server/src/services/work-folder-transport.ts index e6342e9d4e..f375d8dd9f 100644 --- a/server/src/services/work-folder-transport.ts +++ b/server/src/services/work-folder-transport.ts @@ -15,6 +15,9 @@ let source: Promise | undefined; // Requests are base64 encoded twice (file bytes, then JSON). Stay below // Linux's 128 KiB single-argument limit, including a provider shell wrapper. const WRITE_CHUNK_BYTES = 48 * 1024; +const BATCH_BYTES = 4 * 1024 * 1024; +const BATCH_OPERATIONS = 256; +export type WorkFileTransfer = { entry: WorkTreeEntry; body?: Readable }; function transientReadFailure(error: unknown) { if (!(error instanceof Error)) return false; @@ -24,15 +27,19 @@ function transientReadFailure(error: unknown) { } export function workFolderTransport(runner: CommandManagedRuntimeRunner) { - async function command(input: Record): Promise { + // Native file-sync providers already support bounded stdin uploads. SSH + // runners explicitly advertise streaming stdin. Other providers retain the + // small-argv transport without assuming additional capabilities. + const bulkStdin = Boolean(runner.syncIn && runner.syncOut) || runner.supportsSingleStreamStdinProgress === true; + async function command(input: Record, stdin?: string): Promise { source ??= readFile(new URL("./scripts/work-folder-io.mjs", import.meta.url), "utf8"); const args = ["--input-type=module", "-e", await source, Buffer.from(JSON.stringify(input)).toString("base64")]; - const readOnly = ["home", "scan", "read"].includes(String(input.operation)); + const readOnly = ["home", "scan", "read", "read-batch"].includes(String(input.operation)); const deadline = Date.now() + 120_000; let result; for (let attempt = 0; ; attempt++) { try { - result = await runner.execute({ command: "node", args, bypassSession: true, + result = await runner.execute({ command: "node", args, ...(stdin === undefined ? {} : { stdin }), bypassSession: true, timeoutMs: Math.max(1, deadline - Date.now()) }); break; } catch (error) { @@ -57,7 +64,8 @@ export function workFolderTransport(runner: CommandManagedRuntimeRunner) { validateWorkFilePath(filePath); return Readable.from((async function* () { for (let offset = 0; offset < byteSize;) { - const result = z.object({ data: z.string().max(350_000) }).parse(await command({ operation: "read", root, path: filePath, offset })); + const length = bulkStdin ? 1024 * 1024 : 256 * 1024; + const result = z.object({ data: z.string().max(Math.ceil(length / 3) * 4) }).parse(await command({ operation: "read", root, path: filePath, offset, length })); const bytes = Buffer.from(result.data, "base64"); if (bytes.length === 0 || offset + bytes.length > byteSize) throw new Error("Work file changed during transfer"); offset += bytes.length; @@ -65,11 +73,30 @@ export function workFolderTransport(runner: CommandManagedRuntimeRunner) { } })()); } - async function write(root: string, stagingRoot: string, entry: WorkTreeEntry, body: Readable) { + async function readBatch(root: string, entries: WorkTreeEntry[]) { + if (entries.length > 64) throw new Error("Work folder read batch exceeds entry limit"); + let bytes = 0; + for (const entry of entries) { + entrySchema.parse(entry); + if (entry.kind !== "file" || entry.linkTarget) throw new Error("Work folder read batch requires regular files"); + bytes += entry.byteSize; + } + if (bytes > 1024 * 1024) throw new Error("Work folder read batch exceeds byte limit"); + if (!entries.length) return []; + const result = z.array(z.object({ data: z.string().max(Math.ceil(1024 * 1024 / 3) * 4) })).max(64) + .parse(await command({ operation: "read-batch", root, entries: entries.map(({ path, byteSize }) => ({ path, byteSize })) })); + if (result.length !== entries.length) throw new Error("Work folder read batch did not complete"); + return result.map(({ data }, index) => { + const buffer = Buffer.from(data, "base64"); + if (buffer.length !== entries[index]!.byteSize) throw new Error("Work file changed during transfer"); + return buffer; + }); + } + async function writeChunks(root: string, stagingRoot: string, entry: WorkTreeEntry, body: Readable) { const stagingPath = randomUUID(); let offset = 0; for await (const value of body) { - const chunk = Buffer.from(value); + const chunk = Buffer.isBuffer(value) ? value : Buffer.from(value); for (let start = 0; start < chunk.length; start += WRITE_CHUNK_BYTES) { const bytes = chunk.subarray(start, start + WRITE_CHUNK_BYTES); await command({ operation: "write", root: stagingRoot, path: stagingPath, offset, data: bytes.toString("base64") }); @@ -81,7 +108,66 @@ export function workFolderTransport(runner: CommandManagedRuntimeRunner) { await command({ operation: "publish", root, stagingRoot, stagingPath, path: entry.path, sha256: entry.sha256, executable: entry.executable }); } - return { home, scan, read, write, + async function writeMany(root: string, stagingRoot: string, transfers: AsyncIterable, + beforePublish?: (entries: WorkTreeEntry[]) => Promise) { + let publishedEntries: WorkTreeEntry[] = []; + let operations: Array> = []; + let bufferedBytes = 0; + async function flush() { + if (!operations.length) return; + const body = JSON.stringify(operations); + if (Buffer.byteLength(body) > 8 * 1024 * 1024) throw new Error("Work folder batch exceeds transfer limit"); + if (publishedEntries.length) await beforePublish?.(publishedEntries); + const result = z.object({ completed: z.number().int() }).parse(await command({ operation: "batch", root, stagingRoot }, body)); + if (result.completed !== operations.length) throw new Error("Work folder batch did not complete"); + operations = []; bufferedBytes = 0; publishedEntries = []; + } + async function append(operation: Record, bytes = 0, entry?: WorkTreeEntry) { + if (bufferedBytes + bytes > BATCH_BYTES || operations.length >= BATCH_OPERATIONS) await flush(); + operations.push(operation); bufferedBytes += bytes; + if (entry) publishedEntries.push({ ...entry }); + } + for await (const { entry, body } of transfers) { + try { + entrySchema.parse(entry); + if (entry.linkTarget) throw new Error("Work folder batch cannot materialize symlinks"); + if (entry.kind === "directory") { + if (body) throw new Error("Directory transfer cannot have a body"); + if (bulkStdin) await append({ operation: "mkdir", path: entry.path }, 0, entry); + else { + await beforePublish?.([entry]); + await command({ operation: "mkdir", root, path: entry.path }); + } + continue; + } + if (!body) throw new Error("Work file transfer source is missing"); + if (!bulkStdin) { + await beforePublish?.([entry]); + await writeChunks(root, stagingRoot, entry, body); + continue; + } + const stagingPath = randomUUID(); + let offset = 0; + for await (const value of body) { + const chunk = Buffer.isBuffer(value) ? value : Buffer.from(value); + for (let start = 0; start < chunk.length; start += 256 * 1024) { + const bytes = chunk.subarray(start, start + 256 * 1024); + if (offset + bytes.length > entry.byteSize) throw new Error("Work file size changed during transfer"); + await append({ operation: "write", path: stagingPath, offset, data: bytes.toString("base64") }, bytes.length); + offset += bytes.length; + } + } + if (offset === 0) await append({ operation: "write", path: stagingPath, offset: 0, data: "" }); + if (offset !== entry.byteSize) throw new Error("Work file size changed during transfer"); + await append({ operation: "publish", stagingPath, path: entry.path, sha256: entry.sha256, executable: entry.executable }, 0, entry); + } finally { body?.destroy(); } + } + if (bulkStdin) await flush(); + } + async function write(root: string, stagingRoot: string, entry: WorkTreeEntry, body: Readable) { + return writeMany(root, stagingRoot, (async function* () { yield { entry, body }; })()); + } + return { home, scan, read, readBatch: bulkStdin ? readBatch : undefined, write, writeMany, moveRoot: async (source: string, root: string) => { await command({ operation: "move-root", source, root }); }, symlink: async (root: string, stagingRoot: string, entry: WorkTreeEntry) => { await command({ operation: "symlink", root, stagingRoot, stagingPath: randomUUID(), path: entry.path, linkTarget: entry.linkTarget }); }, mkdirRoot: async (root: string) => { await command({ operation: "mkdir-root", root }); },