feat: Adds shell command context to gemini history (#1076)

1 files changed, 274 insertions, 194 deletions

diff --git a/packages/cli/src/ui/hooks/shellCommandProcessor.ts b/packages/cli/src/ui/hooks/shellCommandProcessor.ts
index f7502f3f..fbd62d0c 100644
--- a/packages/cli/src/ui/hooks/shellCommandProcessor.ts
+++ b/packages/cli/src/ui/hooks/shellCommandProcessor.ts
@@ -5,11 +5,13 @@
  */
 
 import { spawn } from 'child_process';
+import { StringDecoder } from 'string_decoder';
 import type { HistoryItemWithoutId } from '../types.js';
-import type { exec as ExecType } from 'child_process';
 import { useCallback } from 'react';
-import { Config } from '@gemini-cli/core';
+import { Config, GeminiClient } from '@gemini-cli/core';
 import { type PartListUnion } from '@google/genai';
+import { formatMemoryUsage } from '../utils/formatters.js';
+import { isBinary } from '../utils/textUtils.js';
 import { UseHistoryManagerReturn } from './useHistoryManager.js';
 import crypto from 'crypto';
 import path from 'path';
@@ -18,10 +20,189 @@ import fs from 'fs';
 import stripAnsi from 'strip-ansi';
 
 const OUTPUT_UPDATE_INTERVAL_MS = 1000;
+const MAX_OUTPUT_LENGTH = 10000;
 
 /**
- * Hook to process shell commands (e.g., !ls, $pwd).
- * Executes the command in the target directory and adds output/errors to history.
+ * A structured result from a shell command execution.
+ */
+interface ShellExecutionResult {
+  rawOutput: Buffer;
+  output: string;
+  exitCode: number | null;
+  signal: NodeJS.Signals | null;
+  error: Error | null;
+  aborted: boolean;
+}
+
+/**
+ * Executes a shell command using `spawn`, capturing all output and lifecycle events.
+ * This is the single, unified implementation for shell execution.
+ *
+ * @param commandToExecute The exact command string to run.
+ * @param cwd The working directory to execute the command in.
+ * @param abortSignal An AbortSignal to terminate the process.
+ * @param onOutputChunk A callback for streaming real-time output.
+ * @param onDebugMessage A callback for logging debug information.
+ * @returns A promise that resolves with the complete execution result.
+ */
+function executeShellCommand(
+  commandToExecute: string,
+  cwd: string,
+  abortSignal: AbortSignal,
+  onOutputChunk: (chunk: string) => void,
+  onDebugMessage: (message: string) => void,
+): Promise<ShellExecutionResult> {
+  return new Promise((resolve) => {
+    const isWindows = os.platform() === 'win32';
+    const shell = isWindows ? 'cmd.exe' : 'bash';
+    const shellArgs = isWindows
+      ? ['/c', commandToExecute]
+      : ['-c', commandToExecute];
+
+    const child = spawn(shell, shellArgs, {
+      cwd,
+      stdio: ['ignore', 'pipe', 'pipe'],
+      detached: !isWindows, // Use process groups on non-Windows for robust killing
+    });
+
+    // Use decoders to handle multi-byte characters safely (for streaming output).
+    const stdoutDecoder = new StringDecoder('utf8');
+    const stderrDecoder = new StringDecoder('utf8');
+
+    let stdout = '';
+    let stderr = '';
+    const outputChunks: Buffer[] = [];
+    let error: Error | null = null;
+    let exited = false;
+
+    let streamToUi = true;
+    const MAX_SNIFF_SIZE = 4096;
+    let sniffedBytes = 0;
+
+    const handleOutput = (data: Buffer, stream: 'stdout' | 'stderr') => {
+      outputChunks.push(data);
+
+      if (streamToUi && sniffedBytes < MAX_SNIFF_SIZE) {
+        // Use a limited-size buffer for the check to avoid performance issues.
+        const sniffBuffer = Buffer.concat(outputChunks.slice(0, 20));
+        sniffedBytes = sniffBuffer.length;
+
+        if (isBinary(sniffBuffer)) {
+          streamToUi = false;
+          // Overwrite any garbled text that may have streamed with a clear message.
+          onOutputChunk('[Binary output detected. Halting stream...]');
+        }
+      }
+
+      const decodedChunk =
+        stream === 'stdout'
+          ? stdoutDecoder.write(data)
+          : stderrDecoder.write(data);
+      if (stream === 'stdout') {
+        stdout += stripAnsi(decodedChunk);
+      } else {
+        stderr += stripAnsi(decodedChunk);
+      }
+
+      if (!exited && streamToUi) {
+        // Send only the new chunk to avoid re-rendering the whole output.
+        const combinedOutput = stdout + (stderr ? `\n${stderr}` : '');
+        onOutputChunk(combinedOutput);
+      } else if (!exited && !streamToUi) {
+        // Send progress updates for the binary stream
+        const totalBytes = outputChunks.reduce(
+          (sum, chunk) => sum + chunk.length,
+          0,
+        );
+        onOutputChunk(
+          `[Receiving binary output... ${formatMemoryUsage(totalBytes)} received]`,
+        );
+      }
+    };
+
+    child.stdout.on('data', (data) => handleOutput(data, 'stdout'));
+    child.stderr.on('data', (data) => handleOutput(data, 'stderr'));
+    child.on('error', (err) => {
+      error = err;
+    });
+
+    const abortHandler = async () => {
+      if (child.pid && !exited) {
+        onDebugMessage(`Aborting shell command (PID: ${child.pid})`);
+        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, 200));
+            if (!exited) {
+              process.kill(-child.pid, 'SIGKILL');
+            }
+          } catch (_e) {
+            // Fallback to killing just the main process if group kill fails.
+            if (!exited) child.kill('SIGKILL');
+          }
+        }
+      }
+    };
+
+    abortSignal.addEventListener('abort', abortHandler, { once: true });
+
+    child.on('exit', (code, signal) => {
+      exited = true;
+      abortSignal.removeEventListener('abort', abortHandler);
+
+      // Handle any final bytes lingering in the decoders
+      stdout += stdoutDecoder.end();
+      stderr += stderrDecoder.end();
+
+      const finalBuffer = Buffer.concat(outputChunks);
+
+      resolve({
+        rawOutput: finalBuffer,
+        output: stdout + (stderr ? `\n${stderr}` : ''),
+        exitCode: code,
+        signal,
+        error,
+        aborted: abortSignal.aborted,
+      });
+    });
+  });
+}
+
+function addShellCommandToGeminiHistory(
+  geminiClient: GeminiClient,
+  rawQuery: string,
+  resultText: string,
+) {
+  const modelContent =
+    resultText.length > MAX_OUTPUT_LENGTH
+      ? resultText.substring(0, MAX_OUTPUT_LENGTH) + '\n... (truncated)'
+      : resultText;
+
+  geminiClient.addHistory({
+    role: 'user',
+    parts: [
+      {
+        text: `I ran the following shell command:
+\`\`\`sh
+${rawQuery}
+\`\`\`
+
+This produced the following result:
+\`\`\`
+${modelContent}
+\`\`\``,
+      },
+    ],
+  });
+}
+
+/**
+ * Hook to process shell commands.
+ * Orchestrates command execution and updates history and agent context.
  */
 export const useShellCommandProcessor = (
   addItemToHistory: UseHistoryManagerReturn['addItem'],
@@ -31,227 +212,126 @@ export const useShellCommandProcessor = (
   onExec: (command: Promise<void>) => void,
   onDebugMessage: (message: string) => void,
   config: Config,
-  executeCommand?: typeof ExecType, // injectable for testing
+  geminiClient: GeminiClient,
 ) => {
-  /**
-   * Checks if the query is a shell command, executes it, and adds results to history.
-   * @returns True if the query was handled as a shell command, false otherwise.
-   */
   const handleShellCommand = useCallback(
     (rawQuery: PartListUnion, abortSignal: AbortSignal): boolean => {
-      if (typeof rawQuery !== 'string') {
+      if (typeof rawQuery !== 'string' || rawQuery.trim() === '') {
         return false;
       }
 
-      const isWindows = os.platform() === 'win32';
-      let commandToExecute: string;
-      let pwdFilePath: string | undefined;
-
-      if (isWindows) {
-        commandToExecute = rawQuery;
-      } else {
-        // wrap command to write pwd to temporary file
-        let command = rawQuery.trim();
-        const pwdFileName = `shell_pwd_${crypto
-          .randomBytes(6)
-          .toString('hex')}.tmp`;
-        pwdFilePath = path.join(os.tmpdir(), pwdFileName);
-        if (!command.endsWith('&')) command += ';';
-        // note here we could also restore a previous pwd with `cd {cwd}; { ... }`
-        commandToExecute = `{ ${command} }; __code=$?; pwd >${pwdFilePath}; exit $__code`;
-      }
-
       const userMessageTimestamp = Date.now();
       addItemToHistory(
         { type: 'user_shell', text: rawQuery },
         userMessageTimestamp,
       );
 
-      if (rawQuery.trim() === '') {
-        addItemToHistory(
-          { type: 'error', text: 'Empty shell command.' },
-          userMessageTimestamp,
-        );
-        return true; // Handled (by showing error)
-      }
-
+      const isWindows = os.platform() === 'win32';
       const targetDir = config.getTargetDir();
-      onDebugMessage(
-        `Executing shell command in ${targetDir}: ${commandToExecute}`,
-      );
-      const execOptions = {
-        cwd: targetDir,
-      };
+      let commandToExecute = rawQuery;
+      let pwdFilePath: string | undefined;
+
+      // On non-windows, wrap the command to capture the final working directory.
+      if (!isWindows) {
+        let command = rawQuery.trim();
+        const pwdFileName = `shell_pwd_${crypto.randomBytes(6).toString('hex')}.tmp`;
+        pwdFilePath = path.join(os.tmpdir(), pwdFileName);
+        // Ensure command ends with a separator before adding our own.
+        if (!command.endsWith(';') && !command.endsWith('&')) {
+          command += ';';
+        }
+        commandToExecute = `{ ${command} }; __code=$?; pwd > "${pwdFilePath}"; exit $__code`;
+      }
 
       const execPromise = new Promise<void>((resolve) => {
-        if (executeCommand) {
-          executeCommand(
-            commandToExecute,
-            execOptions,
-            (error, stdout, stderr) => {
-              if (error) {
-                addItemToHistory(
-                  {
-                    type: 'error',
-                    // remove wrapper from user's command in error message
-                    text: error.message.replace(commandToExecute, rawQuery),
-                  },
-                  userMessageTimestamp,
-                );
-              } else {
-                let output = '';
-                if (stdout) output += stdout;
-                if (stderr) output += (output ? '\n' : '') + stderr; // Include stderr as info
+        let lastUpdateTime = 0;
 
-                addItemToHistory(
-                  {
-                    type: 'info',
-                    text: output || '(Command produced no output)',
-                  },
-                  userMessageTimestamp,
-                );
-              }
-              if (pwdFilePath && fs.existsSync(pwdFilePath)) {
-                const pwd = fs.readFileSync(pwdFilePath, 'utf8').trim();
-                if (pwd !== targetDir) {
-                  addItemToHistory(
-                    {
-                      type: 'info',
-                      text: `WARNING: shell mode is stateless; \`cd ${pwd}\` will not apply to next command`,
-                    },
-                    userMessageTimestamp,
-                  );
-                }
-                fs.unlinkSync(pwdFilePath);
-              }
-              resolve();
-            },
-          );
-        } else {
-          const child = isWindows
-            ? spawn('cmd.exe', ['/c', commandToExecute], {
-                cwd: targetDir,
-                stdio: ['ignore', 'pipe', 'pipe'],
-              })
-            : spawn('bash', ['-c', commandToExecute], {
-                cwd: targetDir,
-                stdio: ['ignore', 'pipe', 'pipe'],
-                detached: true, // Important for process group killing
-              });
+        onDebugMessage(`Executing in ${targetDir}: ${commandToExecute}`);
+        executeShellCommand(
+          commandToExecute,
+          targetDir,
+          abortSignal,
+          (streamedOutput) => {
+            // Throttle pending UI updates to avoid excessive re-renders.
+            if (Date.now() - lastUpdateTime > OUTPUT_UPDATE_INTERVAL_MS) {
+              setPendingHistoryItem({ type: 'info', text: streamedOutput });
+              lastUpdateTime = Date.now();
+            }
+          },
+          onDebugMessage,
+        )
+          .then((result) => {
+            // TODO(abhipatel12) - Consider updating pending item and using timeout to ensure
+            // there is no jump where intermediate output is skipped.
+            setPendingHistoryItem(null);
 
-          let exited = false;
-          let output = '';
-          let lastUpdateTime = Date.now();
-          const handleOutput = (data: Buffer) => {
-            // continue to consume post-exit for background processes
-            // removing listeners can overflow OS buffer and block subprocesses
-            // destroying (e.g. child.stdout.destroy()) can terminate subprocesses via SIGPIPE
-            if (!exited) {
-              output += stripAnsi(data.toString());
-              if (Date.now() - lastUpdateTime > OUTPUT_UPDATE_INTERVAL_MS) {
-                setPendingHistoryItem({
-                  type: 'info',
-                  text: output,
-                });
-                lastUpdateTime = Date.now();
-              }
+            let historyItemType: HistoryItemWithoutId['type'] = 'info';
+            let mainContent: string;
+
+            // The context sent to the model utilizes a text tokenizer which means raw binary data is
+            // cannot be parsed and understood and thus would only pollute the context window and waste
+            // tokens.
+            if (isBinary(result.rawOutput)) {
+              mainContent =
+                '[Command produced binary output, which is not shown.]';
+            } else {
+              mainContent =
+                result.output.trim() || '(Command produced no output)';
             }
-          };
-          child.stdout.on('data', handleOutput);
-          child.stderr.on('data', handleOutput);
 
-          let error: Error | null = null;
-          child.on('error', (err: Error) => {
-            error = err;
-          });
+            let finalOutput = mainContent;
 
-          const abortHandler = async () => {
-            if (child.pid && !exited) {
-              onDebugMessage(
-                `Aborting shell command (PID: ${child.pid}) due to signal.`,
-              );
-              if (os.platform() === 'win32') {
-                // For Windows, use taskkill to kill the process tree
-                spawn('taskkill', ['/pid', child.pid.toString(), '/f', '/t']);
-              } else {
-                try {
-                  // attempt to SIGTERM process group (negative PID)
-                  // fall back to SIGKILL (to group) after 200ms
-                  process.kill(-child.pid, 'SIGTERM');
-                  await new Promise((resolve) => setTimeout(resolve, 200));
-                  if (child.pid && !exited) {
-                    process.kill(-child.pid, 'SIGKILL');
-                  }
-                } catch (_e) {
-                  // if group kill fails, fall back to killing just the main process
-                  try {
-                    if (child.pid) {
-                      child.kill('SIGKILL');
-                    }
-                  } catch (_e) {
-                    console.error(
-                      `failed to kill shell process ${child.pid}: ${_e}`,
-                    );
-                  }
-                }
+            if (result.error) {
+              historyItemType = 'error';
+              finalOutput = `${result.error.message}\n${finalOutput}`;
+            } else if (result.aborted) {
+              finalOutput = `Command was cancelled.\n${finalOutput}`;
+            } else if (result.signal) {
+              historyItemType = 'error';
+              finalOutput = `Command terminated by signal: ${result.signal}.\n${finalOutput}`;
+            } else if (result.exitCode !== 0) {
+              historyItemType = 'error';
+              finalOutput = `Command exited with code ${result.exitCode}.\n${finalOutput}`;
+            }
+
+            if (pwdFilePath && fs.existsSync(pwdFilePath)) {
+              const finalPwd = fs.readFileSync(pwdFilePath, 'utf8').trim();
+              if (finalPwd && finalPwd !== targetDir) {
+                const warning = `WARNING: shell mode is stateless; the directory change to '${finalPwd}' will not persist.`;
+                finalOutput = `${warning}\n\n${finalOutput}`;
               }
             }
-          };
 
-          abortSignal.addEventListener('abort', abortHandler, { once: true });
+            // Add the complete, contextual result to the local UI history.
+            addItemToHistory(
+              { type: historyItemType, text: finalOutput },
+              userMessageTimestamp,
+            );
 
-          child.on('exit', (code, signal) => {
-            exited = true;
-            abortSignal.removeEventListener('abort', abortHandler);
+            // Add the same complete, contextual result to the LLM's history.
+            addShellCommandToGeminiHistory(geminiClient, rawQuery, finalOutput);
+          })
+          .catch((err) => {
             setPendingHistoryItem(null);
-            output = output.trim() || '(Command produced no output)';
-            if (error) {
-              const text = `${error.message.replace(commandToExecute, rawQuery)}\n${output}`;
-              addItemToHistory({ type: 'error', text }, userMessageTimestamp);
-            } else if (code !== null && code !== 0) {
-              const text = `Command exited with code ${code}\n${output}`;
-              addItemToHistory({ type: 'error', text }, userMessageTimestamp);
-            } else if (abortSignal.aborted) {
-              addItemToHistory(
-                {
-                  type: 'info',
-                  text: `Command was cancelled.\n${output}`,
-                },
-                userMessageTimestamp,
-              );
-            } else if (signal) {
-              const text = `Command terminated with signal ${signal}.\n${output}`;
-              addItemToHistory({ type: 'error', text }, userMessageTimestamp);
-            } else {
-              addItemToHistory(
-                { type: 'info', text: output + '\n' },
-                userMessageTimestamp,
-              );
-            }
+            const errorMessage =
+              err instanceof Error ? err.message : String(err);
+            addItemToHistory(
+              {
+                type: 'error',
+                text: `An unexpected error occurred: ${errorMessage}`,
+              },
+              userMessageTimestamp,
+            );
+          })
+          .finally(() => {
             if (pwdFilePath && fs.existsSync(pwdFilePath)) {
-              const pwd = fs.readFileSync(pwdFilePath, 'utf8').trim();
-              if (pwd !== targetDir) {
-                addItemToHistory(
-                  {
-                    type: 'info',
-                    text: `WARNING: shell mode is stateless; \`cd ${pwd}\` will not apply to next command`,
-                  },
-                  userMessageTimestamp,
-                );
-              }
               fs.unlinkSync(pwdFilePath);
             }
             resolve();
           });
-        }
       });
 
-      try {
-        onExec(execPromise);
-      } catch (_e) {
-        // silently ignore errors from this since it's from the caller
-      }
-
+      onExec(execPromise);
       return true; // Command was initiated
     },
     [
@@ -260,7 +340,7 @@ export const useShellCommandProcessor = (
       addItemToHistory,
       setPendingHistoryItem,
       onExec,
-      executeCommand,
+      geminiClient,
     ],
   );