Optimize JavaScript sandbox command startup

js/lib/resources/abstraction/sandbox.ts

@@ -61,6 +61,7 @@ function shellQuote(arg: string): string {
  */
 export class Sandbox extends Pod {
   public syncLocalDir: boolean = false;
+  private runtimePreparation?: Promise<boolean>;
 
   constructor(config: CreateStubConfig, syncLocalDir: boolean = false) {
     super(config);
@@ -192,12 +193,29 @@ export class Sandbox extends Pod {
 
     const ignorePatterns = this.syncLocalDir ? undefined : ["*"];
 
-    const prepared = await this.stub.prepareRuntime(
-      undefined,
-      EStubType.Sandbox,
-      true,
-      ignorePatterns,
-    );
+    if (!this.runtimePreparation) {
+      this.runtimePreparation = this.stub.prepareRuntime(
+        undefined,
+        EStubType.Sandbox,
+        true,
+        ignorePatterns,
+      );
+    }
+
+    const currentPreparation = this.runtimePreparation;
+    let prepared: boolean;
+    try {
+      prepared = await currentPreparation;
+    } catch (error) {
+      if (this.runtimePreparation === currentPreparation) {
+        this.runtimePreparation = undefined;
+      }
+      throw error;
+    }
+
+    if (!prepared && this.runtimePreparation === currentPreparation) {
+      this.runtimePreparation = undefined;
+    }
     if (!prepared) {
       const detail = this.stub.lastError?.message ?? "unknown reason";
       throw new SandboxConnectionError(`Failed to prepare runtime: ${detail}`);
@@ -453,7 +471,11 @@ export class SandboxInstance extends PodInstance {
     cwd?: string,
     env?: Record<string, string>,
   ): Promise<SandboxProcessResponse | SandboxProcess> {
-    const process = await this.exec(["python3", "-c", code], { cwd, env });
+    const process = await this.exec(["python3", "-c", code], {
+      cwd,
+      env,
+      wait: blocking,
+    });
     if (blocking) {
       await process.wait();
       const [stdoutStr, stderrStr] = await Promise.all([
@@ -492,7 +514,7 @@ export class SandboxInstance extends PodInstance {
 
   private async _exec(
     shellCommand: string,
-    opts?: { cwd?: string; env?: Record<string, string> },
+    opts?: ExecOptions,
   ): Promise<SandboxProcess> {
     const resp = await beamClient.request({
       method: "POST",
@@ -501,14 +523,15 @@ export class SandboxInstance extends PodInstance {
         command: shellCommand,
         cwd: opts?.cwd,
         env: opts?.env,
+        wait: opts?.wait ?? false,
       },
     });
     const data = resp.data as PodSandboxExecResponse;
     if (!data.ok || !data.pid || data.pid <= 0) {
       throw new SandboxProcessError(data.errorMsg || "Failed to start process");
     }
 
-    const process = new SandboxProcess(this, data.pid);
+    const process = new SandboxProcess(this, data.pid, data);
     this.processes[data.pid] = process;
     return process;
   }
@@ -594,9 +617,15 @@ export class SandboxProcessStream {
   constructor(
     process: SandboxProcess,
     fetchFn: () => Promise<string> | string,
+    initialOutput?: string,
   ) {
     this.process = process;
     this.fetch_fn = fetchFn;
+    if (initialOutput !== undefined) {
+      this._buffer = initialOutput;
+      this._last_output = initialOutput;
+      this._closed = true;
+    }
   }
 
   public [Symbol.asyncIterator](): AsyncIterableIterator<string> {
@@ -682,6 +711,9 @@ export class SandboxProcessStream {
   public async read(): Promise<string> {
     let data = this._buffer;
     this._buffer = "";
+    if (this._closed) {
+      return data;
+    }
     while (true) {
       const chunk = await this._fetch_next_chunk();
       if (chunk) {
@@ -708,6 +740,9 @@ export class SandboxProcess {
   public env: string[] = [];
   public running: boolean = true;
   private _status: string = "";
+  private _inlineDone: boolean = false;
+  private _inlineStdout: string = "";
+  private _inlineStderr: string = "";
 
   constructor(
     sandboxInstance: SandboxInstance,
@@ -718,6 +753,9 @@ export class SandboxProcess {
       env?: string[];
       exitCode?: number;
       running?: boolean;
+      done?: boolean;
+      stdout?: string;
+      stderr?: string;
     },
   ) {
     this.sandbox_instance = sandboxInstance;
@@ -727,14 +765,28 @@ export class SandboxProcess {
     this.env = info?.env ?? [];
     this.exitCode = info?.exitCode ?? -1;
     this.running = info?.running ?? this.exitCode < 0;
+    if (info?.done) {
+      this._inlineDone = true;
+      this.exitCode = info.exitCode ?? 0;
+      this.running = false;
+      this._status = "done";
+      this._inlineStdout = info.stdout || "";
+      this._inlineStderr = info.stderr || "";
+    }
   }
 
   /** Wait for the process to complete and return the exit code. */
   public async wait(): Promise<number> {
+    if (this.exitCode >= 0) {
+      return this.exitCode;
+    }
+
     [this.exitCode, this._status] = await this.status();
     while (this.exitCode < 0) {
       [this.exitCode, this._status] = await this.status();
-      await new Promise((r) => setTimeout(r, 100));
+      if (this.exitCode < 0) {
+        await new Promise((r) => setTimeout(r, 100));
+      }
     }
     return this.exitCode;
   }
@@ -787,7 +839,7 @@ export class SandboxProcess {
       });
       const data = resp.data as { ok: boolean; stdout?: string };
       return data.stdout || "";
-    });
+    }, this._inlineDone ? this._inlineStdout : undefined);
   }
 
   /** Get a handle to a stream of the process's stderr. */
@@ -800,7 +852,7 @@ export class SandboxProcess {
       });
       const data = resp.data as { ok: boolean; stderr?: string };
       return data.stderr || "";
-    });
+    }, this._inlineDone ? this._inlineStderr : undefined);
   }
 
   /** Returns a combined stream of both stdout and stderr. */
@@ -818,7 +870,7 @@ export class SandboxProcess {
       async _process_stream(streamName: "stdout" | "stderr") {
         const isStdout = streamName === "stdout";
         const stream = isStdout ? this._stdout : this._stderr;
-        const chunk = await (stream as any)._fetch_next_chunk();
+        const chunk = await stream.read();
         if (chunk) {
           if (isStdout) this._stdoutBuffer += chunk;
           else this._stderrBuffer += chunk;
@@ -834,7 +886,8 @@ export class SandboxProcess {
             this._queue.push(line + "\n");
           }
         } else {
-          const [exitCode] = await self.status();
+          const exitCode =
+            self.exitCode >= 0 ? self.exitCode : (await self.status())[0];
           if (exitCode >= 0) {
             const buf = isStdout ? this._stdoutBuffer : this._stderrBuffer;
             if (buf) {

js/lib/types/pod.ts

@@ -50,12 +50,17 @@ export interface PodSandboxExecRequest {
   command: string;
   cwd?: string;
   env?: Record<string, string>;
+  wait?: boolean;
 }
 
 export interface PodSandboxExecResponse {
   ok: boolean;
   errorMsg?: string;
   pid: number;
+  done?: boolean;
+  exitCode?: number;
+  stdout?: string;
+  stderr?: string;
 }
 export interface PodSandboxStatusRequest {
   containerId: string;
@@ -326,6 +331,8 @@ export interface PodSandboxCreateImageFromFilesystemResponse {
 export interface ExecOptions {
   cwd?: string;
   env?: Record<string, string>;
+  /** Wait briefly for completion and return inline status/output when available. */
+  wait?: boolean;
 }
 
 // Store requests here?

js/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@beamcloud/beam-js",
-  "version": "1.0.11",
+  "version": "1.0.13",
   "description": "TypeScript and JavaScript SDK for Beam",
   "main": "dist/index.js",
   "module": "dist/index.mjs",

js/tests/sandbox-network.test.ts

@@ -165,6 +165,117 @@ describe("Sandbox network parity", () => {
       8080: "https://8080.example.com",
     });
   });
+
+  test("shares runtime preparation across concurrent sandbox creates", async () => {
+    const sandbox = new Sandbox({ name: "concurrent-sandbox" });
+    let releasePreparation!: (prepared: boolean) => void;
+    const preparation = new Promise<boolean>((resolve) => {
+      releasePreparation = resolve;
+    });
+    const prepareRuntimeMock = jest
+      .spyOn(sandbox.stub, "prepareRuntime")
+      .mockReturnValue(preparation);
+    let nextContainer = 0;
+    const requestMock = jest
+      .spyOn(beamClient, "request")
+      .mockImplementation(async (config) => {
+        if (config.url === "api/v1/gateway/pods") {
+          nextContainer += 1;
+          return {
+            data: {
+              ok: true,
+              containerId: `sandbox-${nextContainer}`,
+            },
+          };
+        }
+        if (config.url?.endsWith("/connect")) {
+          return { data: { ok: true } };
+        }
+        throw new Error(`Unexpected request: ${config.url}`);
+      });
+
+    const firstCreate = sandbox.create();
+    const secondCreate = sandbox.create();
+
+    expect(prepareRuntimeMock).toHaveBeenCalledTimes(1);
+    releasePreparation(true);
+
+    const instances = await Promise.all([firstCreate, secondCreate]);
+    expect(instances.map((instance) => instance.containerId)).toEqual([
+      "sandbox-1",
+      "sandbox-2",
+    ]);
+    expect(requestMock).toHaveBeenCalledTimes(4);
+  });
+
+  test("returns inline exec results without follow-up requests", async () => {
+    const requestMock = jest.spyOn(beamClient, "request").mockResolvedValue({
+      data: {
+        ok: true,
+        pid: 7,
+        done: true,
+        exitCode: 0,
+        stdout: "v20.0.0\n",
+        stderr: "",
+      },
+    });
+
+    const instance = new SandboxInstance(
+      {
+        containerId: "sandbox-123",
+        stubId: "stub-123",
+        url: "",
+        ok: true,
+        errorMsg: "",
+      },
+      new Sandbox({ name: "networked-sandbox" })
+    );
+
+    const process = await instance.exec(["node", "-v"], { wait: true });
+
+    await expect(process.wait()).resolves.toBe(0);
+    await expect(process.stdout.read()).resolves.toBe("v20.0.0\n");
+    await expect(process.stderr.read()).resolves.toBe("");
+    expect(requestMock).toHaveBeenCalledTimes(1);
+    expect(requestMock).toHaveBeenCalledWith(
+      expect.objectContaining({
+        data: expect.objectContaining({ wait: true }),
+      })
+    );
+  });
+
+  test("iterates inline combined logs without follow-up requests", async () => {
+    const requestMock = jest.spyOn(beamClient, "request").mockResolvedValue({
+      data: {
+        ok: true,
+        pid: 8,
+        done: true,
+        exitCode: 0,
+        stdout: "first\nsecond\n",
+        stderr: "warning\n",
+      },
+    });
+
+    const instance = new SandboxInstance(
+      {
+        containerId: "sandbox-123",
+        stubId: "stub-123",
+        url: "",
+        ok: true,
+        errorMsg: "",
+      },
+      new Sandbox({ name: "networked-sandbox" })
+    );
+
+    const process = await instance.exec(["node", "-v"], { wait: true });
+    const lines: string[] = [];
+    for await (const line of process.logs) {
+      lines.push(line);
+    }
+
+    expect(lines).toEqual(["first\n", "second\n", "warning\n"]);
+    expect(requestMock).toHaveBeenCalledTimes(1);
+  });
 });
 
 describe("prepareRuntime surfaces real errors via lastError", () => {