my harness for niri
1import { randomBytes } from "crypto"
2import fs from "fs/promises"
3import path from "path"
4import {
5 DEFAULT_COMMAND_TIMEOUT_MS,
6 DEFAULT_FILE_TIMEOUT_MS,
7 IMAGE_MAX_BYTES,
8 IMAGE_ROOT,
9 USE_DOCKER_SHELL,
10 normalizeTimeoutMs,
11 resolveMaxLines,
12} from "./config"
13import { currentWorkingDirectory, runOneOff, runRaw } from "./shell"
14import type { EditResult, ImageToolPayload, ModelImageInput } from "./types"
15
16function invokesSudo(command: string): boolean {
17 return /(^|[;&|({]\s*)sudo(\s|$)/.test(String(command ?? ""))
18}
19
20function shouldRedirectStdinToDevNullByDefault(command: string): boolean {
21 // This runner executes in a real PTY, but we still want to avoid accidentally
22 // hanging forever on commands that commonly read from stdin (REPLs, pagers,
23 // editors, etc.). For most non-interactive commands we keep stdin attached.
24 const trimmed = String(command ?? "").trim()
25 if (!trimmed) return true
26
27 // Only inspect the leading program token (optionally prefixed by sudo).
28 // This is intentionally heuristic; it avoids trying to fully parse shell.
29 const m = trimmed.match(/^(?:sudo\s+)?([^\s]+)/)
30 const prog = (m?.[1] ?? "").toLowerCase()
31 if (!prog) return true
32
33 // Common interactive tools.
34 if (
35 [
36 "bash",
37 "sh",
38 "zsh",
39 "fish",
40 "python",
41 "python3",
42 "node",
43 "ruby",
44 "irb",
45 "php",
46 "psql",
47 "mysql",
48 "sqlite3",
49 "less",
50 "more",
51 "man",
52 "top",
53 "htop",
54 "nano",
55 "vim",
56 "vi",
57 "emacs",
58 "ssh",
59 ].includes(prog)
60 ) {
61 return true
62 }
63
64 // `cat` with no args is a very common accidental hang.
65 if (prog === "cat" && !/\bcat\b\s+\S/.test(trimmed)) return true
66
67 return false
68}
69
70/**
71 * Run a shell command.
72 * Output is capped at `maxLines` lines (default 150, 40 for known-verbose commands).
73 * When truncated, a "[truncated]" header is prepended so niri knows lines were dropped.
74 * Pass `maxLines: 0` to disable truncation entirely.
75 *
76 * @param command - Shell command to execute.
77 * @param maxLines - Optional line cap override (`0` disables truncation).
78 * @param timeoutMs - Optional command timeout in milliseconds.
79 * @returns Cleaned command output, potentially truncated with a header.
80 */
81export async function runCommand(command: string, maxLines?: number, timeoutMs?: number): Promise<string> {
82 const cap = resolveMaxLines(command, maxLines)
83 const opTimeoutMs = normalizeTimeoutMs(timeoutMs, DEFAULT_COMMAND_TIMEOUT_MS)
84 const raw = invokesSudo(command)
85 ? await runOneOff(command, await currentWorkingDirectory(opTimeoutMs), { timeoutMs: opTimeoutMs })
86 : await runRaw(command, {
87 timeoutMs: opTimeoutMs,
88 redirectStdinToDevNull: shouldRedirectStdinToDevNullByDefault(command),
89 })
90
91 if (cap === 0) return raw
92
93 const lines = raw.split("\n")
94 if (lines.length <= cap) return raw
95
96 const kept = lines.slice(-cap)
97 return `[truncated — showing last ${cap} of ${lines.length} lines]\n${kept.join("\n")}`
98}
99
100function normalizeImagePath(filePath: string): string {
101 const raw = String(filePath ?? "").trim()
102 if (!raw) throw new Error("image path is required")
103 const normalized = path.posix.normalize(raw)
104 if (!(normalized === IMAGE_ROOT || normalized.startsWith(`${IMAGE_ROOT}/`))) {
105 throw new Error(`image path must be inside ${IMAGE_ROOT}`)
106 }
107 return normalized
108}
109
110async function resolveLocalPath(filePath: string, timeoutMs: number): Promise<string> {
111 const raw = String(filePath ?? "").trim()
112 if (!raw) throw new Error("path is required")
113 if (path.isAbsolute(raw)) return path.normalize(raw)
114 return path.resolve(await currentWorkingDirectory(timeoutMs), raw)
115}
116
117function imageMimeFromBytes(filePath: string, data: Buffer): string | null {
118 if (data.length >= 8 && data.subarray(0, 8).equals(Buffer.from([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]))) {
119 return "image/png"
120 }
121 if (data.length >= 3 && data[0] === 0xff && data[1] === 0xd8 && data[2] === 0xff) return "image/jpeg"
122 if (data.length >= 6 && (data.subarray(0, 6).toString("ascii") === "GIF87a" || data.subarray(0, 6).toString("ascii") === "GIF89a")) {
123 return "image/gif"
124 }
125 if (data.length >= 12 && data.subarray(0, 4).toString("ascii") === "RIFF" && data.subarray(8, 12).toString("ascii") === "WEBP") {
126 return "image/webp"
127 }
128 if (data.length >= 2 && data.subarray(0, 2).toString("ascii") === "BM") return "image/bmp"
129 if (
130 data.length >= 4 &&
131 (data.subarray(0, 4).equals(Buffer.from([0x49, 0x49, 0x2a, 0x00])) ||
132 data.subarray(0, 4).equals(Buffer.from([0x4d, 0x4d, 0x00, 0x2a])))
133 ) {
134 return "image/tiff"
135 }
136
137 const ext = path.extname(filePath).toLowerCase()
138 const byExt: Record<string, string> = {
139 ".jpg": "image/jpeg",
140 ".jpeg": "image/jpeg",
141 ".png": "image/png",
142 ".gif": "image/gif",
143 ".webp": "image/webp",
144 ".bmp": "image/bmp",
145 ".tif": "image/tiff",
146 ".tiff": "image/tiff",
147 }
148 return byExt[ext] ?? null
149}
150
151/**
152 * Read an image from the container and encode it as a data URL for multimodal input.
153 * Only paths inside the configured IMAGE_ROOT are allowed.
154 *
155 * @param filePath - Absolute image path inside the allowed image root.
156 * @param timeoutMs - Optional read timeout in milliseconds.
157 * @returns Parsed image payload suitable for model multimodal input.
158 * @throws If validation, decoding, or path checks fail.
159 */
160export async function readImageForModel(filePath: string, timeoutMs?: number): Promise<ModelImageInput> {
161 const safePath = normalizeImagePath(filePath)
162 const opTimeoutMs = normalizeTimeoutMs(timeoutMs, DEFAULT_FILE_TIMEOUT_MS)
163
164 if (!USE_DOCKER_SHELL) {
165 let st
166 try {
167 st = await fs.stat(safePath)
168 } catch (err) {
169 const message = err instanceof Error ? err.message : String(err)
170 throw new Error(`could not stat ${safePath}: ${message}`)
171 }
172 if (!st.isFile()) throw new Error(`not a regular file: ${safePath}`)
173 if (st.size <= 0) throw new Error(`file is empty: ${safePath}`)
174 if (st.size > IMAGE_MAX_BYTES) throw new Error(`file too large: ${st.size} bytes (max ${IMAGE_MAX_BYTES})`)
175
176 const data = await fs.readFile(safePath)
177 const mime = imageMimeFromBytes(safePath, data)
178 if (!mime) throw new Error(`unsupported image type: ${safePath}`)
179 return {
180 path: safePath,
181 mime,
182 bytes: data.length,
183 dataUrl: `data:${mime};base64,${data.toString("base64")}`,
184 }
185 }
186
187 const py = [
188 "import base64, imghdr, json, mimetypes, os, stat, sys, warnings",
189 "warnings.filterwarnings('ignore', category=DeprecationWarning)",
190 "path = sys.argv[1]",
191 "max_bytes = int(sys.argv[2])",
192 "def out(obj):",
193 " print(json.dumps(obj, ensure_ascii=False))",
194 "try:",
195 " st = os.stat(path)",
196 "except FileNotFoundError:",
197 " out({'ok': False, 'message': f'file not found: {path}'})",
198 " raise SystemExit(0)",
199 "except Exception as e:",
200 " out({'ok': False, 'message': f'could not stat {path}: {e}'})",
201 " raise SystemExit(0)",
202 "if not stat.S_ISREG(st.st_mode):",
203 " out({'ok': False, 'message': f'not a regular file: {path}'})",
204 " raise SystemExit(0)",
205 "if st.st_size <= 0:",
206 " out({'ok': False, 'message': f'file is empty: {path}'})",
207 " raise SystemExit(0)",
208 "if st.st_size > max_bytes:",
209 " out({'ok': False, 'message': f'file too large: {st.st_size} bytes (max {max_bytes})'})",
210 " raise SystemExit(0)",
211 "with open(path, 'rb') as f:",
212 " data = f.read()",
213 "kind = imghdr.what(None, data)",
214 "mime_map = {",
215 " 'jpeg': 'image/jpeg',",
216 " 'png': 'image/png',",
217 " 'gif': 'image/gif',",
218 " 'webp': 'image/webp',",
219 " 'bmp': 'image/bmp',",
220 " 'tiff': 'image/tiff',",
221 "}",
222 "mime = mime_map.get(kind)",
223 "if not mime:",
224 " guessed, _ = mimetypes.guess_type(path)",
225 " if guessed and guessed.startswith('image/'):",
226 " mime = guessed",
227 "if not mime:",
228 " out({'ok': False, 'message': f'unsupported image type: {path}'})",
229 " raise SystemExit(0)",
230 "encoded = base64.b64encode(data).decode('ascii')",
231 "out({",
232 " 'ok': True,",
233 " 'path': path,",
234 " 'mime': mime,",
235 " 'bytes': len(data),",
236 " 'data_url': 'data:' + mime + ';base64,' + encoded,",
237 "})",
238 ].join("\n")
239
240 const raw = await runRaw(pythonCommand(py, [safePath, String(IMAGE_MAX_BYTES)]), { timeoutMs: opTimeoutMs })
241
242 let parsed: ImageToolPayload
243
244 try {
245 parsed = JSON.parse(raw.trim()) as ImageToolPayload
246 } catch {
247 throw new Error(`image_tool failed for ${safePath}: ${raw.trim() || "unknown error"}`)
248 }
249
250 if (!parsed?.ok) {
251 throw new Error(parsed?.message ?? `image_tool failed for ${safePath}`)
252 }
253
254 const bytes = parsed.bytes
255 if (!parsed.path || !parsed.mime || !parsed.data_url || typeof bytes !== "number" || !Number.isFinite(bytes)) {
256 throw new Error(`image_tool returned invalid payload for ${safePath}`)
257 }
258
259 return {
260 path: parsed.path,
261 mime: parsed.mime,
262 bytes,
263 dataUrl: parsed.data_url,
264 }
265}
266
267/**
268 * Escape a string for safe interpolation into a POSIX shell command.
269 */
270function shellQuote(str: string): string {
271 return "'" + str.replace(/'/g, "'\\''") + "'"
272}
273
274function pythonCommand(source: string, args: string[] = []): string {
275 const token = `NIRI_PY_${randomBytes(8).toString("hex").toUpperCase()}`
276 const scriptPath = `/tmp/niri-${randomBytes(8).toString("hex")}.py`
277 return [
278 `cat > ${shellQuote(scriptPath)} << '${token}'`,
279 source,
280 token,
281 `python3 ${shellQuote(scriptPath)} ${args.map(shellQuote).join(" ")}`,
282 `rm -f ${shellQuote(scriptPath)}`,
283 ].join("\n")
284}
285
286function wrappedBase64(str: string): string {
287 return Buffer.from(str, "utf8").toString("base64").match(/.{1,76}/g)?.join("\n") ?? ""
288}
289
290/**
291 * Read a file from the container with optional line-range selection.
292 * Returns content with a metadata header showing the line range and total line count.
293 * Defaults to lines 1–100 when no range is given.
294 *
295 * @param filePath - Path to the file inside the container.
296 * @param startLine - Inclusive 1-based start line.
297 * @param endLine - Optional inclusive 1-based end line.
298 * @param timeoutMs - Optional read timeout in milliseconds.
299 * @returns Header plus selected file content.
300 */
301export async function readFile(filePath: string, startLine = 1, endLine?: number, timeoutMs?: number): Promise<string> {
302 const start = Math.max(1, Math.trunc(Number(startLine)) || 1)
303 const end = endLine !== undefined ? Math.max(start, Math.trunc(Number(endLine)) || start) : undefined
304 const opTimeoutMs = normalizeTimeoutMs(timeoutMs, DEFAULT_FILE_TIMEOUT_MS)
305
306 if (!Number.isFinite(start) || (end !== undefined && !Number.isFinite(end))) {
307 throw new Error(`readFile: invalid line range (${startLine}, ${endLine})`)
308 }
309
310 if (!USE_DOCKER_SHELL) {
311 const resolvedPath = await resolveLocalPath(filePath, opTimeoutMs)
312 const content = await fs.readFile(resolvedPath, "utf8")
313 const lines = content.split("\n")
314 if (lines.at(-1) === "") lines.pop()
315 const totalLines = lines.length
316 const effectiveEnd = end ?? Math.min(start + 99, totalLines > 0 ? totalLines : start + 99)
317 const selected = lines.slice(start - 1, effectiveEnd).join("\n")
318 const rangeStr = `lines ${start}–${effectiveEnd}`
319 const totalStr = totalLines > 0 ? ` of ${totalLines} total` : ""
320 return `[${filePath} ${rangeStr}${totalStr}]\n${selected}`
321 }
322
323 const quoted = shellQuote(filePath)
324
325 const countRaw = await runRaw(`wc -l < ${quoted} 2>/dev/null || echo 0`, { timeoutMs: opTimeoutMs })
326 const totalLines = parseInt(countRaw.trim(), 10) || 0
327
328 const effectiveEnd = end ?? Math.min(start + 99, totalLines > 0 ? totalLines : start + 99)
329
330 const content = await runRaw(`sed -n '${start},${effectiveEnd}p' ${quoted} 2>&1`, { timeoutMs: opTimeoutMs })
331
332 const rangeStr = `lines ${start}–${effectiveEnd}`
333 const totalStr = totalLines > 0 ? ` of ${totalLines} total` : ""
334 const header = `[${filePath} ${rangeStr}${totalStr}]\n`
335
336 return header + content
337}
338
339/**
340 * Edit a file by replacing an exact snippet of text.
341 * old_text must appear exactly once; errors clearly if 0 or >1 matches.
342 *
343 * Uses an in-container python transform so we do not stream huge file bodies
344 * through the PTY (which is fragile on very large files).
345 *
346 * @param filePath - File path to edit.
347 * @param oldText - Exact text to replace (must match exactly once).
348 * @param newText - Replacement text.
349 * @param timeoutMs - Optional operation timeout in milliseconds.
350 * @returns Structured edit result describing success/failure.
351 */
352export async function editFile(filePath: string, oldText: string, newText: string, timeoutMs?: number): Promise<EditResult> {
353 if (oldText.length === 0) {
354 return { ok: false, message: "old_text must not be empty" }
355 }
356
357 const opTimeoutMs = normalizeTimeoutMs(timeoutMs, DEFAULT_FILE_TIMEOUT_MS)
358
359 if (!USE_DOCKER_SHELL) {
360 const resolvedPath = await resolveLocalPath(filePath, opTimeoutMs)
361 let content: string
362 try {
363 content = await fs.readFile(resolvedPath, "utf8")
364 } catch (err) {
365 const message = err instanceof Error ? err.message : String(err)
366 return { ok: false, message: `could not read ${filePath}: ${message}` }
367 }
368
369 const count = content.split(oldText).length - 1
370 if (count === 0) return { ok: false, message: `old_text not found in ${filePath}` }
371 if (count > 1) return { ok: false, message: `old_text found ${count} times in ${filePath} — must be unique` }
372
373 try {
374 await fs.writeFile(resolvedPath, content.replace(oldText, newText), "utf8")
375 } catch (err) {
376 const message = err instanceof Error ? err.message : String(err)
377 return { ok: false, message: `could not write ${filePath}: ${message}` }
378 }
379
380 const linesDelta = newText.split("\n").length - oldText.split("\n").length
381 const sign = linesDelta >= 0 ? "+" : ""
382 return {
383 ok: true,
384 message: `edited ${filePath} (${sign}${linesDelta} lines)`,
385 }
386 }
387
388 const payload = wrappedBase64(JSON.stringify({ path: filePath, old_text: oldText, new_text: newText }))
389 const payloadToken = `NIRI_EDIT_PAYLOAD_${randomBytes(8).toString("hex").toUpperCase()}`
390 const payloadPath = `/tmp/niri-edit-${randomBytes(8).toString("hex")}.b64`
391
392 const py = [
393 "import base64, json, sys",
394 "def out(obj):",
395 " print(json.dumps(obj, ensure_ascii=False))",
396 "with open(sys.argv[1], 'r', encoding='ascii') as f:",
397 " payload = json.loads(base64.b64decode(f.read()).decode('utf-8'))",
398 "path = payload.get('path', '')",
399 "old = payload.get('old_text', '')",
400 "new = payload.get('new_text', '')",
401 "if old == '':",
402 " out({'ok': False, 'message': 'old_text must not be empty'})",
403 " raise SystemExit(0)",
404 "try:",
405 " with open(path, 'r', encoding='utf-8') as f:",
406 " content = f.read()",
407 "except FileNotFoundError:",
408 " out({'ok': False, 'message': f'could not read {path}: file not found'})",
409 " raise SystemExit(0)",
410 "except Exception as e:",
411 " out({'ok': False, 'message': f'could not read {path}: {e}'})",
412 " raise SystemExit(0)",
413 "count = content.count(old)",
414 "if count == 0:",
415 " out({'ok': False, 'message': f'old_text not found in {path}'})",
416 " raise SystemExit(0)",
417 "if count > 1:",
418 " out({'ok': False, 'message': f'old_text found {count} times in {path} — must be unique'})",
419 " raise SystemExit(0)",
420 "updated = content.replace(old, new, 1)",
421 "try:",
422 " with open(path, 'w', encoding='utf-8') as f:",
423 " f.write(updated)",
424 "except Exception as e:",
425 " out({'ok': False, 'message': f'could not write {path}: {e}'})",
426 " raise SystemExit(0)",
427 "out({'ok': True})",
428 ].join("\n")
429
430 const raw = await runRaw(
431 [
432 `cat > ${shellQuote(payloadPath)} << '${payloadToken}'`,
433 payload,
434 payloadToken,
435 pythonCommand(py, [payloadPath]),
436 `rm -f ${shellQuote(payloadPath)}`,
437 ].join("\n"),
438 { timeoutMs: opTimeoutMs },
439 )
440
441 let parsed: { ok?: boolean; message?: string } | null = null
442 try {
443 parsed = JSON.parse(raw.trim()) as { ok?: boolean; message?: string }
444 } catch {
445 return { ok: false, message: `edit failed for ${filePath}: ${raw.trim() || "unknown error"}` }
446 }
447
448 if (!parsed.ok) {
449 return { ok: false, message: parsed.message ?? `edit failed for ${filePath}` }
450 }
451
452 const linesDelta = newText.split("\n").length - oldText.split("\n").length
453 const sign = linesDelta >= 0 ? "+" : ""
454 return {
455 ok: true,
456 message: `edited ${filePath} (${sign}${linesDelta} lines)`,
457 }
458}