feat(core): refactor shell execution to use node-pty (#6088)

1 files changed, 132 insertions, 145 deletions

diff --git a/packages/core/src/services/shellExecutionService.ts b/packages/core/src/services/shellExecutionService.ts
index 3749fcf6..26d884b4 100644
--- a/packages/core/src/services/shellExecutionService.ts
+++ b/packages/core/src/services/shellExecutionService.ts
@@ -4,29 +4,35 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { spawn } from 'child_process';
+import * as pty from '@lydell/node-pty';
 import { TextDecoder } from 'util';
 import os from 'os';
-import stripAnsi from 'strip-ansi';
 import { getCachedEncodingForBuffer } from '../utils/systemEncoding.js';
 import { isBinary } from '../utils/textUtils.js';
+import pkg from '@xterm/headless';
+const { Terminal } = pkg;
 
-const SIGKILL_TIMEOUT_MS = 200;
+// @ts-expect-error getFullText is not a public API.
+const getFullText = (terminal: Terminal) => {
+  const buffer = terminal.buffer.active;
+  const lines: string[] = [];
+  for (let i = 0; i < buffer.length; i++) {
+    const line = buffer.getLine(i);
+    lines.push(line ? line.translateToString(true) : '');
+  }
+  return lines.join('\n').trim();
+};
 
 /** A structured result from a shell command execution. */
 export interface ShellExecutionResult {
   /** The raw, unprocessed output buffer. */
   rawOutput: Buffer;
-  /** The combined, decoded stdout and stderr as a string. */
+  /** The combined, decoded output as a string. */
   output: string;
-  /** The decoded stdout as a string. */
-  stdout: string;
-  /** The decoded stderr as a string. */
-  stderr: string;
   /** The process exit code, or null if terminated by a signal. */
   exitCode: number | null;
   /** The signal that terminated the process, if any. */
-  signal: NodeJS.Signals | null;
+  signal: number | null;
   /** An error object if the process failed to spawn. */
   error: Error | null;
   /** A boolean indicating if the command was aborted by the user. */
@@ -50,8 +56,6 @@ export type ShellOutputEvent =
   | {
       /** The event contains a chunk of output data. */
       type: 'data';
-      /** The stream from which the data originated. */
-      stream: 'stdout' | 'stderr';
       /** The decoded string chunk. */
       chunk: string;
     }
@@ -73,7 +77,7 @@ export type ShellOutputEvent =
  */
 export class ShellExecutionService {
   /**
-   * Executes a shell command using `spawn`, capturing all output and lifecycle events.
+   * Executes a shell command using `node-pty`, capturing all output and lifecycle events.
    *
    * @param commandToExecute The exact command string to run.
    * @param cwd The working directory to execute the command in.
@@ -87,167 +91,150 @@ export class ShellExecutionService {
     cwd: string,
     onOutputEvent: (event: ShellOutputEvent) => void,
     abortSignal: AbortSignal,
+    terminalColumns?: number,
+    terminalRows?: number,
   ): ShellExecutionHandle {
     const isWindows = os.platform() === 'win32';
+    const shell = isWindows ? 'cmd.exe' : 'bash';
+    const args = isWindows
+      ? ['/c', commandToExecute]
+      : ['-c', commandToExecute];
 
-    const child = spawn(commandToExecute, [], {
-      cwd,
-      stdio: ['ignore', 'pipe', 'pipe'],
-      // Use bash unless in Windows (since it doesn't support bash).
-      // For windows, just use the default.
-      shell: isWindows ? true : 'bash',
-      // Use process groups on non-Windows for robust killing.
-      // Windows process termination is handled by `taskkill /t`.
-      detached: !isWindows,
-      env: {
-        ...process.env,
-        GEMINI_CLI: '1',
-      },
-    });
+    let ptyProcess;
+    try {
+      ptyProcess = pty.spawn(shell, args, {
+        cwd,
+        name: 'xterm-color',
+        cols: terminalColumns ?? 200,
+        rows: terminalRows ?? 20,
+        env: {
+          ...process.env,
+          GEMINI_CLI: '1',
+        },
+        handleFlowControl: true,
+      });
+    } catch (e) {
+      const error = e as Error;
+      return {
+        pid: undefined,
+        result: Promise.resolve({
+          rawOutput: Buffer.from(''),
+          output: '',
+          exitCode: 1,
+          signal: null,
+          error,
+          aborted: false,
+          pid: undefined,
+        }),
+      };
+    }
 
     const result = new Promise<ShellExecutionResult>((resolve) => {
-      // Use decoders to handle multi-byte characters safely (for streaming output).
-      let stdoutDecoder: TextDecoder | null = null;
-      let stderrDecoder: TextDecoder | null = null;
-
-      let stdout = '';
-      let stderr = '';
+      const headlessTerminal = new Terminal({
+        allowProposedApi: true,
+        cols: terminalColumns ?? 200,
+        rows: terminalRows ?? 20,
+      });
+      let processingChain = Promise.resolve();
+      let decoder: TextDecoder | null = null;
+      let output = '';
       const outputChunks: Buffer[] = [];
-      let error: Error | null = null;
+      const error: Error | null = null;
       let exited = false;
 
       let isStreamingRawContent = true;
       const MAX_SNIFF_SIZE = 4096;
       let sniffedBytes = 0;
 
-      const handleOutput = (data: Buffer, stream: 'stdout' | 'stderr') => {
-        if (!stdoutDecoder || !stderrDecoder) {
-          const encoding = getCachedEncodingForBuffer(data);
-          try {
-            stdoutDecoder = new TextDecoder(encoding);
-            stderrDecoder = new TextDecoder(encoding);
-          } catch {
-            // If the encoding is not supported, fall back to utf-8.
-            // This can happen on some platforms for certain encodings like 'utf-32le'.
-            stdoutDecoder = new TextDecoder('utf-8');
-            stderrDecoder = new TextDecoder('utf-8');
-          }
-        }
-
-        outputChunks.push(data);
-
-        // Binary detection logic. This only runs until we've made a determination.
-        if (isStreamingRawContent && sniffedBytes < MAX_SNIFF_SIZE) {
-          const sniffBuffer = Buffer.concat(outputChunks.slice(0, 20));
-          sniffedBytes = sniffBuffer.length;
+      const handleOutput = (data: Buffer) => {
+        // NOTE: The migration from `child_process` to `node-pty` means we
+        // no longer have separate `stdout` and `stderr` streams. The `data`
+        // buffer contains the merged output. If a drop in LLM quality is
+        // observed after this change, we may need to revisit this and
+        // explore ways to re-introduce that distinction.
+        processingChain = processingChain.then(
+          () =>
+            new Promise<void>((resolve) => {
+              if (!decoder) {
+                const encoding = getCachedEncodingForBuffer(data);
+                try {
+                  decoder = new TextDecoder(encoding);
+                } catch {
+                  decoder = new TextDecoder('utf-8');
+                }
+              }
 
-          if (isBinary(sniffBuffer)) {
-            // Change state to stop streaming raw content.
-            isStreamingRawContent = false;
-            onOutputEvent({ type: 'binary_detected' });
-          }
-        }
+              outputChunks.push(data);
 
-        const decodedChunk =
-          stream === 'stdout'
-            ? stdoutDecoder.decode(data, { stream: true })
-            : stderrDecoder.decode(data, { stream: true });
-        const strippedChunk = stripAnsi(decodedChunk);
+              // First, check if we need to switch to binary mode.
+              if (isStreamingRawContent && sniffedBytes < MAX_SNIFF_SIZE) {
+                const sniffBuffer = Buffer.concat(outputChunks.slice(0, 20));
+                sniffedBytes = sniffBuffer.length;
 
-        if (stream === 'stdout') {
-          stdout += strippedChunk;
-        } else {
-          stderr += strippedChunk;
-        }
+                if (isBinary(sniffBuffer)) {
+                  isStreamingRawContent = false;
+                  onOutputEvent({ type: 'binary_detected' });
+                }
+              }
 
-        if (isStreamingRawContent) {
-          onOutputEvent({ type: 'data', stream, chunk: strippedChunk });
-        } else {
-          const totalBytes = outputChunks.reduce(
-            (sum, chunk) => sum + chunk.length,
-            0,
-          );
-          onOutputEvent({ type: 'binary_progress', bytesReceived: totalBytes });
-        }
+              // Now, based on the *current* state, either process as text or binary.
+              if (isStreamingRawContent) {
+                const decodedChunk = decoder.decode(data, { stream: true });
+                headlessTerminal.write(decodedChunk, () => {
+                  const newStrippedOutput = getFullText(headlessTerminal);
+                  output = newStrippedOutput;
+                  onOutputEvent({ type: 'data', chunk: newStrippedOutput });
+                  resolve();
+                });
+              } else {
+                // Once in binary mode, we only emit progress events.
+                const totalBytes = outputChunks.reduce(
+                  (sum, chunk) => sum + chunk.length,
+                  0,
+                );
+                onOutputEvent({
+                  type: 'binary_progress',
+                  bytesReceived: totalBytes,
+                });
+                resolve();
+              }
+            }),
+        );
       };
 
-      child.stdout.on('data', (data) => handleOutput(data, 'stdout'));
-      child.stderr.on('data', (data) => handleOutput(data, 'stderr'));
-      child.on('error', (err) => {
-        const { stdout, stderr, finalBuffer } = cleanup();
-        error = err;
-        resolve({
-          error,
-          stdout,
-          stderr,
-          rawOutput: finalBuffer,
-          output: stdout + (stderr ? `\n${stderr}` : ''),
-          exitCode: 1,
-          signal: null,
-          aborted: false,
-          pid: child.pid,
-        });
+      ptyProcess.onData((data) => {
+        const bufferData = Buffer.from(data, 'utf-8');
+        handleOutput(bufferData);
       });
 
-      const abortHandler = async () => {
-        if (child.pid && !exited) {
-          if (isWindows) {
-            spawn('taskkill', ['/pid', child.pid.toString(), '/f', '/t']);
-          } else {
-            try {
-              // Kill the entire process group (negative PID).
-              // SIGTERM first, then SIGKILL if it doesn't die.
-              process.kill(-child.pid, 'SIGTERM');
-              await new Promise((res) => setTimeout(res, SIGKILL_TIMEOUT_MS));
-              if (!exited) {
-                process.kill(-child.pid, 'SIGKILL');
-              }
-            } catch (_e) {
-              // Fall back to killing just the main process if group kill fails.
-              if (!exited) child.kill('SIGKILL');
-            }
-          }
-        }
-      };
-
-      abortSignal.addEventListener('abort', abortHandler, { once: true });
+      ptyProcess.onExit(({ exitCode, signal }) => {
+        exited = true;
+        abortSignal.removeEventListener('abort', abortHandler);
 
-      child.on('exit', (code: number, signal: NodeJS.Signals) => {
-        const { stdout, stderr, finalBuffer } = cleanup();
+        processingChain.then(() => {
+          const finalBuffer = Buffer.concat(outputChunks);
 
-        resolve({
-          rawOutput: finalBuffer,
-          output: stdout + (stderr ? `\n${stderr}` : ''),
-          stdout,
-          stderr,
-          exitCode: code,
-          signal,
-          error,
-          aborted: abortSignal.aborted,
-          pid: child.pid,
+          resolve({
+            rawOutput: finalBuffer,
+            output,
+            exitCode,
+            signal: signal ?? null,
+            error,
+            aborted: abortSignal.aborted,
+            pid: ptyProcess.pid,
+          });
         });
       });
 
-      /**
-       * Cleans up a process (and it's accompanying state) that is exiting or
-       * erroring and returns output formatted output buffers and strings
-       */
-      function cleanup() {
-        exited = true;
-        abortSignal.removeEventListener('abort', abortHandler);
-        if (stdoutDecoder) {
-          stdout += stripAnsi(stdoutDecoder.decode());
-        }
-        if (stderrDecoder) {
-          stderr += stripAnsi(stderrDecoder.decode());
+      const abortHandler = async () => {
+        if (ptyProcess.pid && !exited) {
+          ptyProcess.kill('SIGHUP');
         }
+      };
 
-        const finalBuffer = Buffer.concat(outputChunks);
-
-        return { stdout, stderr, finalBuffer };
-      }
+      abortSignal.addEventListener('abort', abortHandler, { once: true });
     });
 
-    return { pid: child.pid, result };
+    return { pid: ptyProcess.pid, result };
   }
 }