@prismatic-io/lux 0.0.1

package/src/drivers/subprocess/index.ts

@@ -0,0 +1,490 @@
+import type { ChildProcessWithoutNullStreams } from "node:child_process";
+import { createInterface } from "node:readline";
+import { fileURLToPath } from "node:url";
+import { z } from "zod";
+import {
+  type AgentDriver,
+  type Answer,
+  type Artifact,
+  DoneSchema,
+  type DriverEventMap,
+  type DriverPlugin,
+  ErrorEventSchema,
+  type FormatAnswerFn,
+  formatWireAnswer,
+  InterruptSchema,
+  type ManagedProcess,
+  type ParsedDriverEvent,
+  type ParseEventFn,
+  ProgressSchema,
+  parseWireEvent,
+  type ReadyState,
+  ReadyStateSchema,
+  type StartContext,
+  spawnManagedProcess,
+  TypedEmitter,
+} from "../../core/index.js";
+import { walkArtifacts } from "../shared/artifacts.js";
+
+export const SubprocessDriverConfigSchema = z
+  .object({
+    command: z.string(),
+    args: z.array(z.string()).optional(),
+    cwd: z.string().optional(),
+    env: z.record(z.string(), z.string()).optional(),
+    /** Inherit the host environment. Experiment mode always forces this off. */
+    inheritEnv: z.boolean().optional(),
+    protocol: z.enum(["lux-jsonl", "custom"]).default("lux-jsonl"),
+    idleTimeoutMs: z.number().positive().default(120_000),
+    maxInterrupts: z.number().positive().default(20),
+    /**
+     * Maximum files inventoried from LUX_ARTIFACTS_DIR. Keep the conservative
+     * default unless the child deliberately emits bounded full snapshots.
+     */
+    maxArtifactFiles: z.number().int().positive().max(250_000).optional(),
+    /**
+     * Cap on a single stdout line's bytes — readline buffers a newline-less
+     * stream without bound, so a misbehaving agent could OOM the host.
+     */
+    maxLineBytes: z.number().positive().optional(),
+    /**
+     * For "custom" protocols: parse a stdout line into a Lux driver event,
+     * or null to ignore it. Required for "custom", rejected for "lux-jsonl".
+     */
+    parseEvent: z.custom<ParseEventFn>((v) => typeof v === "function").optional(),
+    /**
+     * For "custom" protocols: serialize a Lux Answer into a stdin line.
+     * Defaults to the lux-jsonl format.
+     */
+    formatAnswer: z.custom<FormatAnswerFn>((v) => typeof v === "function").optional(),
+    /**
+     * For "custom" protocols: serialize the initial prompt into a stdin line
+     * written once at startup, for agents that read the opening turn from
+     * stdin rather than the LUX_PROMPT env var or argv.
+     */
+    formatPrompt: z.custom<(text: string) => string>((v) => typeof v === "function").optional(),
+    /** Close stdin after the opening prompt for one-shot CLIs that wait for EOF. */
+    closeStdinAfterPrompt: z.boolean().optional(),
+  })
+  .strict()
+  .refine((cfg) => cfg.protocol !== "custom" || cfg.parseEvent !== undefined, {
+    path: ["parseEvent"],
+    message: 'protocol "custom" requires a parseEvent function',
+  })
+  .refine((cfg) => cfg.protocol !== "lux-jsonl" || cfg.parseEvent === undefined, {
+    path: ["parseEvent"],
+    message: 'parseEvent is only used with protocol "custom"',
+  });
+
+export type SubprocessDriverConfig = z.infer<typeof SubprocessDriverConfigSchema>;
+
+const KILL_GRACE_MS = 5_000;
+const DEFAULT_MAX_LINE_BYTES = 10_000_000;
+const MAX_STDERR_TAIL_CHARS = 8_192;
+const DEFAULT_MAX_ARTIFACT_FILES = 500;
+const ARTIFACT_SNAPSHOT_MODULE = fileURLToPath(
+  new URL("../../core/artifact-snapshot.js", import.meta.url),
+);
+// Grace for a still-open stdout to flush buffered lines after the child's
+// exit before the read end is force-destroyed.
+const EXIT_FLUSH_GRACE_MS = 1_000;
+
+const ParsedDriverEventSchema = z.discriminatedUnion("tag", [
+  z.object({ tag: z.literal("ready"), ready: ReadyStateSchema }),
+  z.object({ tag: z.literal("interrupt"), interrupt: InterruptSchema }),
+  z.object({ tag: z.literal("progress"), progress: ProgressSchema }),
+  z.object({ tag: z.literal("done"), done: DoneSchema }),
+  z.object({ tag: z.literal("error"), error: ErrorEventSchema }),
+]);
+
+class SubprocessDriverImpl implements AgentDriver {
+  private readonly emitter = new TypedEmitter<DriverEventMap>();
+  private child: ChildProcessWithoutNullStreams | null = null;
+  private managed: ManagedProcess | null = null;
+  private readyResolved: ((ready: ReadyState) => void) | null = null;
+  private readyRejected: ((err: Error) => void) | null = null;
+  private idleTimer: NodeJS.Timeout | null = null;
+  private exitFlushTimer: NodeJS.Timeout | null = null;
+  private interruptsSeen = 0;
+  private pendingAnswers = 0;
+  private exited = false;
+  private stdoutDone = false;
+  private pipeHeldPastExit = false;
+  private terminalSent = false;
+  private exitDescription = "";
+  private exitSettled: Promise<void> = Promise.resolve();
+  // Chains events so handlers are awaited in emission order (see AgentDriver);
+  // readline delivers buffered lines synchronously.
+  private dispatchChain: Promise<void> = Promise.resolve();
+  private readonly parse: ParseEventFn;
+  private readonly format: FormatAnswerFn;
+  private readonly maxLineBytes: number;
+  private readonly config: SubprocessDriverConfig;
+  private artifactsDir: string | null = null;
+
+  constructor(config: SubprocessDriverConfig) {
+    this.config = config;
+    this.parse = config.parseEvent ?? parseWireEvent;
+    this.format = config.formatAnswer ?? formatWireAnswer;
+    this.maxLineBytes = config.maxLineBytes ?? DEFAULT_MAX_LINE_BYTES;
+  }
+
+  on<K extends keyof DriverEventMap>(event: K, handler: (e: DriverEventMap[K]) => void): void {
+    this.emitter.on(event, handler);
+  }
+
+  off<K extends keyof DriverEventMap>(event: K, handler: (e: DriverEventMap[K]) => void): void {
+    this.emitter.off(event, handler);
+  }
+
+  async start(ctx: StartContext): Promise<ReadyState> {
+    this.artifactsDir = ctx.artifactsDir;
+    const managed = spawnManagedProcess(this.config.command, this.config.args ?? [], {
+      cwd: this.config.cwd ?? ctx.runDir,
+      env: {
+        ...(this.config.inheritEnv === false ? {} : process.env),
+        ...this.config.env,
+        ...(ctx.caseId ? { LUX_CASE_ID: ctx.caseId } : {}),
+        LUX_RUN_DIR: ctx.runDir,
+        LUX_ARTIFACTS_DIR: ctx.artifactsDir,
+        LUX_PROMPT: ctx.prompt,
+        LUX_ARTIFACT_SNAPSHOT_MODULE: ARTIFACT_SNAPSHOT_MODULE,
+      },
+      killGraceMs: KILL_GRACE_MS,
+    });
+    const { child } = managed;
+    this.managed = managed;
+    this.child = child;
+    this.exitSettled = managed.exited.then(() => {});
+    // The child error handler below turns spawn failures into driver events.
+    void managed.spawned.catch(() => {});
+    this.bumpIdleTimer();
+
+    // A child that exits before draining stdin turns the next write into an
+    // EPIPE 'error' event that would crash the whole process.
+    child.stdin.on("error", () => {});
+
+    // setEncoding keeps a multibyte character straddling a pipe-chunk
+    // boundary from being mangled.
+    child.stderr.setEncoding("utf8");
+    let stderrTail = "";
+    child.stderr.on("data", (chunk: string) => {
+      const lines = (stderrTail + chunk).split("\n");
+      // Cap the retained partial line so unterminated stderr can't grow
+      // without bound.
+      stderrTail = (lines.pop() ?? "").slice(0, MAX_STDERR_TAIL_CHARS);
+      for (const line of lines) process.stderr.write(`[subprocess] ${line}\n`);
+    });
+    child.stderr.on("close", () => {
+      if (stderrTail) process.stderr.write(`[subprocess] ${stderrTail}\n`);
+      stderrTail = "";
+    });
+
+    // Tracks the unterminated tail: an over-cap partial line must die before
+    // readline's unbounded buffering OOMs the host. Complete lines are
+    // capped in handleLine.
+    let lineBytes = 0;
+    child.stdout.on("data", (chunk: Buffer) => {
+      const lastNewline = chunk.lastIndexOf(0x0a);
+      lineBytes = lastNewline === -1 ? lineBytes + chunk.length : chunk.length - lastNewline - 1;
+      if (lineBytes > this.maxLineBytes) this.failLineCap();
+    });
+
+    const stdoutLines = createInterface({ input: child.stdout });
+    stdoutLines.on("line", (line) => this.handleLine(line));
+    // readline 'close' fires only after every buffered line delivers, so a
+    // done line racing the child's exit wins over the synthetic exit error.
+    stdoutLines.on("close", () => {
+      this.stdoutDone = true;
+      this.emitExitErrorIfSilent();
+    });
+
+    if (this.config.formatPrompt) {
+      child.stdin.write(this.config.formatPrompt(ctx.prompt));
+      if (this.config.closeStdinAfterPrompt) child.stdin.end();
+    }
+
+    // Removed once the child is gone, or a suite sharing one AbortSignal
+    // would retain every exited driver for its life.
+    const onAbort = (): void => this.terminate();
+
+    // A missing/unexecutable binary emits `error`, not a throw; unhandled,
+    // it crashes the whole `lux run` instead of failing one case.
+    child.on("error", (err) => {
+      ctx.abortSignal.removeEventListener("abort", onAbort);
+      this.exited = true;
+      this.clearIdleTimer();
+      this.fail({
+        exitReason: "error",
+        reason: `failed to spawn '${this.config.command}': ${err.message}`,
+      });
+    });
+
+    child.on("exit", (code, signal) => {
+      ctx.abortSignal.removeEventListener("abort", onAbort);
+      this.exited = true;
+      this.clearIdleTimer();
+      this.exitDescription = `code=${code}, signal=${signal}`;
+      if (this.readyRejected) {
+        this.fail({
+          exitReason: "error",
+          reason: `subprocess exited before ready (${this.exitDescription})`,
+        });
+        return;
+      }
+      this.emitExitErrorIfSilent();
+      // readline's close never fires while a grandchild holds our stdout,
+      // which would hang the run: let buffered lines flush, then close the
+      // interface ourselves (destroying the stream alone never ends readline).
+      if (!this.stdoutDone) {
+        this.exitFlushTimer = setTimeout(() => {
+          if (this.stdoutDone) return;
+          this.pipeHeldPastExit = true;
+          stdoutLines.close();
+          child.stdout.destroy();
+          child.stderr.destroy();
+        }, EXIT_FLUSH_GRACE_MS);
+      }
+    });
+
+    return new Promise<ReadyState>((resolveStart, reject) => {
+      this.readyResolved = resolveStart;
+      this.readyRejected = reject;
+      // An already-aborted signal never fires "abort"; terminate directly.
+      if (ctx.abortSignal.aborted) this.terminate();
+      else ctx.abortSignal.addEventListener("abort", onAbort, { once: true });
+    });
+  }
+
+  /**
+   * Route a fatal condition to whoever is waiting: reject a pending start(),
+   * otherwise emit a terminal error event — never both, never twice.
+   */
+  private fail(error: DriverEventMap["error"]): void {
+    if (this.terminalSent) return;
+    this.terminalSent = true;
+    this.clearIdleTimer();
+    if (this.readyRejected) {
+      this.readyRejected(new Error(error.reason));
+      this.readyRejected = null;
+      this.readyResolved = null;
+      return;
+    }
+    this.emitInOrder("error", error);
+  }
+
+  /**
+   * All emits ride one promise chain so every handler is awaited before the
+   * next event dispatches (the AgentDriver ordering contract). A throwing
+   * handler becomes the terminal error and kills the child.
+   */
+  private emitInOrder<K extends keyof DriverEventMap>(event: K, payload: DriverEventMap[K]): void {
+    this.dispatchChain = this.dispatchChain
+      .then(() => this.emitter.emit(event, payload))
+      .catch((err: unknown) => {
+        this.fail({
+          exitReason: "error",
+          reason: `${String(event)} handler failed: ${err instanceof Error ? err.message : String(err)}`,
+        });
+        this.terminate();
+      });
+  }
+
+  // A child that exits without emitting done/error would hang the run
+  // forever. Requires both 'exit' and stdout close so no trailing line is
+  // still in flight.
+  private emitExitErrorIfSilent(): void {
+    if (!this.exited || !this.stdoutDone) return;
+    this.fail({
+      exitReason: "error",
+      reason: `subprocess exited before done (${this.exitDescription})`,
+    });
+  }
+
+  /** Request graceful shutdown; the shared supervisor owns escalation. */
+  private terminate(): void {
+    if (this.exited) void this.managed?.forceTerminate();
+    else void this.managed?.terminate();
+  }
+
+  private failLineCap(): void {
+    this.fail({
+      exitReason: "error",
+      reason: `stdout line exceeded ${this.maxLineBytes} bytes`,
+    });
+    this.terminate();
+  }
+
+  private handleLine(line: string): void {
+    // Post-terminal chatter must not re-arm the idle timer on a dead run.
+    if (this.terminalSent) return;
+    if (!line.trim()) return;
+    if (Buffer.byteLength(line) > this.maxLineBytes) {
+      this.failLineCap();
+      return;
+    }
+    // parseEvent is user-supplied: a throw inside readline's synchronous
+    // 'line' handler would take down the whole lux process, not just this case.
+    let parsed: ReturnType<ParseEventFn>;
+    try {
+      parsed = this.parse(line);
+    } catch (err) {
+      this.fail({
+        exitReason: "error",
+        reason: `parseEvent threw: ${err instanceof Error ? err.message : String(err)}`,
+      });
+      this.terminate();
+      return;
+    }
+    if (!parsed) return;
+    // A single line may carry multiple events (e.g. text + tool_use blocks).
+    const events = (Array.isArray(parsed) ? parsed : [parsed]).flatMap((event) => {
+      const validated = ParsedDriverEventSchema.safeParse(event);
+      return validated.success ? [validated.data] : [];
+    });
+    if (events.length === 0) return;
+    // Only protocol activity demonstrates that the agent is alive. Logs and
+    // malformed JSON must not keep a wedged agent running indefinitely.
+    if (this.pendingAnswers === 0) this.bumpIdleTimer();
+    for (const event of events) {
+      this.dispatchEvent(event);
+    }
+  }
+
+  private dispatchEvent(parsed: ParsedDriverEvent): void {
+    // Everything past the first done/error is dropped, even from a
+    // malformed agent that keeps emitting.
+    if (this.terminalSent) return;
+    // A terminal event or interrupt before ready is a protocol violation:
+    // a terminal here would strand the pending start(), and a pre-ready
+    // interrupt would pause the idle clock's never-became-ready rescue.
+    // Progress stays allowed — boot-time telemetry is harmless.
+    if (this.readyRejected && parsed.tag === "error") {
+      this.fail(parsed.error);
+      return;
+    }
+    if (this.readyRejected && parsed.tag !== "ready" && parsed.tag !== "progress") {
+      this.fail({
+        exitReason: "error",
+        reason: `agent emitted ${parsed.tag} before ready`,
+      });
+      return;
+    }
+    switch (parsed.tag) {
+      case "ready":
+        if (this.readyResolved) {
+          this.readyResolved(parsed.ready);
+          this.readyResolved = null;
+          this.readyRejected = null;
+        }
+        return;
+      case "error":
+        this.terminalSent = true;
+        this.clearIdleTimer();
+        this.emitInOrder("error", parsed.error);
+        return;
+      case "done":
+        this.terminalSent = true;
+        this.clearIdleTimer();
+        this.emitInOrder("done", parsed.done);
+        return;
+      case "interrupt":
+        this.interruptsSeen++;
+        if (this.interruptsSeen > this.config.maxInterrupts) {
+          this.fail({
+            exitReason: "max-interactions",
+            reason: `exceeded maxInterrupts (${this.config.maxInterrupts})`,
+          });
+          return;
+        }
+        // The answerer's round-trip is the user's time, not the agent's:
+        // the idle clock stops here and restarts in respond(); the
+        // answerer's own wait budget still bounds the window.
+        this.pendingAnswers++;
+        this.clearIdleTimer();
+        this.emitInOrder("interrupt", parsed.interrupt);
+        return;
+      case "progress":
+        this.emitInOrder("progress", parsed.progress);
+        return;
+    }
+  }
+
+  async respond(answer: Answer): Promise<void> {
+    if (!this.child || this.exited) {
+      throw new Error("cannot respond: subprocess is not running");
+    }
+    this.pendingAnswers = Math.max(0, this.pendingAnswers - 1);
+    if (this.pendingAnswers === 0) this.bumpIdleTimer();
+    const line = this.format(answer);
+    await new Promise<void>((res, rej) => {
+      this.child?.stdin.write(line, (err) => (err ? rej(err) : res()));
+    });
+  }
+
+  async collect(): Promise<Artifact[]> {
+    await this.quiesce();
+    return this.artifactsDir
+      ? walkArtifacts(this.artifactsDir, this.config.maxArtifactFiles ?? DEFAULT_MAX_ARTIFACT_FILES)
+      : [];
+  }
+
+  async quiesce(): Promise<void> {
+    this.clearIdleTimer();
+    this.clearExitFlushTimer();
+    if (this.child && !this.exited) {
+      this.child.stdin.end();
+      this.terminate();
+    } else if (this.child && (this.pipeHeldPastExit || !this.stdoutDone)) {
+      // The child is gone but something it spawned still holds our stdout:
+      // reap group survivors and drop the pipe so it can't outlive the suite.
+      await this.managed?.forceTerminate();
+      this.child.stdout.destroy();
+      this.child.stderr.destroy();
+    }
+    await this.exitSettled;
+  }
+
+  async close(): Promise<void> {
+    await this.quiesce();
+    this.emitter.removeAllListeners();
+  }
+
+  private bumpIdleTimer(): void {
+    this.clearIdleTimer();
+    this.idleTimer = setTimeout(() => {
+      const reason = `no event for ${this.config.idleTimeoutMs}ms`;
+      // Idle before ready means start() is still pending; reject it (and kill
+      // the child) or the orchestrator waits forever.
+      const preReady = this.readyRejected !== null;
+      this.fail(
+        preReady
+          ? { exitReason: "error", reason: `subprocess never became ready: ${reason}` }
+          : { exitReason: "idle-timeout", reason },
+      );
+      // Kill immediately: the run is lost and the child would keep doing
+      // billable work until close().
+      this.terminate();
+    }, this.config.idleTimeoutMs);
+  }
+
+  private clearIdleTimer(): void {
+    if (this.idleTimer) {
+      clearTimeout(this.idleTimer);
+      this.idleTimer = null;
+    }
+  }
+
+  private clearExitFlushTimer(): void {
+    if (this.exitFlushTimer) {
+      clearTimeout(this.exitFlushTimer);
+      this.exitFlushTimer = null;
+    }
+  }
+}
+
+export const subprocessDriver: DriverPlugin<SubprocessDriverConfig> = {
+  name: "subprocess",
+  configSchema: SubprocessDriverConfigSchema,
+  create: (config) => new SubprocessDriverImpl(config),
+};