fix: retry interrupted work-folder uploads with verified fresh streams

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

server/src/__tests__/sandbox-work-folders.test.ts

@@ -1,6 +1,7 @@
 import { afterAll, beforeAll, describe, expect, it, vi } from "vitest";
 import { execFile } from "node:child_process";
 import { promisify } from "node:util";
+import { Readable } from "node:stream";
 import { createHash, randomUUID } from "node:crypto";
 import { eq, inArray, sql } from "drizzle-orm";
 import fs from "node:fs/promises";
@@ -347,6 +348,54 @@ 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 () => {
+    const leaseId = randomUUID();
+    const run = await prepare(path.join(root, "replayed-repository-upload"), leaseId);
+    await run.flush();
+    const filename = path.join(run.primaryRepo, "retry-upload");
+    const content = "replay the entire repository file";
+    await fs.writeFile(filename, content);
+    const targetHash = createHash("sha256").update(content).digest("hex");
+    const put = storage.putObject.bind(storage);
+    const streams: Readable[] = [];
+    let attempts = 0;
+    const recover = vi.spyOn(storage, "putObject").mockImplementation(async (input) => {
+      if (input.objectKey.endsWith("/" + targetHash)) {
+        streams.push(input.body as Readable);
+        if (++attempts === 1) {
+          await (input.body as Readable).iterator({ destroyOnReturn: false }).next();
+          throw Object.assign(new Error("socket hang up"), { code: "ECONNRESET" });
+        }
+      }
+      return put(input);
+    });
+    try { await run.flush(); } finally { recover.mockRestore(); }
+    expect(attempts).toBe(2);
+    expect(new Set(streams).size).toBe(2);
+    expect(streams.every((stream) => stream.destroyed)).toBe(true);
+    const bindingId = run.manifest.repositories[0]!.bindingId;
+    const [before] = await db.select().from(taskRepositoryBindings).where(eq(taskRepositoryBindings.id, bindingId));
+    await fs.writeFile(filename, "retain this unsaved edit");
+    const nextHash = createHash("sha256").update("retain this unsaved edit").digest("hex");
+    let failedAttempts = 0;
+    const fail = vi.spyOn(storage, "putObject").mockImplementation(async (input) => {
+      if (input.objectKey.endsWith("/" + nextHash)) {
+        failedAttempts++;
+        for await (const _chunk of input.body as Readable) { /* Lost response after consuming the body. */ }
+        throw Object.assign(new Error("socket hang up"), { code: "ECONNRESET" });
+      }
+      return put(input);
+    });
+    try {
+      await expect(run.stop()).rejects.toThrow("socket hang up");
+      expect(failedAttempts).toBe(3);
+      const [after] = await db.select().from(taskRepositoryBindings).where(eq(taskRepositoryBindings.id, bindingId));
+      expect(after!.checkpointKey).toBe(before!.checkpointKey);
+      expect(await fs.readFile(filename, "utf8")).toBe("retain this unsaved edit");
+      expect(await retainUnsavedWorkFolderLease(db, { id: leaseId, companyId })).toBe(true);
+    } finally { fail.mockRestore(); active.splice(active.indexOf(run), 1); }
+  }, 120_000);
+
   it("does not publish a partial repository checkpoint and recovers a failed final save in a new run", async () => {
     const leaseId = randomUUID();
     const run = await prepare(path.join(root, "interrupted-checkpoint"), leaseId);

server/src/__tests__/work-folder-upload.test.ts

@@ -0,0 +1,109 @@
+import { createHash } from "node:crypto";
+import { Readable } from "node:stream";
+import { describe, expect, it, vi } from "vitest";
+import { uploadWorkFolderObject } from "../services/work-folder-upload.js";
+
+const data = Buffer.from("verified file contents\n");
+const reset = () => Object.assign(new Error("socket hang up"), { name: "TimeoutError", code: "ECONNRESET" });
+const digest = (bytes: Buffer) => createHash("sha256").update(bytes).digest("hex");
+const input = { objectKey: "company/repository/blob", contentType: "application/octet-stream", contentLength: data.length, sha256: digest(data) };
+async function bytes(body: Buffer | Readable) {
+  const chunks: Buffer[] = [];
+  for await (const chunk of body) chunks.push(Buffer.from(chunk));
+  return Buffer.concat(chunks);
+}
+
+describe("replayable verified work-folder uploads", () => {
+  it.each(["partial", "complete"])("reopens and verifies the whole source after a reset following %s consumption", async (consumption) => {
+    const sources: Readable[] = [], bodies: Readable[] = [];
+    let attempt = 0;
+    const putObject = vi.fn(async ({ body, objectKey }) => {
+      expect(objectKey).toBe(input.objectKey);
+      bodies.push(body);
+      if (++attempt === 1) {
+        if (consumption === "complete") expect(await bytes(body)).toEqual(data);
+        else await body.iterator({ destroyOnReturn: false }).next();
+        throw reset();
+      }
+      expect(await bytes(body)).toEqual(data);
+    });
+    await uploadWorkFolderObject({ putObject }, { ...input, createSource: () => {
+      const stream = Readable.from([data.subarray(0, 3), data.subarray(3)]); sources.push(stream); return stream;
+    } });
+    expect(putObject).toHaveBeenCalledTimes(2);
+    expect(new Set(bodies).size).toBe(2);
+    expect(sources.every((source) => source.destroyed)).toBe(true);
+    expect(bodies.every((body) => body.destroyed)).toBe(true);
+  });
+
+  it("waits for the previous PUT to settle before reopening a failed source", async () => {
+    let release!: () => void;
+    let failed!: () => void;
+    const held = new Promise<void>((resolve) => { release = resolve; });
+    const sourceFailed = new Promise<void>((resolve) => { failed = resolve; });
+    let attempts = 0;
+    const createSource = vi.fn(() => ++attempts === 1
+      ? Readable.from((async function* () { yield data.subarray(0, 3); failed(); throw reset(); })())
+      : Readable.from([data]));
+    const putObject = vi.fn(async ({ body }) => {
+      if (attempts === 1) { await held; throw reset(); }
+      expect(await bytes(body)).toEqual(data);
+    });
+    const upload = uploadWorkFolderObject({ putObject }, { ...input, createSource });
+    await sourceFailed;
+    // Longer than the retry delay: an unsettled old request must still own the key.
+    await new Promise((resolve) => setTimeout(resolve, 300));
+    expect(createSource).toHaveBeenCalledTimes(1);
+    release();
+    await upload;
+    expect(createSource).toHaveBeenCalledTimes(2);
+  });
+
+  it("stops after three failed attempts and destroys every source", async () => {
+    const sources: Readable[] = [];
+    const putObject = vi.fn(async ({ body }) => { await bytes(body); throw reset(); });
+    await expect(uploadWorkFolderObject({ putObject }, { ...input, createSource: () => {
+      const source = Readable.from([data]); sources.push(source); return source;
+    } })).rejects.toMatchObject({ code: "ECONNRESET" });
+    expect(putObject).toHaveBeenCalledTimes(3);
+    expect(sources.every((source) => source.destroyed)).toBe(true);
+  });
+
+  it("rejects a changed source after a reset without publishing or another retry", async () => {
+    let attempt = 0;
+    const putObject = vi.fn(async ({ body }) => { await bytes(body); throw reset(); });
+    await expect(uploadWorkFolderObject({ putObject }, { ...input, createSource: () => Readable.from([
+      ++attempt === 1 ? data : Buffer.alloc(data.length, 120),
+    ]) })).rejects.toThrow("Work file changed during upload");
+    expect(putObject).toHaveBeenCalledTimes(2);
+  });
+
+  it.each([Buffer.from("short"), Buffer.alloc(data.length + 1), Buffer.alloc(data.length, 120)])("rejects size or hash changes without retrying", async (body) => {
+    const putObject = vi.fn(async ({ body: stream }) => { await bytes(stream); });
+    await expect(uploadWorkFolderObject({ putObject }, { ...input, createSource: () => Readable.from([body]) }))
+      .rejects.toThrow("Work file changed during upload");
+    expect(putObject).toHaveBeenCalledTimes(1);
+  });
+
+  it.each([
+    Object.assign(new Error("denied"), { $metadata: { httpStatusCode: 403 } }),
+    Object.assign(new Error("missing"), { code: "ENOENT" }),
+    new Error("owner was deleted"),
+  ])("does not replay a permanent or unclassified error", async (error) => {
+    const source = Readable.from([data]);
+    const putObject = vi.fn(async () => { throw error; });
+    await expect(uploadWorkFolderObject({ putObject }, { ...input, createSource: () => source })).rejects.toBe(error);
+    expect(putObject).toHaveBeenCalledTimes(1);
+    expect(source.destroyed).toBe(true);
+  });
+
+  it("retries a classified temporary HTTP failure and supports empty files", async () => {
+    let attempt = 0;
+    const putObject = vi.fn(async ({ body }) => {
+      expect(await bytes(body)).toEqual(Buffer.alloc(0));
+      if (++attempt === 1) throw Object.assign(new Error("unavailable"), { $metadata: { httpStatusCode: 503 } });
+    });
+    await uploadWorkFolderObject({ putObject }, { ...input, contentLength: 0, sha256: digest(Buffer.alloc(0)), createSource: () => Readable.from([]) });
+    expect(putObject).toHaveBeenCalledTimes(2);
+  });
+});

server/src/__tests__/work-folders.test.ts

@@ -1,4 +1,4 @@
-import { afterAll, beforeAll, describe, expect, it } from "vitest";
+import { afterAll, beforeAll, describe, expect, it, vi } from "vitest";
 import { randomUUID } from "node:crypto";
 import { mkdtemp, rm } from "node:fs/promises";
 import path from "node:path";
@@ -97,6 +97,36 @@ describe("durable work folders", () => {
     await expect(svc.write(f, { path: "large", body: Buffer.from("large"), maxBytes: 2, operationId: "large" })).rejects.toMatchObject({ status: 413 });
     expect((await svc.list(f)).files).toHaveLength(0);
   });
+  it("replays a failed spool upload without publishing a receipt or replacing the previous file", async () => {
+    const f = await folder();
+    await svc.write(f, { path: "note", body: Buffer.from("old"), operationId: "old" });
+    const old = await svc.get(f, "note");
+    const put = storage.putObject.bind(storage);
+    const bodies: Readable[] = [];
+    const failed = vi.spyOn(storage, "putObject").mockImplementation(async (input) => {
+      bodies.push(input.body as Readable);
+      for await (const _chunk of input.body as Readable) { /* Consume the uncertain request. */ }
+      throw Object.assign(new Error("socket hang up"), { code: "ECONNRESET" });
+    });
+    try {
+      await expect(svc.write(f, { path: "note", body: Buffer.from("new"), operationId: "new" })).rejects.toThrow("socket hang up");
+      expect(failed).toHaveBeenCalledTimes(3);
+      expect(new Set(bodies).size).toBe(3);
+      expect(bodies.every((body) => body.destroyed)).toBe(true);
+      expect((await svc.get(f, "note")).objectKey).toBe(old.objectKey);
+      expect(await textContent(f, "note")).toBe("old");
+    } finally { failed.mockRestore(); }
+    let attempts = 0;
+    const recovered = vi.spyOn(storage, "putObject").mockImplementation(async (input) => {
+      await put(input);
+      if (++attempts === 1) throw Object.assign(new Error("lost response"), { code: "ECONNRESET" });
+    });
+    try {
+      expect(await svc.write(f, { path: "note", body: Buffer.from("new"), operationId: "new" })).toEqual({ applied: true });
+      expect(recovered).toHaveBeenCalledTimes(2);
+      expect(await textContent(f, "note")).toBe("new");
+    } finally { recovered.mockRestore(); }
+  });
   it("serializes conflicting parent/file creation", async () => {
     const f = await folder();
     const results = await Promise.allSettled([