1  import type { ChildProcess } from 'child_process'
  2  import { stat } from 'fs/promises'
  3  import type { Readable } from 'stream'
  4  import treeKill from 'tree-kill'
  5  import { generateTaskId } from '../Task.js'
  6  import { formatDuration } from './format.js'
  7  import {
  8    MAX_TASK_OUTPUT_BYTES,
  9    MAX_TASK_OUTPUT_BYTES_DISPLAY,
 10  } from './task/diskOutput.js'
 11  import { TaskOutput } from './task/TaskOutput.js'
 12  
 13  export type ExecResult = {
 14    stdout: string
 15    stderr: string
 16    code: number
 17    interrupted: boolean
 18    backgroundTaskId?: string
 19    backgroundedByUser?: boolean
 20    /** Set when assistant-mode auto-backgrounded a long-running blocking command. */
 21    assistantAutoBackgrounded?: boolean
 22    /** Set when stdout was too large to fit inline — points to the output file on disk. */
 23    outputFilePath?: string
 24    /** Total size of the output file in bytes (set when outputFilePath is set). */
 25    outputFileSize?: number
 26    /** The task ID for the output file (set when outputFilePath is set). */
 27    outputTaskId?: string
 28    /** Error message when the command failed before spawning (e.g., deleted cwd). */
 29    preSpawnError?: string
 30  }
 31  
 32  export type ShellCommand = {
 33    background: (backgroundTaskId: string) => boolean
 34    result: Promise<ExecResult>
 35    kill: () => void
 36    status: 'running' | 'backgrounded' | 'completed' | 'killed'
 37    /**
 38     * Cleans up stream resources (event listeners).
 39     * Should be called after the command completes or is killed to prevent memory leaks.
 40     */
 41    cleanup: () => void
 42    onTimeout?: (
 43      callback: (backgroundFn: (taskId: string) => boolean) => void,
 44    ) => void
 45    /** The TaskOutput instance that owns all stdout/stderr data and progress. */
 46    taskOutput: TaskOutput
 47  }
 48  
 49  const SIGKILL = 137
 50  const SIGTERM = 143
 51  
 52  // Background tasks write stdout/stderr directly to a file fd (no JS involvement),
 53  // so a stuck append loop can fill the disk. Poll file size and kill when exceeded.
 54  const SIZE_WATCHDOG_INTERVAL_MS = 5_000
 55  
 56  function prependStderr(prefix: string, stderr: string): string {
 57    return stderr ? `${prefix} ${stderr}` : prefix
 58  }
 59  
 60  /**
 61   * Thin pipe from a child process stream into TaskOutput.
 62   * Used in pipe mode (hooks) for stdout and stderr.
 63   * In file mode (bash commands), both fds go to the output file —
 64   * the child process streams are null and no wrappers are created.
 65   */
 66  class StreamWrapper {
 67    #stream: Readable | null
 68    #isCleanedUp = false
 69    #taskOutput: TaskOutput | null
 70    #isStderr: boolean
 71    #onData = this.#dataHandler.bind(this)
 72  
 73    constructor(stream: Readable, taskOutput: TaskOutput, isStderr: boolean) {
 74      this.#stream = stream
 75      this.#taskOutput = taskOutput
 76      this.#isStderr = isStderr
 77      // Emit strings instead of Buffers - avoids repeated .toString() calls
 78      stream.setEncoding('utf-8')
 79      stream.on('data', this.#onData)
 80    }
 81  
 82    #dataHandler(data: Buffer | string): void {
 83      const str = typeof data === 'string' ? data : data.toString()
 84  
 85      if (this.#isStderr) {
 86        this.#taskOutput!.writeStderr(str)
 87      } else {
 88        this.#taskOutput!.writeStdout(str)
 89      }
 90    }
 91  
 92    cleanup(): void {
 93      if (this.#isCleanedUp) {
 94        return
 95      }
 96      this.#isCleanedUp = true
 97      this.#stream!.removeListener('data', this.#onData)
 98      // Release references so the stream, its StringDecoder, and
 99      // the TaskOutput can be GC'd independently of this wrapper.
100      this.#stream = null
101      this.#taskOutput = null
102      this.#onData = () => {}
103    }
104  }
105  
106  /**
107   * Implementation of ShellCommand that wraps a child process.
108   *
109   * For bash commands: both stdout and stderr go to a file fd via
110   * stdio[1] and stdio[2] — no JS involvement. Progress is extracted
111   * by polling the file tail.
112   * For hooks: pipe mode with StreamWrappers for real-time detection.
113   */
114  class ShellCommandImpl implements ShellCommand {
115    #status: 'running' | 'backgrounded' | 'completed' | 'killed' = 'running'
116    #backgroundTaskId: string | undefined
117    #stdoutWrapper: StreamWrapper | null
118    #stderrWrapper: StreamWrapper | null
119    #childProcess: ChildProcess
120    #timeoutId: NodeJS.Timeout | null = null
121    #sizeWatchdog: NodeJS.Timeout | null = null
122    #killedForSize = false
123    #maxOutputBytes: number
124    #abortSignal: AbortSignal
125    #onTimeoutCallback:
126      | ((backgroundFn: (taskId: string) => boolean) => void)
127      | undefined
128    #timeout: number
129    #shouldAutoBackground: boolean
130    #resultResolver: ((result: ExecResult) => void) | null = null
131    #exitCodeResolver: ((code: number) => void) | null = null
132    #boundAbortHandler: (() => void) | null = null
133    readonly taskOutput: TaskOutput
134  
135    static #handleTimeout(self: ShellCommandImpl): void {
136      if (self.#shouldAutoBackground && self.#onTimeoutCallback) {
137        self.#onTimeoutCallback(self.background.bind(self))
138      } else {
139        self.#doKill(SIGTERM)
140      }
141    }
142  
143    readonly result: Promise<ExecResult>
144    readonly onTimeout?: (
145      callback: (backgroundFn: (taskId: string) => boolean) => void,
146    ) => void
147  
148    constructor(
149      childProcess: ChildProcess,
150      abortSignal: AbortSignal,
151      timeout: number,
152      taskOutput: TaskOutput,
153      shouldAutoBackground = false,
154      maxOutputBytes = MAX_TASK_OUTPUT_BYTES,
155    ) {
156      this.#childProcess = childProcess
157      this.#abortSignal = abortSignal
158      this.#timeout = timeout
159      this.#shouldAutoBackground = shouldAutoBackground
160      this.#maxOutputBytes = maxOutputBytes
161      this.taskOutput = taskOutput
162  
163      // In file mode (bash commands), both stdout and stderr go to the
164      // output file fd — childProcess.stdout/.stderr are both null.
165      // In pipe mode (hooks), wrap streams to funnel data into TaskOutput.
166      this.#stderrWrapper = childProcess.stderr
167        ? new StreamWrapper(childProcess.stderr, taskOutput, true)
168        : null
169      this.#stdoutWrapper = childProcess.stdout
170        ? new StreamWrapper(childProcess.stdout, taskOutput, false)
171        : null
172  
173      if (shouldAutoBackground) {
174        this.onTimeout = (callback): void => {
175          this.#onTimeoutCallback = callback
176        }
177      }
178  
179      this.result = this.#createResultPromise()
180    }
181  
182    get status(): 'running' | 'backgrounded' | 'completed' | 'killed' {
183      return this.#status
184    }
185  
186    #abortHandler(): void {
187      // On 'interrupt' (user submitted a new message), don't kill — let the
188      // caller background the process so the model can see partial output.
189      if (this.#abortSignal.reason === 'interrupt') {
190        return
191      }
192      this.kill()
193    }
194  
195    #exitHandler(code: number | null, signal: NodeJS.Signals | null): void {
196      const exitCode =
197        code !== null && code !== undefined
198          ? code
199          : signal === 'SIGTERM'
200            ? 144
201            : 1
202      this.#resolveExitCode(exitCode)
203    }
204  
205    #errorHandler(): void {
206      this.#resolveExitCode(1)
207    }
208  
209    #resolveExitCode(code: number): void {
210      if (this.#exitCodeResolver) {
211        this.#exitCodeResolver(code)
212        this.#exitCodeResolver = null
213      }
214    }
215  
216    // Note: exit/error listeners are NOT removed here — they're needed for
217    // the result promise to resolve. They clean up when the child process exits.
218    #cleanupListeners(): void {
219      this.#clearSizeWatchdog()
220      const timeoutId = this.#timeoutId
221      if (timeoutId) {
222        clearTimeout(timeoutId)
223        this.#timeoutId = null
224      }
225      const boundAbortHandler = this.#boundAbortHandler
226      if (boundAbortHandler) {
227        this.#abortSignal.removeEventListener('abort', boundAbortHandler)
228        this.#boundAbortHandler = null
229      }
230    }
231  
232    #clearSizeWatchdog(): void {
233      if (this.#sizeWatchdog) {
234        clearInterval(this.#sizeWatchdog)
235        this.#sizeWatchdog = null
236      }
237    }
238  
239    #startSizeWatchdog(): void {
240      this.#sizeWatchdog = setInterval(() => {
241        void stat(this.taskOutput.path).then(
242          s => {
243            // Bail if the watchdog was cleared while this stat was in flight
244            // (process exited on its own) — otherwise we'd mislabel stderr.
245            if (
246              s.size > this.#maxOutputBytes &&
247              this.#status === 'backgrounded' &&
248              this.#sizeWatchdog !== null
249            ) {
250              this.#killedForSize = true
251              this.#clearSizeWatchdog()
252              this.#doKill(SIGKILL)
253            }
254          },
255          () => {
256            // ENOENT before first write, or unlinked mid-run — skip this tick
257          },
258        )
259      }, SIZE_WATCHDOG_INTERVAL_MS)
260      this.#sizeWatchdog.unref()
261    }
262  
263    #createResultPromise(): Promise<ExecResult> {
264      this.#boundAbortHandler = this.#abortHandler.bind(this)
265      this.#abortSignal.addEventListener('abort', this.#boundAbortHandler, {
266        once: true,
267      })
268  
269      // Use 'exit' not 'close': 'close' waits for stdio to close, which includes
270      // grandchild processes that inherit file descriptors (e.g. `sleep 30 &`).
271      // 'exit' fires when the shell itself exits, returning control immediately.
272      this.#childProcess.once('exit', this.#exitHandler.bind(this))
273      this.#childProcess.once('error', this.#errorHandler.bind(this))
274  
275      this.#timeoutId = setTimeout(
276        ShellCommandImpl.#handleTimeout,
277        this.#timeout,
278        this,
279      ) as NodeJS.Timeout
280  
281      const exitPromise = new Promise<number>(resolve => {
282        this.#exitCodeResolver = resolve
283      })
284  
285      return new Promise<ExecResult>(resolve => {
286        this.#resultResolver = resolve
287        void exitPromise.then(this.#handleExit.bind(this))
288      })
289    }
290  
291    async #handleExit(code: number): Promise<void> {
292      this.#cleanupListeners()
293      if (this.#status === 'running' || this.#status === 'backgrounded') {
294        this.#status = 'completed'
295      }
296  
297      const stdout = await this.taskOutput.getStdout()
298      const result: ExecResult = {
299        code,
300        stdout,
301        stderr: this.taskOutput.getStderr(),
302        interrupted: code === SIGKILL,
303        backgroundTaskId: this.#backgroundTaskId,
304      }
305  
306      if (this.taskOutput.stdoutToFile && !this.#backgroundTaskId) {
307        if (this.taskOutput.outputFileRedundant) {
308          // Small file — full content is in result.stdout, delete the file
309          void this.taskOutput.deleteOutputFile()
310        } else {
311          // Large file — tell the caller where the full output lives
312          result.outputFilePath = this.taskOutput.path
313          result.outputFileSize = this.taskOutput.outputFileSize
314          result.outputTaskId = this.taskOutput.taskId
315        }
316      }
317  
318      if (this.#killedForSize) {
319        result.stderr = prependStderr(
320          `Background command killed: output file exceeded ${MAX_TASK_OUTPUT_BYTES_DISPLAY}`,
321          result.stderr,
322        )
323      } else if (code === SIGTERM) {
324        result.stderr = prependStderr(
325          `Command timed out after ${formatDuration(this.#timeout)}`,
326          result.stderr,
327        )
328      }
329  
330      const resultResolver = this.#resultResolver
331      if (resultResolver) {
332        this.#resultResolver = null
333        resultResolver(result)
334      }
335    }
336  
337    #doKill(code?: number): void {
338      this.#status = 'killed'
339      if (this.#childProcess.pid) {
340        treeKill(this.#childProcess.pid, 'SIGKILL')
341      }
342      this.#resolveExitCode(code ?? SIGKILL)
343    }
344  
345    kill(): void {
346      this.#doKill()
347    }
348  
349    background(taskId: string): boolean {
350      if (this.#status === 'running') {
351        this.#backgroundTaskId = taskId
352        this.#status = 'backgrounded'
353        this.#cleanupListeners()
354        if (this.taskOutput.stdoutToFile) {
355          // File mode: child writes directly to the fd with no JS involvement.
356          // The foreground timeout is gone, so watch file size to prevent
357          // a stuck append loop from filling the disk (768GB incident).
358          this.#startSizeWatchdog()
359        } else {
360          // Pipe mode: spill the in-memory buffer so readers can find it on disk.
361          this.taskOutput.spillToDisk()
362        }
363        return true
364      }
365      return false
366    }
367  
368    cleanup(): void {
369      this.#stdoutWrapper?.cleanup()
370      this.#stderrWrapper?.cleanup()
371      this.taskOutput.clear()
372      // Must run before nulling #abortSignal — #cleanupListeners() calls
373      // removeEventListener on it. Without this, a kill()+cleanup() sequence
374      // crashes: kill() queues #handleExit as a microtask, cleanup() nulls
375      // #abortSignal, then #handleExit runs #cleanupListeners() on the null ref.
376      this.#cleanupListeners()
377      // Release references to allow GC of ChildProcess internals and AbortController chain
378      this.#childProcess = null!
379      this.#abortSignal = null!
380      this.#onTimeoutCallback = undefined
381    }
382  }
383  
384  /**
385   * Wraps a child process to enable flexible handling of shell command execution.
386   */
387  export function wrapSpawn(
388    childProcess: ChildProcess,
389    abortSignal: AbortSignal,
390    timeout: number,
391    taskOutput: TaskOutput,
392    shouldAutoBackground = false,
393    maxOutputBytes = MAX_TASK_OUTPUT_BYTES,
394  ): ShellCommand {
395    return new ShellCommandImpl(
396      childProcess,
397      abortSignal,
398      timeout,
399      taskOutput,
400      shouldAutoBackground,
401      maxOutputBytes,
402    )
403  }
404  
405  /**
406   * Static ShellCommand implementation for commands that were aborted before execution.
407   */
408  class AbortedShellCommand implements ShellCommand {
409    readonly status = 'killed' as const
410    readonly result: Promise<ExecResult>
411    readonly taskOutput: TaskOutput
412  
413    constructor(opts?: {
414      backgroundTaskId?: string
415      stderr?: string
416      code?: number
417    }) {
418      this.taskOutput = new TaskOutput(generateTaskId('local_bash'), null)
419      this.result = Promise.resolve({
420        code: opts?.code ?? 145,
421        stdout: '',
422        stderr: opts?.stderr ?? 'Command aborted before execution',
423        interrupted: true,
424        backgroundTaskId: opts?.backgroundTaskId,
425      })
426    }
427  
428    background(): boolean {
429      return false
430    }
431  
432    kill(): void {}
433  
434    cleanup(): void {}
435  }
436  
437  export function createAbortedCommand(
438    backgroundTaskId?: string,
439    opts?: { stderr?: string; code?: number },
440  ): ShellCommand {
441    return new AbortedShellCommand({
442      backgroundTaskId,
443      ...opts,
444    })
445  }
446  
447  export function createFailedCommand(preSpawnError: string): ShellCommand {
448    const taskOutput = new TaskOutput(generateTaskId('local_bash'), null)
449    return {
450      status: 'completed' as const,
451      result: Promise.resolve({
452        code: 1,
453        stdout: '',
454        stderr: preSpawnError,
455        interrupted: false,
456        preSpawnError,
457      }),
458      taskOutput,
459      background(): boolean {
460        return false
461      },
462      kill(): void {},
463      cleanup(): void {},
464    }
465  }