feat: batch writes (#175)

Author: christianalfoni

package-lock.json

@@ -108,6 +108,7 @@
     "chalk": "^5.4.1",
     "cli-table3": "^0.6.3",
     "date-fns": "^4.1.0",
+    "fflate": "^0.8.2",
     "isbinaryfile": "^5.0.4",
     "ora": "^8.2.0",
     "path": "^0.12.7",

package.json

@@ -3,6 +3,7 @@ import { type IAgentClient } from "../AgentClient/agent-client-interface";
 import { Disposable } from "../utils/disposable";
 import { Emitter, type Event } from "../utils/event";
 import { Tracer, SpanStatusCode } from "@opentelemetry/api";
+import { zip } from "fflate";
 
 export type FSStatResult = {
   type: "file" | "directory";
@@ -39,6 +40,11 @@ export type Watcher = {
   onEvent: Event<WatchEvent>;
 };
 
+export type BatchWriteFile = {
+  path: string;
+  content: string | Uint8Array;
+};
+
 export class FileSystem {
   private disposable = new Disposable();
   private tracer?: Tracer;
@@ -137,6 +143,97 @@ export class FileSystem {
     );
   }
 
+  /**
+   * Batch write multiple files by zipping them, uploading the zip, and extracting it on the sandbox.
+   * This is more efficient than writing many files individually.
+   * Files will be created/overwritten as needed.
+   */
+  async batchWrite(files: BatchWriteFile[]): Promise<void> {
+    return this.withSpan(
+      "fs.batchWrite",
+      {
+        "fs.fileCount": files.length,
+      },
+      async () => {
+        if (files.length === 0) {
+          return;
+        }
+
+        // Create a zip file containing all the files
+        const zipData: Record<string, Uint8Array> = {};
+
+        for (const file of files) {
+          const content =
+            typeof file.content === "string"
+              ? new TextEncoder().encode(file.content)
+              : file.content;
+          zipData[file.path] = content;
+        }
+
+        // Create the zip using fflate
+        const zipBytes = await new Promise<Uint8Array>((resolve, reject) => {
+          zip(zipData, (err, data) => {
+            if (err) reject(err);
+            else resolve(data);
+          });
+        });
+
+        // Write the zip file to a temporary location
+        const tempZipPath = `/tmp/batch_write_${Date.now()}.zip`;
+        await this.writeFile(tempZipPath, zipBytes);
+
+        try {
+          // Extract the zip file using unzip command
+          const result = await this.agentClient.shells.create(
+            this.agentClient.workspacePath,
+            { cols: 128, rows: 24 },
+            `cd ${this.agentClient.workspacePath} && unzip -o ${tempZipPath}`,
+            "COMMAND",
+            true
+          );
+
+          if (result.status === "ERROR" || result.status === "KILLED") {
+            throw new Error(
+              `Failed to extract batch files: ${
+                result.buffer?.join("\n") || "Unknown error"
+              }`
+            );
+          }
+
+          // Wait for the command to complete if it's still running
+          if (result.status === "RUNNING") {
+            // Wait for shell exit event
+            await new Promise<void>((resolve, reject) => {
+              const disposable = this.agentClient.shells.onShellExited(
+                ({ shellId, exitCode }) => {
+                  if (shellId === result.shellId) {
+                    disposable.dispose();
+                    if (exitCode === 0) {
+                      resolve();
+                    } else {
+                      reject(
+                        new Error(
+                          `Unzip command failed with exit code ${exitCode}`
+                        )
+                      );
+                    }
+                  }
+                }
+              );
+            });
+          }
+        } finally {
+          // Always clean up the temporary zip file, regardless of success or failure
+          try {
+            await this.remove(tempZipPath);
+          } catch {
+            // Ignore cleanup errors - file might already be deleted or not exist
+          }
+        }
+      }
+    );
+  }
+
   /**
    * Create a directory.
    */