pre-release commit

packages/cli/src/ui/hooks/shellCommandProcessor.ts

@@ -0,0 +1,348 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { spawn } from 'child_process';
+import { StringDecoder } from 'string_decoder';
+import type { HistoryItemWithoutId } from '../types.js';
+import { useCallback } from 'react';
+import { Config, GeminiClient } from '@qwen/qwen-code-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';
+import os from 'os';
+import fs from 'fs';
+import stripAnsi from 'strip-ansi';
+
+const OUTPUT_UPDATE_INTERVAL_MS = 1000;
+const MAX_OUTPUT_LENGTH = 10000;
+
+/**
+ * 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'],
+  setPendingHistoryItem: React.Dispatch<
+    React.SetStateAction<HistoryItemWithoutId | null>
+  >,
+  onExec: (command: Promise<void>) => void,
+  onDebugMessage: (message: string) => void,
+  config: Config,
+  geminiClient: GeminiClient,
+) => {
+  const handleShellCommand = useCallback(
+    (rawQuery: PartListUnion, abortSignal: AbortSignal): boolean => {
+      if (typeof rawQuery !== 'string' || rawQuery.trim() === '') {
+        return false;
+      }
+
+      const userMessageTimestamp = Date.now();
+      addItemToHistory(
+        { type: 'user_shell', text: rawQuery },
+        userMessageTimestamp,
+      );
+
+      const isWindows = os.platform() === 'win32';
+      const targetDir = config.getTargetDir();
+      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) => {
+        let lastUpdateTime = 0;
+
+        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 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)';
+            }
+
+            let finalOutput = mainContent;
+
+            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}`;
+              }
+            }
+
+            // Add the complete, contextual result to the local UI history.
+            addItemToHistory(
+              { type: historyItemType, text: finalOutput },
+              userMessageTimestamp,
+            );
+
+            // Add the same complete, contextual result to the LLM's history.
+            addShellCommandToGeminiHistory(geminiClient, rawQuery, finalOutput);
+          })
+          .catch((err) => {
+            setPendingHistoryItem(null);
+            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)) {
+              fs.unlinkSync(pwdFilePath);
+            }
+            resolve();
+          });
+      });
+
+      onExec(execPromise);
+      return true; // Command was initiated
+    },
+    [
+      config,
+      onDebugMessage,
+      addItemToHistory,
+      setPendingHistoryItem,
+      onExec,
+      geminiClient,
+    ],
+  );
+
+  return { handleShellCommand };
+};