test(sandbox): cover relay lifecycle races

Author: VAIBHAVSING

packages/agentos-sandbox/src/provider.ts

@@ -16,7 +16,7 @@ export type SandboxAgentProviderOptions = Omit<
 interface SandboxAgentTransportInternals {
 	baseUrl: string;
 	token?: string;
-	defaultHeaders?: HeadersInit;
+	defaultHeaders?: RequestInit["headers"];
 	fetcher?: typeof globalThis.fetch;
 	awaitHealthy?(signal?: AbortSignal): Promise<void>;
 }
@@ -62,7 +62,9 @@ async function requestThroughSandboxAgent(
 		);
 	}
 	const headers = new Headers(transport.defaultHeaders);
-	new Headers(init.headers).forEach((value, name) => headers.set(name, value));
+	new Headers(init.headers).forEach((value, name) => {
+		headers.set(name, value);
+	});
 	if (transport.token) {
 		headers.set("authorization", `Bearer ${transport.token}`);
 	}

packages/agentos-sandbox/tests/vm-integration.test.ts

@@ -47,7 +47,7 @@ describe("VM integration", () => {
 			software: [common],
 			sandbox: {
 				mountPath: "/sandbox",
-				idleTimeoutMs: 25,
+				idleTimeoutMs: 500,
 				provider: {
 					start: async () => {
 						providerStarts += 1;
@@ -119,9 +119,7 @@ describe("VM integration", () => {
 				new TextEncoder().encode("first"),
 			);
 			expect(
-				new TextDecoder().decode(
-					await vm.readFile("/sandbox/generation.txt"),
-				),
+				new TextDecoder().decode(await vm.readFile("/sandbox/generation.txt")),
 			).toBe("first");
 			expect(providerStarts).toBe(1);
 
@@ -130,15 +128,32 @@ describe("VM integration", () => {
 				{ path: "/generation.txt" },
 				new TextEncoder().encode("second"),
 			);
-			for (let attempt = 0; attempt < 50 && providerDisposals === 0; attempt++) {
+			for (
+				let attempt = 0;
+				attempt < 100 && providerDisposals === 0;
+				attempt++
+			) {
 				await new Promise((resolve) => setTimeout(resolve, 10));
 			}
 			expect(providerDisposals).toBe(1);
-			expect(
-				new TextDecoder().decode(
-					await vm.readFile("/sandbox/generation.txt"),
+			const [secondContent, bindingResult] = await Promise.all([
+				vm.readFile("/sandbox/generation.txt"),
+				vm.process.execFile(
+					"agentos-sandbox",
+					["run-command", "--command", "echo", "--args", "second-generation"],
+					{ output: { capture: "all" } },
 				),
-			).toBe("second");
+			]);
+			expect(new TextDecoder().decode(secondContent)).toBe("second");
+			expect(bindingResult.exitCode).toBe(0);
+			expect(JSON.parse(bindingResult.stdout)).toEqual(
+				expect.objectContaining({
+					ok: true,
+					result: expect.objectContaining({
+						stdout: expect.stringContaining("second-generation"),
+					}),
+				}),
+			);
 			expect(providerStarts).toBe(2);
 		} finally {
 			await replacement.stop();

packages/core/src/index.ts

@@ -1,44 +1,43 @@
 // @rivet-dev/agentos
 
 export { AgentOs, AgentOsSidecar } from "./agent-os.js";
-export type * from "./language-execution.js";
+export {
+	isPackageDescriptor,
+	OPT_AGENTOS_BIN,
+	OPT_AGENTOS_ROOT,
+	tryReadAgentosPackageManifest,
+} from "./agentos-package.js";
+export type { Binding, BindingExample, Bindings } from "./bindings.js";
+export {
+	binding,
+	bindings,
+	MAX_BINDING_DESCRIPTION_LENGTH,
+	validateBindings,
+} from "./bindings.js";
 export {
 	CronManager,
 	InvalidScheduleError,
 	PastScheduleError,
 	TimerScheduleDriver,
 } from "./cron/index.js";
 export { createHostDirBackend, nodeModulesMount } from "./host-dir-mount.js";
-export {
-	binding,
-	MAX_BINDING_DESCRIPTION_LENGTH,
-	bindings,
-	validateBindings,
-} from "./bindings.js";
-export type { Binding, BindingExample, Bindings } from "./bindings.js";
+export type * from "./language-execution.js";
+export { createSnapshotExport } from "./layers.js";
 export {
 	agentOsLimitsSchema,
 	agentOsOptionFieldSchemas,
 	agentOsOptionsSchema,
 	bindingSchema,
+	bindingsSchema,
 	mountConfigSchema,
 	nativeMountConfigSchema,
 	parseAgentOsOptions,
 	permissionsSchema,
 	rootFilesystemConfigSchema,
 	sharedSidecarConfigSchema,
 	sidecarConfigSchema,
-	bindingsSchema,
 } from "./options-schema.js";
-export { createSnapshotExport } from "./layers.js";
 export { defineSoftware } from "./packages.js";
-export {
-	isPackageDescriptor,
-	OPT_AGENTOS_BIN,
-	OPT_AGENTOS_ROOT,
-	tryReadAgentosPackageManifest,
-} from "./agentos-package.js";
-export { KernelError } from "./runtime-compat.js";
 export type {
 	ExecOptions,
 	ExecResult,
@@ -48,6 +47,7 @@ export type {
 	VirtualDirEntry,
 	VirtualStat,
 } from "./runtime.js";
+export { KernelError } from "./runtime-compat.js";
 export {
 	createSandboxBindings,
 	createSandboxFs,

packages/core/src/sandbox-relay.ts

@@ -162,7 +162,9 @@ function mergeUpstreamHeaders(
 	const serializable = client as AgentOsSandboxClient &
 		SerializableSandboxClient;
 	const headers = new Headers(serializable.defaultHeaders);
-	requestHeaders.forEach((value, name) => headers.set(name, value));
+	requestHeaders.forEach((value, name) => {
+		headers.set(name, value);
+	});
 	if (serializable.token) {
 		headers.set("authorization", `Bearer ${serializable.token}`);
 	}
@@ -236,7 +238,10 @@ export async function createSandboxRelay(
 	const token = randomBytes(32).toString("base64url");
 	const maxConcurrentRequests =
 		options.maxConcurrentRequests ?? DEFAULT_MAX_RELAY_REQUESTS;
-	if (!Number.isSafeInteger(maxConcurrentRequests) || maxConcurrentRequests <= 0) {
+	if (
+		!Number.isSafeInteger(maxConcurrentRequests) ||
+		maxConcurrentRequests <= 0
+	) {
 		throw new Error("sandbox.maxRelayRequests must be a positive safe integer");
 	}
 
@@ -283,8 +288,7 @@ export async function createSandboxRelay(
 			activeRequests += 1;
 			if (
 				!warnedNearCapacity &&
-				activeRequests * 100 >=
-					maxConcurrentRequests * RELAY_WARNING_PERCENT
+				activeRequests * 100 >= maxConcurrentRequests * RELAY_WARNING_PERCENT
 			) {
 				warnedNearCapacity = true;
 				console.warn(
@@ -308,7 +312,7 @@ export async function createSandboxRelay(
 							signal: abortController.signal,
 							...(hasBody
 								? {
-										body: Readable.toWeb(request) as unknown as BodyInit,
+										body: Readable.toWeb(request) as never,
 										duplex: "half" as const,
 									}
 								: {}),

packages/core/src/sandbox.ts

@@ -203,9 +203,9 @@ class SandboxClientController implements SandboxRelayClientController {
 		if (!this.#provider) {
 			throw new Error("Sandbox client is not available");
 		}
-		if (this.#stopPromise) await this.#stopPromise;
+		if (this.#stopPromise !== undefined) await this.#stopPromise;
 		if (this.#current) return this.#current;
-		if (this.#startPromise) return await this.#startPromise;
+		if (this.#startPromise !== undefined) return await this.#startPromise;
 
 		const startPromise = this.#startProvider();
 		this.#startPromise = startPromise;
@@ -307,13 +307,15 @@ class SandboxClientController implements SandboxRelayClientController {
 			this.#disposed ||
 			this.#activeOperations !== 0 ||
 			!this.#current ||
-			this.#stopPromise
+			this.#stopPromise !== undefined
 		) {
 			return;
 		}
 		const client = this.#current;
 		this.#current = undefined;
-		const stopPromise = Promise.resolve(client.dispose?.()).then(() => undefined);
+		const stopPromise = Promise.resolve(client.dispose?.()).then(
+			() => undefined,
+		);
 		this.#stopPromise = stopPromise;
 		try {
 			await stopPromise;
@@ -416,9 +418,11 @@ function normalizeHeaders(
 	);
 }
 
-function getSerializableClientConfig(
-	client: AgentOsSandboxClient,
-): Pick<SandboxMountPluginConfig, "baseUrl" | "token" | "headers"> {
+function getSerializableClientConfig(client: AgentOsSandboxClient): {
+	baseUrl: string;
+	token?: string;
+	headers?: Record<string, string>;
+} {
 	const serializable = client as unknown as SerializableSandboxClient;
 	const baseUrl = serializable.baseUrl?.trim().replace(/\/+$/, "");
 	if (!baseUrl) {