my harness for niri
1import * as pty from "node-pty"
2import { randomBytes } from "crypto"
3import { spawn } from "child_process"
4import {
5 CONTAINER_NAME,
6 CONTAINER_USER,
7 DEFAULT_COMMAND_TIMEOUT_MS,
8 USE_DOCKER_SHELL,
9 normalizeTimeoutMs,
10} from "./config.js"
11import type { RunRawOptions } from "./types.js"
12
13/**
14 * Strip ANSI/VT escape sequences and normalize line endings from PTY output.
15 * PTYs emit CRLF and control sequences that we don't want in command results.
16 */
17function cleanOutput(str: string): string {
18 return str
19 .replace(/\x1b\[[0-9;?]*[a-zA-Z]/g, "") // CSI sequences (colors, cursor, etc.)
20 .replace(/\x1b\][^\x07\x1b]*(?:\x07|\x1b\\)/g, "") // OSC sequences (title, etc.)
21 .replace(/\x1b[^[\]]/g, "") // other 2-char ESC sequences
22 .replace(/(?:\x03|\^C)\s*/g, "") // Ctrl+C echo after interrupt
23 .replace(/\r\n/g, "\n") // CRLF -> LF
24 .replace(/\r/g, "\n") // stray CR -> LF
25 .replace(/\x00/g, "") // null bytes
26}
27
28let bash: pty.IPty | null = null
29
30function spawnBash(): { proc: pty.IPty; backend: string } {
31 const env = process.env as Record<string, string>
32 const options = {
33 name: "xterm-256color",
34 cols: 220, // wide enough to avoid line-wrapping sentinels
35 rows: 50,
36 env,
37 }
38
39 if (USE_DOCKER_SHELL) {
40 // docker exec -it allocates a PTY inside the container so bash runs
41 // interactively with job control. Combined with node-pty on the host
42 // this gives us a proper interactive shell where Ctrl+C interrupts
43 // the running command rather than killing bash itself.
44 return {
45 proc: pty.spawn("docker", ["exec", "-it", "-u", CONTAINER_USER, CONTAINER_NAME, "bash"], options),
46 backend: `docker:${CONTAINER_NAME}`,
47 }
48 }
49
50 return {
51 proc: pty.spawn("bash", ["--noprofile", "--norc", "-i"], {
52 ...options,
53 cwd: process.cwd(),
54 }),
55 backend: "local",
56 }
57}
58
59/**
60 * Opens and initializes the persistent PTY bash session inside the configured container.
61 *
62 * This performs one-time terminal setup (disable echo, normalize prompt, source bashrc)
63 * so subsequent `runRaw` calls can rely on sentinel-based output capture.
64 *
65 * @returns A promise that resolves when the shell is ready for commands.
66 * @throws If the container shell cannot be started or initialized.
67 */
68export async function openBash(): Promise<void> {
69 if (bash) return
70
71 const { proc, backend } = spawnBash()
72
73 proc.onExit(({ exitCode }) => {
74 console.log(`[bash:${backend}] exited with code ${exitCode}`)
75 if (bash === proc) bash = null
76 })
77
78 bash = proc
79
80 // Liveness check: if bash exits within 500ms the container is probably down.
81 await new Promise<void>((resolve, reject) => {
82 const timer = setTimeout(resolve, 500)
83 const d = proc.onExit(() => {
84 clearTimeout(timer)
85 d.dispose()
86 const detail = USE_DOCKER_SHELL ? ` — is the '${CONTAINER_NAME}' container running?` : ""
87 reject(new Error(`bash exited immediately${detail}`))
88 })
89 })
90
91 // We CANNOT use sentinel-based detection yet because bash echo is still on:
92 // if sending `echo SENTINEL` bash will echo the line back before running it
93 // so the sentinel appears in the output immediately as a false positive.
94 //
95 // Strategy: disable echo via a PROMPT-BASED signal.
96 // 1. Send `stty -echo` without any unique token on that input line.
97 // 2. Send `export PS1='<token>' PS2=''` and wait for the prompt token.
98 // 3. Now echo is off meaning sentinel detection is safe for all subsequent calls.
99 // 4. Use runRaw to source .bashrc and clear the prompt.
100
101 const initToken = `NIRI_INIT_${randomBytes(4).toString("hex")}_`
102
103 await new Promise<void>((resolve, reject) => {
104 let buf = ""
105 let dataDisposable: { dispose(): void } | null = null
106 const timer = setTimeout(() => {
107 dataDisposable?.dispose()
108 reject(new Error("bash init timed out"))
109 }, 10_000)
110 dataDisposable = proc.onData((chunk: string) => {
111 buf += chunk
112 const clean = cleanOutput(buf)
113 if (clean.includes(initToken)) {
114 clearTimeout(timer)
115 dataDisposable?.dispose()
116 resolve()
117 }
118 })
119 // Keep the token off the stty line. Otherwise the terminal echo can make
120 // readiness detection fire before echo has actually been disabled.
121 proc.write("stty -echo\n")
122 setTimeout(() => {
123 proc.write(`export PS1='${initToken}' PS2=''\n`)
124 }, 25)
125 })
126
127 // Echo is now off. Clear the token prompt, source .bashrc, done.
128 await runRaw("export PS1='' PS2=''")
129 await runRaw("source ~/.bashrc 2>/dev/null || true; export PS1='' PS2=''")
130 console.log(`[bash:${backend}] session ready`)
131}
132
133/**
134 * Closes the active PTY shell session if one exists.
135 */
136export function closeBash(): void {
137 if (bash) {
138 bash.kill()
139 bash = null
140 }
141}
142
143export async function currentWorkingDirectory(timeoutMs?: number): Promise<string> {
144 return (await runRaw("pwd -P", { timeoutMs, redirectStdinToDevNull: true })).trim()
145}
146
147/**
148 * Runs a command in the persistent PTY session and returns cleaned output.
149 *
150 * This is the low-level primitive used by higher-level container tools.
151 *
152 * @param command - Raw shell command to execute.
153 * @param options - Timeout and stdin-redirection behavior.
154 * @returns Combined stdout/stderr output with PTY control noise removed.
155 * @throws If command execution times out.
156 */
157export async function runRaw(command: string, options: RunRawOptions = {}): Promise<string> {
158 const timeoutMs = normalizeTimeoutMs(options.timeoutMs, DEFAULT_COMMAND_TIMEOUT_MS)
159 // Default: keep stdin attached to the PTY for more natural command behavior.
160 // Higher-level helpers (e.g. runCommand) can opt into /dev/null for commands
161 // that are likely to block waiting for stdin.
162 const redirectStdinToDevNull = options.redirectStdinToDevNull ?? false
163
164 // Reconnect lazily if the session was lost.
165 if (!bash) {
166 console.log("[bash] no session — attempting to reconnect...")
167 await openBash()
168 }
169
170 // Capture a stable local reference — the module-level `bash` may be nulled
171 // by the exit handler while we're waiting for output.
172 const session = bash!
173
174 // Two sentinels: start + end. Any PTY output buffered from a previous
175 // command arrives before the start sentinel and is discarded, preventing it
176 // from being prepended to this command's output.
177 const startSentinel = `__NIRI_START_${randomBytes(4).toString("hex")}__`
178 const endSentinel = `__NIRI_DONE_${randomBytes(4).toString("hex")}__`
179 let raw = ""
180 let settled = false
181
182 return new Promise((resolve, reject) => {
183 const dataDisposable = session.onData((chunk: string) => {
184 raw += chunk
185 const cleaned = cleanOutput(raw)
186 if (cleaned.includes(startSentinel) && cleaned.includes(endSentinel)) {
187 if (settled) return
188 settled = true
189 clearTimeout(timer)
190 dataDisposable.dispose()
191 const start = cleaned.indexOf(startSentinel) + startSentinel.length
192 const end = cleaned.indexOf(endSentinel)
193 // Drop the single newline that echo adds after the start sentinel,
194 // then trim trailing whitespace from the command output.
195 resolve(cleaned.slice(start, end).replace(/^\n/, "").trimEnd())
196 }
197 })
198
199 // Wrap command in a { } group. When requested, redirect stdin from
200 // /dev/null so accidental interactive reads do not hang forever.
201 // The closing } is on its own line so heredocs inside the command still
202 // get their correct delimiter line.
203 const groupedCommand = redirectStdinToDevNull ? `{ ${command}\n} < /dev/null` : `{ ${command}\n}`
204
205 session.write(`echo ${startSentinel}\n${groupedCommand}\necho ${endSentinel}\n`)
206
207 const timer = setTimeout(() => {
208 if (settled) return
209 settled = true
210 dataDisposable.dispose()
211 session.write("\x03\n")
212 // After a timeout we no longer know whether bash consumed Ctrl+C,
213 // returned to a prompt, or still has a foreground process attached.
214 // Reuse would interleave the next command with a potentially poisoned
215 // PTY, so force a fresh docker exec session next time.
216 session.kill()
217 if (bash === session) bash = null
218 reject(new Error(`Command timed out after ${timeoutMs}ms: ${command}`))
219 }, timeoutMs)
220 })
221}
222
223/**
224 * Runs a command as a one-off child process instead of through the persistent
225 * PTY shell. This follows Codex's shell-tool execution model for commands
226 * that should not inherit interactive stdin, notably sudo.
227 */
228export async function runOneOff(command: string, cwd: string, options: RunRawOptions = {}): Promise<string> {
229 const timeoutMs = normalizeTimeoutMs(options.timeoutMs, DEFAULT_COMMAND_TIMEOUT_MS)
230 const args = USE_DOCKER_SHELL
231 ? ["exec", "-i", "-u", CONTAINER_USER, "-w", cwd, CONTAINER_NAME, "bash", "-c", command]
232 : ["-c", command]
233 const program = USE_DOCKER_SHELL ? "docker" : "bash"
234
235 return new Promise((resolve, reject) => {
236 let raw = ""
237 let settled = false
238 const child = spawn(program, args, {
239 cwd: USE_DOCKER_SHELL ? undefined : cwd,
240 env: process.env,
241 stdio: ["ignore", "pipe", "pipe"],
242 })
243
244 child.stdout.on("data", (chunk) => {
245 raw += chunk.toString("utf8")
246 })
247 child.stderr.on("data", (chunk) => {
248 raw += chunk.toString("utf8")
249 })
250 child.on("error", (err) => {
251 if (settled) return
252 settled = true
253 clearTimeout(timer)
254 reject(err)
255 })
256 child.on("close", () => {
257 if (settled) return
258 settled = true
259 clearTimeout(timer)
260 resolve(cleanOutput(raw).trimEnd())
261 })
262
263 const timer = setTimeout(() => {
264 if (settled) return
265 settled = true
266 child.kill("SIGKILL")
267 reject(new Error(`Command timed out after ${timeoutMs}ms: ${command}`))
268 }, timeoutMs)
269 })
270}