| 1 | import { spawn, type ChildProcess } from "node:child_process"; |
| 2 | import { createWriteStream, existsSync, type WriteStream } from "node:fs"; |
| 3 | import { mkdir, open, stat } from "node:fs/promises"; |
| 4 | import path from "node:path"; |
| 5 | import type { Masker } from "./mask.js"; |
| 6 | import { killTree } from "./util.js"; |
| 7 | |
| 8 | /** |
| 9 | * Background processes in a sandbox: a coding agent's run can last hours, far |
| 10 | * longer than one /exec request. Each process writes stdout and stderr to its |
| 11 | * own files; callers read them from an offset while it runs, and after. |
| 12 | * |
| 13 | * Values passed in `secretEnv` (a person's model API key) reach only that |
| 14 | * process's environment, and are added to the sandbox's masker so no output |
| 15 | * the sandbox returns ever contains them. |
| 16 | */ |
| 17 | |
| 18 | export interface ProcStart { |
| 19 | id: string; |
| 20 | cmd: string; |
| 21 | cwd: string; |
| 22 | env: NodeJS.ProcessEnv; |
| 23 | secretEnv?: Record<string, string>; |
| 24 | timeoutMs?: number; |
| 25 | } |
| 26 | |
| 27 | export interface ProcRead { |
| 28 | id: string; |
| 29 | running: boolean; |
| 30 | exitCode: number | null; |
| 31 | timedOut: boolean; |
| 32 | stdout: string; |
| 33 | stderr: string; |
| 34 | /** Offsets to pass to the next read. */ |
| 35 | stdoutOffset: number; |
| 36 | stderrOffset: number; |
| 37 | } |
| 38 | |
| 39 | interface Proc { |
| 40 | child: ChildProcess; |
| 41 | out: string; |
| 42 | err: string; |
| 43 | files: WriteStream[]; |
| 44 | exitCode: number | null; |
| 45 | running: boolean; |
| 46 | timedOut: boolean; |
| 47 | timer?: NodeJS.Timeout; |
| 48 | done: Promise<void>; |
| 49 | } |
| 50 | |
| 51 | const ID = /^[A-Za-z0-9_-]{1,64}$/; |
| 52 | /** Each read returns at most this much of each stream. */ |
| 53 | export const MAX_READ = 512 * 1024; |
| 54 | |
| 55 | export class ProcTable { |
| 56 | private procs = new Map<string, Proc>(); |
| 57 | |
| 58 | constructor( |
| 59 | private dir: string, |
| 60 | private masker: Masker, |
| 61 | ) {} |
| 62 | |
| 63 | async start(o: ProcStart): Promise<{ id: string; pid: number | null }> { |
| 64 | if (!ID.test(o.id)) throw new Error("`id` must be 1-64 letters, digits, - or _"); |
| 65 | const prev = this.procs.get(o.id); |
| 66 | if (prev?.running) throw new Error(`process ${o.id} is already running`); |
| 67 | await mkdir(this.dir, { recursive: true }); |
| 68 | for (const v of Object.values(o.secretEnv ?? {})) this.masker.add(v); |
| 69 | const out = path.join(this.dir, `${o.id}.out`); |
| 70 | const err = path.join(this.dir, `${o.id}.err`); |
| 71 | const files = [createWriteStream(out, { flags: "w" }), createWriteStream(err, { flags: "w" })]; |
| 72 | const child = spawn("bash", ["-c", o.cmd], { cwd: o.cwd, env: { ...o.env, ...o.secretEnv }, detached: true, stdio: ["ignore", "pipe", "pipe"] }); |
| 73 | child.stdout!.pipe(files[0]!); |
| 74 | child.stderr!.pipe(files[1]!); |
| 75 | let finish!: () => void; |
| 76 | const p: Proc = { child, out, err, files, exitCode: null, running: true, timedOut: false, done: new Promise((r) => (finish = r)) }; |
| 77 | const end = (code: number) => { |
| 78 | if (!p.running) return; |
| 79 | p.running = false; |
| 80 | p.exitCode = p.timedOut ? 124 : code; |
| 81 | if (p.timer) clearTimeout(p.timer); |
| 82 | // Let the pipes drain into the files before readers see the exit. |
| 83 | let open = files.length; |
| 84 | for (const f of files) f.end(() => --open === 0 && finish()); |
| 85 | }; |
| 86 | child.on("error", (e) => { |
| 87 | files[1]!.write(`${e.message}\n`); |
| 88 | end(127); |
| 89 | }); |
| 90 | child.on("close", (code, signal) => end(code ?? (signal ? 128 + 15 : 1))); |
| 91 | if (o.timeoutMs && o.timeoutMs > 0) { |
| 92 | p.timer = setTimeout(() => { |
| 93 | p.timedOut = true; |
| 94 | killTree(child); |
| 95 | }, o.timeoutMs); |
| 96 | } |
| 97 | this.procs.set(o.id, p); |
| 98 | return { id: o.id, pid: child.pid ?? null }; |
| 99 | } |
| 100 | |
| 101 | async read(id: string, stdoutOffset = 0, stderrOffset = 0): Promise<ProcRead> { |
| 102 | const p = this.procs.get(id); |
| 103 | if (!p) throw new Error(`no process ${id}`); |
| 104 | // Once it has exited, wait for the last output to reach the files. |
| 105 | if (!p.running) await p.done; |
| 106 | const [so, se] = await Promise.all([this.chunk(p.out, stdoutOffset, p.running), this.chunk(p.err, stderrOffset, p.running)]); |
| 107 | return { id, running: p.running, exitCode: p.exitCode, timedOut: p.timedOut, stdout: this.masker.mask(so.text), stderr: this.masker.mask(se.text), stdoutOffset: so.next, stderrOffset: se.next }; |
| 108 | } |
| 109 | |
| 110 | kill(id: string, signal: "SIGTERM" | "SIGKILL" | "SIGINT" = "SIGTERM"): boolean { |
| 111 | const p = this.procs.get(id); |
| 112 | if (!p?.running) return false; |
| 113 | if (signal === "SIGTERM") killTree(p.child); |
| 114 | else { |
| 115 | try { |
| 116 | process.kill(-p.child.pid!, signal); |
| 117 | } catch { |
| 118 | p.child.kill(signal); |
| 119 | } |
| 120 | } |
| 121 | return true; |
| 122 | } |
| 123 | |
| 124 | killAll(): void { |
| 125 | for (const id of this.procs.keys()) this.kill(id); |
| 126 | } |
| 127 | |
| 128 | /** |
| 129 | * Bytes from `offset`, ending at a line break while the process runs so a |
| 130 | * masked value is never split between two reads (a single line longer than |
| 131 | * MAX_READ is the exception). |
| 132 | */ |
| 133 | private async chunk(file: string, offset: number, running: boolean): Promise<{ text: string; next: number }> { |
| 134 | if (!existsSync(file)) return { text: "", next: offset }; |
| 135 | const size = (await stat(file)).size; |
| 136 | if (offset >= size) return { text: "", next: offset }; |
| 137 | const len = Math.min(size - offset, MAX_READ); |
| 138 | const fh = await open(file, "r"); |
| 139 | try { |
| 140 | const buf = Buffer.alloc(len); |
| 141 | await fh.read(buf, 0, len, offset); |
| 142 | let end = len; |
| 143 | const partial = running || offset + len < size; |
| 144 | if (partial) { |
| 145 | const nl = buf.lastIndexOf(0x0a); |
| 146 | if (nl >= 0) end = nl + 1; |
| 147 | else if (len < MAX_READ) end = 0; |
| 148 | } |
| 149 | return { text: buf.subarray(0, end).toString("utf8"), next: offset + end }; |
| 150 | } finally { |
| 151 | await fh.close(); |
| 152 | } |
| 153 | } |
| 154 | } |