Refactor startup to be faster

Author: cramforce

src/sandbox/just-bash.test.ts

@@ -71,7 +71,7 @@ describe("wrapJustBash", () => {
     expect(content).toBe("content of /test.txt");
   });
 
-  it("wraps writeFile via fs.writeFile", async () => {
+  it("wraps writeFiles via fs.writeFile", async () => {
     const writtenFiles: Array<{ path: string; content: string }> = [];
     const mockBash = {
       exec: async () => ({ stdout: "", stderr: "", exitCode: 0 }),
@@ -84,10 +84,14 @@ describe("wrapJustBash", () => {
     };
 
     const sandbox = wrapJustBash(mockBash);
-    await sandbox.writeFile("/dir/test.txt", "my content");
+    await sandbox.writeFiles([
+      { path: "/dir/test.txt", content: "my content" },
+      { path: "/other/file.txt", content: "other content" },
+    ]);
 
     expect(writtenFiles).toEqual([
       { path: "/dir/test.txt", content: "my content" },
+      { path: "/other/file.txt", content: "other content" },
     ]);
   });
 });
@@ -121,7 +125,7 @@ describe("duck-typing disambiguation", () => {
     const customSandbox = {
       executeCommand: async () => ({ stdout: "", stderr: "", exitCode: 0 }),
       readFile: async () => "",
-      writeFile: async () => {},
+      writeFiles: async () => {},
     };
 
     expect(isJustBash(customSandbox)).toBe(false);

src/sandbox/just-bash.ts

@@ -63,8 +63,12 @@ export async function createJustBashSandbox(
       return bashEnv.fs.readFile(filePath);
     },
 
-    async writeFile(filePath: string, content: string): Promise<void> {
-      await bashEnv.fs.writeFile(filePath, content);
+    async writeFiles(
+      files: Array<{ path: string; content: string }>,
+    ): Promise<void> {
+      for (const file of files) {
+        await bashEnv.fs.writeFile(file.path, file.content);
+      }
     },
   };
 }
@@ -97,8 +101,12 @@ export function wrapJustBash(bashInstance: JustBashLike): Sandbox {
       return bashInstance.fs.readFile(filePath);
     },
 
-    async writeFile(filePath: string, content: string): Promise<void> {
-      await bashInstance.fs.writeFile(filePath, content);
+    async writeFiles(
+      files: Array<{ path: string; content: string }>,
+    ): Promise<void> {
+      for (const file of files) {
+        await bashInstance.fs.writeFile(file.path, file.content);
+      }
     },
   };
 }

src/sandbox/type-compatibility.test-d.ts

@@ -22,7 +22,7 @@ function acceptsVercelSandbox(sandbox: VercelSandbox) {
 function ourSandboxIsValid(sandbox: Sandbox) {
   sandbox.executeCommand("ls");
   sandbox.readFile("/file");
-  sandbox.writeFile("/file", "content");
+  sandbox.writeFiles([{ path: "/file", content: "content" }]);
 }
 
 // Suppress unused variable warnings

src/sandbox/vercel.test.ts

@@ -92,7 +92,7 @@ describe("wrapVercelSandbox", () => {
     );
   });
 
-  it("wraps writeFile using writeFiles with Buffer", async () => {
+  it("wraps writeFiles using native writeFiles with Buffer", async () => {
     const mockWriteFiles = vi.fn().mockResolvedValue(undefined);
 
     const mockVercelSandbox: VercelSandboxLike = {
@@ -103,10 +103,14 @@ describe("wrapVercelSandbox", () => {
     };
 
     const sandbox = wrapVercelSandbox(mockVercelSandbox);
-    await sandbox.writeFile("/test.txt", "content");
+    await sandbox.writeFiles([
+      { path: "/test.txt", content: "content" },
+      { path: "/other.txt", content: "other" },
+    ]);
 
     expect(mockWriteFiles).toHaveBeenCalledWith([
       { path: "/test.txt", content: Buffer.from("content", "utf-8") },
+      { path: "/other.txt", content: Buffer.from("other", "utf-8") },
     ]);
   });
 });

src/sandbox/vercel.ts

@@ -70,10 +70,15 @@ export function wrapVercelSandbox(vercelSandbox: VercelSandboxLike): Sandbox {
       return streamToString(stream);
     },
 
-    async writeFile(filePath: string, content: string): Promise<void> {
-      await vercelSandbox.writeFiles([
-        { path: filePath, content: Buffer.from(content, "utf-8") },
-      ]);
+    async writeFiles(
+      files: Array<{ path: string; content: string }>,
+    ): Promise<void> {
+      await vercelSandbox.writeFiles(
+        files.map((f) => ({
+          path: f.path,
+          content: Buffer.from(f.content, "utf-8"),
+        })),
+      );
     },
   };
 }

src/tool.test.ts

@@ -86,7 +86,7 @@ describe("createBashTool", () => {
     expect(sandbox).toBeDefined();
     expect(typeof sandbox.executeCommand).toBe("function");
     expect(typeof sandbox.readFile).toBe("function");
-    expect(typeof sandbox.writeFile).toBe("function");
+    expect(typeof sandbox.writeFiles).toBe("function");
   });
 
   it("writes inline files to destination", async () => {
@@ -295,7 +295,7 @@ describe("createBashTool", () => {
         .fn()
         .mockResolvedValue({ stdout: "custom", stderr: "", exitCode: 0 }),
       readFile: vi.fn().mockResolvedValue("custom content"),
-      writeFile: vi.fn().mockResolvedValue(undefined),
+      writeFiles: vi.fn().mockResolvedValue(undefined),
       stop: vi.fn().mockResolvedValue(undefined),
     };
 
@@ -306,11 +306,10 @@ describe("createBashTool", () => {
 
     expect(sandbox).toBe(customSandbox);
 
-    // Files should be written to custom sandbox
-    expect(customSandbox.writeFile).toHaveBeenCalledWith(
-      "/workspace/test.txt",
-      "content",
-    );
+    // Files should be written to custom sandbox in a single call
+    expect(customSandbox.writeFiles).toHaveBeenCalledWith([
+      { path: "/workspace/test.txt", content: "content" },
+    ]);
 
     // Tools should use custom sandbox
     assert(tools.bash.execute, "bash.execute should be defined");

src/tool.ts

@@ -79,9 +79,12 @@ export async function createBashTool(
       sandbox = options.sandbox as Sandbox;
     }
 
-    // Write files to existing sandbox
-    for (const [filePath, content] of Object.entries(filesWithDestination)) {
-      await sandbox.writeFile(filePath, content);
+    // Write files to existing sandbox in one call
+    const filesToWrite = Object.entries(filesWithDestination).map(
+      ([filePath, content]) => ({ path: filePath, content }),
+    );
+    if (filesToWrite.length > 0) {
+      await sandbox.writeFiles(filesToWrite);
     }
   } else {
     // Create just-bash sandbox with files

src/tools-prompt.test.ts

@@ -22,7 +22,7 @@ function createMockSandbox(tools: string[]): Sandbox {
       return { stdout: "", stderr: "Unknown command", exitCode: 1 };
     },
     readFile: async () => "",
-    writeFile: async () => {},
+    writeFiles: async () => {},
   };
 }
 
@@ -129,7 +129,7 @@ describe("discoverAvailableTools", () => {
         exitCode: 0,
       }),
       readFile: async () => "",
-      writeFile: async () => {},
+      writeFiles: async () => {},
     };
 
     const tools = await discoverAvailableTools(sandbox);

src/tools/bash.test.ts

@@ -13,7 +13,7 @@ vi.mock("ai", () => ({
 const mockSandbox = {
   executeCommand: vi.fn(),
   readFile: vi.fn(),
-  writeFile: vi.fn(),
+  writeFiles: vi.fn(),
   stop: vi.fn(),
 };
 

src/tools/write-file.ts

@@ -23,7 +23,7 @@ export function createWriteFileTool(options: CreateWriteFileToolOptions) {
     inputSchema: writeFileSchema,
     execute: async ({ path, content }) => {
       const resolvedPath = nodePath.posix.resolve(cwd, path);
-      await sandbox.writeFile(resolvedPath, content);
+      await sandbox.writeFiles([{ path: resolvedPath, content }]);
       return { success: true };
     },
   });