|
| 1 | +import { type ChildProcess, spawn } from "node:child_process"; |
| 2 | + |
| 3 | +const MAX_BUFFER_BYTES = 64 * 1024; // 64KB rolling buffer per shell |
| 4 | + |
| 5 | +export interface BackgroundShellRecord { |
| 6 | + id: string; |
| 7 | + command: string; |
| 8 | + cwd: string; |
| 9 | + startedAt: number; |
| 10 | + endedAt?: number; |
| 11 | + status: "running" | "exited" | "killed"; |
| 12 | + exitCode?: number; |
| 13 | + signal?: NodeJS.Signals; |
| 14 | + /** Captured stdout+stderr, head-truncated when past MAX_BUFFER_BYTES. */ |
| 15 | + output: string; |
| 16 | + /** Total bytes ever emitted (across head-truncated history) — lets callers detect truncation. */ |
| 17 | + bytesEmitted: number; |
| 18 | +} |
| 19 | + |
| 20 | +export type BackgroundShellListener = (shells: readonly BackgroundShellRecord[]) => void; |
| 21 | + |
| 22 | +/** |
| 23 | + * Tracks long-running shell processes that the agent spawned with |
| 24 | + * `shell({ background: true })`. The agent's tool turn returns |
| 25 | + * immediately with a `task_id`; the process keeps running in the |
| 26 | + * background, output streams into a rolling buffer here, and the |
| 27 | + * agent can poll via `shell_output` or terminate via `shell_kill`. |
| 28 | + * |
| 29 | + * Cleanup invariant: every spawned process is SIGTERM'd on app exit |
| 30 | + * via terminal-restore so we don't leak children when the CLI quits. |
| 31 | + */ |
| 32 | +export class BackgroundShellStore { |
| 33 | + private records = new Map<string, BackgroundShellRecord>(); |
| 34 | + private processes = new Map<string, ChildProcess>(); |
| 35 | + private readonly listeners = new Set<BackgroundShellListener>(); |
| 36 | + private nextId = 1; |
| 37 | + |
| 38 | + /** |
| 39 | + * Spawn a detached shell. Returns the new record immediately — |
| 40 | + * caller doesn't wait for the process to do anything. The shell |
| 41 | + * runs via `sh -c` so the command string can contain pipes / |
| 42 | + * redirects / etc., the same surface a normal `shell` tool call has. |
| 43 | + */ |
| 44 | + spawn(command: string, cwd: string): BackgroundShellRecord { |
| 45 | + const id = `bg-${this.nextId++}`; |
| 46 | + const record: BackgroundShellRecord = { |
| 47 | + id, |
| 48 | + command, |
| 49 | + cwd, |
| 50 | + startedAt: Date.now(), |
| 51 | + status: "running", |
| 52 | + output: "", |
| 53 | + bytesEmitted: 0, |
| 54 | + }; |
| 55 | + const child = spawn(command, { |
| 56 | + shell: true, |
| 57 | + cwd, |
| 58 | + env: process.env, |
| 59 | + detached: false, // we want to track and kill on exit, not orphan |
| 60 | + stdio: ["ignore", "pipe", "pipe"], |
| 61 | + }); |
| 62 | + |
| 63 | + const onChunk = (chunk: Buffer): void => { |
| 64 | + const text = chunk.toString("utf8"); |
| 65 | + record.bytesEmitted += chunk.byteLength; |
| 66 | + record.output += text; |
| 67 | + // Trim head when the rolling buffer exceeds the cap. Lose the |
| 68 | + // oldest output, keep the most recent — that's almost always |
| 69 | + // what's relevant (errors and recent state). |
| 70 | + if (record.output.length > MAX_BUFFER_BYTES) { |
| 71 | + record.output = record.output.slice(record.output.length - MAX_BUFFER_BYTES); |
| 72 | + } |
| 73 | + this.notify(); |
| 74 | + }; |
| 75 | + child.stdout?.on("data", onChunk); |
| 76 | + child.stderr?.on("data", onChunk); |
| 77 | + child.on("exit", (code, signal) => { |
| 78 | + record.endedAt = Date.now(); |
| 79 | + record.status = signal === "SIGTERM" || signal === "SIGKILL" ? "killed" : "exited"; |
| 80 | + record.exitCode = code ?? undefined; |
| 81 | + record.signal = signal ?? undefined; |
| 82 | + this.processes.delete(id); |
| 83 | + this.notify(); |
| 84 | + }); |
| 85 | + child.on("error", (err) => { |
| 86 | + record.output += `\n[spawn error: ${err.message}]\n`; |
| 87 | + record.endedAt = Date.now(); |
| 88 | + record.status = "exited"; |
| 89 | + record.exitCode = -1; |
| 90 | + this.processes.delete(id); |
| 91 | + this.notify(); |
| 92 | + }); |
| 93 | + |
| 94 | + this.records.set(id, record); |
| 95 | + this.processes.set(id, child); |
| 96 | + this.notify(); |
| 97 | + return { ...record }; |
| 98 | + } |
| 99 | + |
| 100 | + /** |
| 101 | + * Read the current state of a tracked shell. Returns a copy so |
| 102 | + * callers can't mutate the live record. Returns undefined for |
| 103 | + * unknown ids. |
| 104 | + */ |
| 105 | + get(id: string): BackgroundShellRecord | undefined { |
| 106 | + const record = this.records.get(id); |
| 107 | + return record ? { ...record } : undefined; |
| 108 | + } |
| 109 | + |
| 110 | + /** All known shells, oldest-first. Returns copies. */ |
| 111 | + list(): readonly BackgroundShellRecord[] { |
| 112 | + return Array.from(this.records.values()).map((r) => ({ ...r })); |
| 113 | + } |
| 114 | + |
| 115 | + /** |
| 116 | + * Terminate a running shell. SIGTERM first; if the process is still |
| 117 | + * around after `gracePeriodMs`, SIGKILL. Resolves once the exit |
| 118 | + * handler has fired. No-op if the id is unknown or already exited. |
| 119 | + */ |
| 120 | + async kill(id: string, gracePeriodMs = 2000): Promise<void> { |
| 121 | + const child = this.processes.get(id); |
| 122 | + if (!child) return; |
| 123 | + const record = this.records.get(id); |
| 124 | + if (!record || record.status !== "running") return; |
| 125 | + return new Promise<void>((resolve) => { |
| 126 | + const onExit = () => { |
| 127 | + clearTimeout(killTimer); |
| 128 | + resolve(); |
| 129 | + }; |
| 130 | + child.once("exit", onExit); |
| 131 | + try { |
| 132 | + child.kill("SIGTERM"); |
| 133 | + } catch { |
| 134 | + // Process already gone; the exit handler may have fired |
| 135 | + // or be about to. Resolve once child emits or after the |
| 136 | + // timer trips. |
| 137 | + } |
| 138 | + const killTimer = setTimeout(() => { |
| 139 | + try { |
| 140 | + child.kill("SIGKILL"); |
| 141 | + } catch { |
| 142 | + // Same — best effort. |
| 143 | + } |
| 144 | + }, gracePeriodMs); |
| 145 | + }); |
| 146 | + } |
| 147 | + |
| 148 | + /** |
| 149 | + * SIGTERM all running shells. Used by terminal-restore on app exit |
| 150 | + * so we don't leak children to the parent shell. Doesn't wait for |
| 151 | + * graceful exits — the parent process is already going down. |
| 152 | + */ |
| 153 | + killAllSync(): void { |
| 154 | + for (const [_id, child] of this.processes) { |
| 155 | + try { |
| 156 | + child.kill("SIGTERM"); |
| 157 | + } catch { |
| 158 | + // Best effort during shutdown. |
| 159 | + } |
| 160 | + } |
| 161 | + } |
| 162 | + |
| 163 | + /** Subscribe to mutation events. Returns an unsubscribe function. */ |
| 164 | + subscribe(listener: BackgroundShellListener): () => void { |
| 165 | + this.listeners.add(listener); |
| 166 | + listener(this.list()); |
| 167 | + return () => { |
| 168 | + this.listeners.delete(listener); |
| 169 | + }; |
| 170 | + } |
| 171 | + |
| 172 | + private notify(): void { |
| 173 | + const snapshot = this.list(); |
| 174 | + for (const fn of this.listeners) fn(snapshot); |
| 175 | + } |
| 176 | +} |
0 commit comments