fix(sandbox-adapter): stop Pi env calls when the run aborts

packages/pi/tests/actor.test.ts

@@ -1,5 +1,5 @@
 import { randomUUID } from "node:crypto";
-import { mkdtemp, readdir, readFile, rm } from "node:fs/promises";
+import { mkdir, mkdtemp, readdir, readFile, rm } from "node:fs/promises";
 import { tmpdir } from "node:os";
 import { join } from "node:path";
 import { BACKGROUND_CONTEXT } from "@earendil-works/chord/context";
@@ -56,6 +56,9 @@ const toolRuns = { safe: 0, unsafe: 0 };
 /** Sandboxes the unreachable provider was asked to create. */
 let sandboxCreates = 0;
 
+/** When set, the cold-start provider's connect waits on it, as a sandbox that is still starting does. */
+let coldStart: (() => Promise<void>) | undefined;
+
 const Todos = defineDoc<{ items: string[] }>({
 	kind: "app.todos",
 	version: 1,
@@ -234,6 +237,30 @@ function buildRegistry(mock: MockModel, root: string) {
 		},
 		connect: async () => undefined,
 	};
+	// One sandbox directory whose working directory is known before it connects, like E2B's.
+	const coldRoot = join(root, "cold-start");
+	const coldSandboxes = localSandboxProvider(coldRoot);
+	const coldStarting: SandboxProvider = {
+		name: "cold-start",
+		cwd: join(coldRoot, "box"),
+		create: async () => {
+			await mkdir(join(coldRoot, "box"), { recursive: true });
+			return "box";
+		},
+		connect: async (c, id) => {
+			await coldStart?.();
+			return coldSandboxes.connect(c, id);
+		},
+	};
+	const coldCoder = pi({
+		...models,
+		registry: (() => {
+			const coding = createRegistry();
+			coding.install(CodingTools);
+			return coding;
+		})(),
+		sandbox: coldStarting,
+	});
 	const lookups = createRegistry();
 	lookups.install(
 		defineExtension({
@@ -264,6 +291,7 @@ function buildRegistry(mock: MockModel, root: string) {
 		use: {
 			agent,
 			coder,
+			coldCoder,
 			files,
 			backoff,
 			backoffWithoutSchedules,
@@ -523,6 +551,37 @@ describe("pi actor", () => {
 		expect(JSON.stringify(read?.content)).not.toContain("hi from pi");
 	});
 
+	test("a run stopped while its sandbox starts writes nothing once the sandbox is up", async (c) => {
+		const { client } = await setupTest(c, registry);
+		const handle = client.coldCoder.getOrCreate([randomUUID()]);
+		const root = await handle.harness.root();
+		const connecting = Promise.withResolvers<void>();
+		const started = Promise.withResolvers<void>();
+		coldStart = () => {
+			started.resolve();
+			return connecting.promise;
+		};
+		try {
+			const prompt = handle.prompt("create hello.txt");
+			await started.promise;
+			const stopping = handle.conversation.abort(root.id);
+			// The abort waits for the write that holds the connect, so the test polls for Pi's abort mark before the sandbox comes up.
+			await vi.waitFor(async () => {
+				const { tasks } = await handle.harness.inspect();
+				expect(tasks.some((task) => task.record.abortRequested)).toBe(true);
+			});
+			connecting.resolve();
+			await stopping;
+			await prompt;
+		} finally {
+			coldStart = undefined;
+		}
+
+		await expect(
+			readFile(join(workdir, "cold-start", "box", "hello.txt"), "utf8"),
+		).rejects.toThrow();
+	});
+
 	test("without a sandbox, files live in the actor's database: they outlast sleep, a ranged read returns its lines, and there is no shell", async (c) => {
 		const { client } = await setupTest(c, registry);
 		const key = ["files", randomUUID()];

packages/sandbox-adapter/src/pi/database-env.ts

@@ -1,4 +1,5 @@
 import { posix } from "node:path";
+import type { Context } from "@earendil-works/chord";
 import {
 	type BinaryReader,
 	type ExecutionEnv,
@@ -64,30 +65,31 @@ export function databaseEnv(
 		absolutePath: async (path) => ok(posix.resolve(env.cwd, path)),
 		joinPath: async (parts) => ok(posix.join(...parts)),
 		canonicalPath: async (path) => ok(posix.resolve(env.cwd, path)),
-		exists: (path) =>
+		exists: (path, context) =>
 			attempt(
 				path,
+				context,
 				async (resolved) => (await entry(db, resolved)) !== undefined,
 			),
-		fileInfo: (path) =>
-			attempt(path, async (resolved) => {
+		fileInfo: (path, context) =>
+			attempt(path, context, async (resolved) => {
 				const info = await entry(db, resolved);
 				if (!info) throw notFound(resolved);
 				return info;
 			}),
-		readTextFile: (path) =>
-			attempt(path, async (resolved) =>
+		readTextFile: (path, context) =>
+			attempt(path, context, async (resolved) =>
 				new TextDecoder().decode(await readFile(db, resolved)),
 			),
-		readBinaryFile: (path) =>
-			attempt(path, (resolved) => readFile(db, resolved)),
-		openBinaryReader: (path) =>
-			attempt(path, async (resolved) => {
+		readBinaryFile: (path, context) =>
+			attempt(path, context, (resolved) => readFile(db, resolved)),
+		openBinaryReader: (path, _options, context) =>
+			attempt(path, context, async (resolved) => {
 				await requireFile(db, resolved);
 				return fileReader(db, resolved);
 			}),
-		writeFile: (path, content) =>
-			attempt(path, async (resolved) => {
+		writeFile: (path, content, context) =>
+			attempt(path, context, async (resolved) => {
 				const bytes =
 					typeof content === "string"
 						? new TextEncoder().encode(content)
@@ -128,12 +130,17 @@ export function databaseEnv(
 		cleanup: async () => {},
 	};
 
-	/** Resolves `path`, runs `operation`, and returns a thrown error as a `FileError`. */
+	/**
+	 * Resolves `path`, runs `operation` unless the run is aborted, and returns a
+	 * thrown error as a `FileError`.
+	 */
 	async function attempt<T>(
 		path: string,
+		context: Context,
 		operation: (resolved: string) => Promise<T>,
 	): Promise<Result<T, FileError>> {
 		const resolved = posix.resolve(env.cwd, path);
+		if (context.abortSignal?.aborted) return err(aborted(resolved));
 		try {
 			await ensureTable(db);
 			return ok(await operation(resolved));
@@ -214,8 +221,10 @@ function fileReader(db: FileDatabase, path: string): BinaryReader {
 		return new Uint8Array(row.chunk ?? new Uint8Array());
 	};
 	const guarded = async <T>(
+		context: Context,
 		operation: () => Promise<T>,
 	): Promise<Result<T, FileError>> => {
+		if (context.abortSignal?.aborted) return err(aborted(path));
 		try {
 			return ok(await operation());
 		} catch (error) {
@@ -225,17 +234,19 @@ function fileReader(db: FileDatabase, path: string): BinaryReader {
 		}
 	};
 	return {
-		info: () => guarded(() => requireFile(db, path)),
-		read: (offset, length) => guarded(() => range(offset, length)),
-		scanLines: ({ startLine, endLine }) =>
-			guarded(async () => {
+		info: (context) => guarded(context, () => requireFile(db, path)),
+		read: (offset, length, context) =>
+			guarded(context, () => range(offset, length)),
+		scanLines: ({ startLine, endLine }, context) =>
+			guarded(context, async () => {
 				let scanner: LineScanner;
 				try {
 					scanner = new LineScanner(startLine, endLine);
 				} catch {
 					throw new FileError("invalid", "Invalid line range", path);
 				}
 				for (let offset = 0; ; offset += SCAN_CHUNK_BYTES) {
+					if (context.abortSignal?.aborted) throw aborted(path);
 					const chunk = await range(offset, SCAN_CHUNK_BYTES);
 					if (chunk.length > 0) scanner.push(chunk);
 					if (chunk.length < SCAN_CHUNK_BYTES) return scanner.finish();
@@ -245,6 +256,10 @@ function fileReader(db: FileDatabase, path: string): BinaryReader {
 	};
 }
 
+function aborted(path: string): FileError {
+	return new FileError("aborted", "aborted", path);
+}
+
 function notFound(path: string): FileError {
 	return new FileError("not_found", `${path} does not exist`, path);
 }

packages/sandbox-adapter/src/pi/sandbox-env.ts

@@ -34,21 +34,33 @@ export function sandboxEnv(
 	connect: () => Promise<Sandbox>,
 	cwd: string | undefined,
 ): ExecutionEnv {
-	/** Connects on first use. A failed connect becomes this call's error result. */
+	/**
+	 * Connects on first use. A failed connect becomes this call's error result.
+	 * An aborted run stops before and after the connect. The connect itself
+	 * keeps going, because other calls may wait on it.
+	 */
 	const withSandbox = async <T, E extends Error>(
+		context: Context,
 		fail: (cause: Error) => E,
 		use: (sandbox: Sandbox) => Promise<Result<T, E>>,
 	): Promise<Result<T, E>> => {
+		if (context.abortSignal?.aborted) return err(fail(new Error("aborted")));
 		let sandbox: Sandbox;
 		try {
 			sandbox = await connect();
 		} catch (error) {
 			return err(fail(toError(error)));
 		}
+		if (context.abortSignal?.aborted) return err(fail(new Error("aborted")));
 		return use(sandbox);
 	};
-	const fileFailure = (path: string) => (cause: Error) =>
-		new FileError("unknown", cause.message, path, cause);
+	const fileFailure = (path: string, context: Context) => (cause: Error) =>
+		new FileError(
+			context.abortSignal?.aborted ? "aborted" : "unknown",
+			cause.message,
+			path,
+			cause,
+		);
 	const inside = (path: string): Result<string, FileError> => {
 		try {
 			return ok(resolveSandboxPath(root, path));
@@ -68,34 +80,34 @@ export function sandboxEnv(
 
 		absolutePath: async (path) => inside(posix.resolve(env.cwd, path)),
 		joinPath: async (parts) => ok(posix.join(...parts)),
-		exists: async (path) => {
+		exists: async (path, context) => {
 			const resolved = inside(path);
 			if (!resolved.ok) return resolved;
-			return withSandbox(fileFailure(resolved.value), (sandbox) =>
+			return withSandbox(context, fileFailure(resolved.value, context), (sandbox) =>
 				onPath(resolved, (target) => sandbox.exists(target)),
 			);
 		},
-		readTextFile: async (path) => {
+		readTextFile: async (path, context) => {
 			const resolved = inside(path);
 			if (!resolved.ok) return resolved;
-			return withSandbox(fileFailure(resolved.value), (sandbox) =>
+			return withSandbox(context, fileFailure(resolved.value, context), (sandbox) =>
 				read(sandbox, resolved.value, async () =>
 					new TextDecoder().decode(await sandbox.readFile(resolved.value)),
 				),
 			);
 		},
-		readBinaryFile: async (path) => {
+		readBinaryFile: async (path, context) => {
 			const resolved = inside(path);
 			if (!resolved.ok) return resolved;
-			return withSandbox(fileFailure(resolved.value), (sandbox) =>
+			return withSandbox(context, fileFailure(resolved.value, context), (sandbox) =>
 				read(sandbox, resolved.value, () => sandbox.readFile(resolved.value)),
 			);
 		},
 		// The sandbox reads whole files, so the reader holds the file it opened.
-		openBinaryReader: async (path) => {
+		openBinaryReader: async (path, _options, context) => {
 			const resolved = inside(path);
 			if (!resolved.ok) return resolved;
-			return withSandbox(fileFailure(resolved.value), async (sandbox) => {
+			return withSandbox(context, fileFailure(resolved.value, context), async (sandbox) => {
 				const info = await fileInfo(sandbox, resolved.value);
 				if (!info.ok) return info;
 				if (info.value.kind === "directory") {
@@ -119,23 +131,23 @@ export function sandboxEnv(
 				);
 			});
 		},
-		writeFile: async (path, content) => {
+		writeFile: async (path, content, context) => {
 			if (typeof content !== "string") {
 				return notSupported("writing binary content", path);
 			}
 			const resolved = inside(path);
 			if (!resolved.ok) return resolved;
-			return withSandbox(fileFailure(resolved.value), (sandbox) =>
+			return withSandbox(context, fileFailure(resolved.value, context), (sandbox) =>
 				onPath(resolved, async (target) => {
 					await sandbox.mkdir(posix.dirname(target));
 					await sandbox.writeFile(target, content);
 				}),
 			);
 		},
-		fileInfo: async (path) => {
+		fileInfo: async (path, context) => {
 			const resolved = inside(path);
 			if (!resolved.ok) return resolved;
-			return withSandbox(fileFailure(resolved.value), (sandbox) =>
+			return withSandbox(context, fileFailure(resolved.value, context), (sandbox) =>
 				fileInfo(sandbox, resolved.value),
 			);
 		},
@@ -150,7 +162,13 @@ export function sandboxEnv(
 				);
 			}
 			return withSandbox(
-				(cause) => new ExecutionError("spawn_error", cause.message, cause),
+				context,
+				(cause) =>
+					new ExecutionError(
+						context.abortSignal?.aborted ? "aborted" : "spawn_error",
+						cause.message,
+						cause,
+					),
 				(sandbox) => exec(sandbox, command, options, context, cwd.value),
 			);
 		},