Process tweaks

src/computer/index.ts

@@ -15,6 +15,47 @@ import { DEFAULT_WORKDIR } from "./definitions";
 import { newProcess,
   Process } from "./process";
 
+const VALID_ENV_VAR_NAME = /^[A-Za-z_][A-Za-z0-9_]*$/;
+
+function buildShellCommand(
+  cmd: string,
+  cwd: string,
+  env: Record<string, string>,
+  emitPidMarker: boolean,
+): Result<string[]> {
+  for (const key of Object.keys(env)) {
+    if (!VALID_ENV_VAR_NAME.test(key)) {
+      return err(
+        "invalid_parameters_error",
+        `Invalid environment variable name: ${key}`,
+      );
+    }
+  }
+
+  const script = [
+    "set -e",
+    'while [ "$#" -gt 2 ]; do',
+    '  export "$1=$2"',
+    '  shift 2',
+    "done",
+    'cd -- "$1"',
+    emitPidMarker ? 'echo "SRCHD_PID:$$" >&2' : undefined,
+    'exec /bin/bash -lc "$2"',
+  ].filter((line): line is string => line !== undefined).join("\n");
+
+  const envArgs = Object.entries(env).flatMap(([key, value]) => [key, value]);
+
+  return ok([
+    "/bin/bash",
+    "-lc",
+    script,
+    "srchd",
+    ...envArgs,
+    cwd,
+    cmd,
+  ]);
+}
+
 export function computerId(
   experiment: ExperimentResource,
   agent: AgentResource,
@@ -217,41 +258,30 @@ export class Computer {
     const env = options?.env ?? {};
     const process = newProcess(cmd, cwd, env, options?.tty);
 
-    // Build the command with environment variables and working directory
-    // We wrap the command to capture the PID and run it in the foreground
-    // This keeps the K8s exec connection alive for the duration of the command
-    let fullCmd = "";
-    if (options?.env) {
-      const envVars = Object.entries(env)
-        .map(([k, v]) => `export ${k}="${v.replace(/"/g, '\\"')}"`)
-        .join("; ");
-      fullCmd += envVars + "; ";
+    const shellCommand = buildShellCommand(cmd, cwd, env, true);
+    if (shellCommand.isErr()) {
+      return shellCommand;
     }
 
-    // Run the command in foreground to keep the K8s exec connection alive.
-    // We use a wrapper script that:
-    // 1. Changes to the specified directory
-    // 2. Prints the PID to stderr for tracking
-    // 3. Runs the command directly (not in background)
-    // This ensures stdin/stdout/stderr remain connected for the duration of the process.
-    const escapedCmd = cmd.replace(/'/g, "'\\''");
-    const escapedCwd = cwd.replace(/'/g, "'\\''");
-    fullCmd += `cd '\''${escapedCwd}'\'' && echo "SRCHD_PID:$$" >&2 && ${escapedCmd}`;
-
     const res = await k8sSpawn(
-      ["/bin/bash", "-lc", fullCmd],
+      shellCommand.value,
       this.namespace,
       this.computerId,
       process,
       options?.tty ? undefined : options?.timeoutMs,
     );
 
-    // Extract PID from captured output
-    // The PID marker is written to stderr (or stdout in TTY mode)
-    const outputToSearch = options?.tty ? process.stdout : process.stderr;
-    const pidMatch = outputToSearch.match(/SRCHD_PID:(\d+)/);
-    if (pidMatch && pidMatch[1]) {
-      process.pid = parseInt(pidMatch[1], 10);
+    // Extract PID from captured output.
+    // We prefer the PID detected while streaming so it remains available even if
+    // the output buffer has already truncated older data.
+    if (process.detectedPid !== undefined) {
+      process.pid = process.detectedPid;
+    } else {
+      const outputToSearch = options?.tty ? process.stdout : process.stderr;
+      const pidMatch = outputToSearch.match(/SRCHD_PID:(\d+)/);
+      if (pidMatch?.[1]) {
+        process.pid = parseInt(pidMatch[1], 10);
+      }
     }
 
     this.processes.set(process.pid, process);
@@ -304,18 +334,13 @@ export class Computer {
 
     const startTs = Date.now();
 
-    // Build the command with environment variables and working directory
-    let fullCmd = "";
-    if (options?.env) {
-      const envVars = Object.entries(options.env)
-        .map(([k, v]) => `export ${k}="${v.replace(/"/g, '\\"')}"`)
-        .join("; ");
-      fullCmd += envVars + "; ";
+    const shellCommand = buildShellCommand(cmd, cwd, options?.env ?? {}, false);
+    if (shellCommand.isErr()) {
+      return shellCommand;
     }
-    fullCmd += `cd "${cwd.replace(/"/g, '\\"')}" && ${cmd}`;
 
     const execPromise = computerExec(
-      ["/bin/bash", "-lc", fullCmd],
+      shellCommand.value,
       this.namespace,
       this.computerId,
       options?.timeoutMs,

src/computer/process.ts

@@ -1,6 +1,8 @@
-import { PassThrough } from "stream";
+import { PassThrough, Writable } from "stream";
 import { Terminal } from "@xterm/xterm";
 
+const MAX_OUTPUT_CHARS = 1024 * 1024; // Keep only the last 1MB per stream.
+
 export type ProcessStatus =
   | 'running'
   | 'terminated';
@@ -9,11 +11,12 @@
   pid: number;
   status: ProcessStatus;
   stdinStream: PassThrough;
-  stdoutStream: PassThrough;
-  stderrStream: PassThrough;
+  stdoutStream: Writable;
+  stderrStream: Writable;
   output: { stdout: string; stderr: string };
   get stdout(): string;
   get stderr(): string;
+  detectedPid?: number;
   exitCode?: number;
   createdAt: Date;
   command: string;
@@ -25,17 +28,53 @@
   getTerminalBuffer(): string;
 }
 
+function appendToBoundedBuffer(
+  current: string,
+  incoming: string,
+  maxChars: number,
+): { value: string; truncatedChars: number } {
+  if (incoming.length >= maxChars) {
+    return {
+      value: incoming.slice(-maxChars),
+      truncatedChars: current.length + incoming.length - maxChars,
+    };
+  }
+
+  const overflow = current.length + incoming.length - maxChars;
+  if (overflow <= 0) {
+    return {
+      value: current + incoming,
+      truncatedChars: 0,
+    };
+  }
+
+  return {
+    value: current.slice(overflow) + incoming,
+    truncatedChars: overflow,
+  };
+}
+
+function formatBufferedOutput(value: string, truncatedChars: number): string {
+  if (truncatedChars <= 0) {
+    return value;
+  }
+
+  return `[srchd truncated ${truncatedChars} chars of earlier output; showing the most recent ${value.length} chars]\n${value}`;
+}
+
 export function newProcess(
   command: string,
   cwd: string,
   env: Record<string, string>,
   tty: boolean = false,
 ): Process {
 
-  const stdoutStream = new PassThrough();
-  const stderrStream = new PassThrough();
-  // Since JS strings are passed by value, we need to use getters to capture the output
-  const output = { stdout: "", stderr: "" };
+  const output = {
+    stdout: "",
+    stderr: "",
+    stdoutTruncatedChars: 0,
+    stderrTruncatedChars: 0,
+  };
 
   // Create terminal instance for TTY mode
   let terminal: Terminal | undefined;
@@ -50,45 +89,79 @@
     });
   }
 
-  const originalStdoutWrite = stdoutStream.write.bind(stdoutStream);
-  stdoutStream.write = (chunk: any, ...args: any[]) => {
-    if (chunk && chunk !== null) {
-      const text = chunk.toString();
-      output.stdout += text;
+  let process: Process | undefined;
 
-      // Write to terminal if in TTY mode
-      if (terminal) {
-        terminal.write(text);
-      }
+  const detectPid = (text: string) => {
+    const pidMatch = text.match(/SRCHD_PID:(\d+)/);
+    if (!pidMatch?.[1] || !process) {
+      return;
     }
-    return originalStdoutWrite(chunk, ...args);
+
+    const pid = parseInt(pidMatch[1], 10);
+    process.detectedPid = pid;
+    process.pid = pid;
   };
 
-  const originalStderrWrite = stderrStream.write.bind(stderrStream);
-  stderrStream.write = (chunk: any, ...args: any[]) => {
-    if (chunk && chunk !== null) {
-      const text = chunk.toString();
-      output.stderr += text;
+  const stdoutStream = new Writable({
+    write(chunk: any, _encoding, callback) {
+      try {
+        if (chunk !== undefined && chunk !== null) {
+          const text = chunk.toString();
+          const next = appendToBoundedBuffer(output.stdout, text, MAX_OUTPUT_CHARS);
+          output.stdout = next.value;
+          output.stdoutTruncatedChars += next.truncatedChars;
+          detectPid(text);
 
-      // In TTY mode, stderr also goes to terminal (like real TTY behavior)
-      if (terminal) {
-        terminal.write(text);
+          // Write to terminal if in TTY mode
+          if (terminal) {
+            terminal.write(text);
+          }
+        }
+        callback();
+      } catch (e) {
+        callback(e as Error);
       }
-    }
-    return originalStderrWrite(chunk, ...args);
-  };
+    },
+  });
+
+  const stderrStream = new Writable({
+    write(chunk: any, _encoding, callback) {
+      try {
+        if (chunk !== undefined && chunk !== null) {
+          const text = chunk.toString();
+          const next = appendToBoundedBuffer(output.stderr, text, MAX_OUTPUT_CHARS);
+          output.stderr = next.value;
+          output.stderrTruncatedChars += next.truncatedChars;
+          detectPid(text);
+
+          // In TTY mode, stderr also goes to terminal (like real TTY behavior)
+          if (terminal) {
+            terminal.write(text);
+          }
+        }
+        callback();
+      } catch (e) {
+        callback(e as Error);
+      }
+    },
+  });
 
   const stdinStream = new PassThrough();
 
-  return {
+  process = {
     pid: -1, // Nonexistent when starting process
     status: 'running',
     stdinStream,
     stdoutStream,
     stderrStream,
     output,
-    get stdout() { return output.stdout; },
-    get stderr() { return output.stderr; },
+    get stdout() {
+      return formatBufferedOutput(output.stdout, output.stdoutTruncatedChars);
+    },
+    get stderr() {
+      return formatBufferedOutput(output.stderr, output.stderrTruncatedChars);
+    },
+    detectedPid: undefined,
     exitCode: undefined,
     createdAt: new Date(),
     command,
@@ -98,7 +171,7 @@
     terminal,
     getTerminalBuffer(): string {
       if (!terminal) {
-        return output.stdout;
+        return formatBufferedOutput(output.stdout, output.stdoutTruncatedChars);
       }
 
       // Serialize the terminal buffer
@@ -115,4 +188,6 @@
       return lines.join('\n');
     },
   };
+
+  return process;
 }