diff options
Diffstat (limited to 'packages/cli/src')
| -rw-r--r-- | packages/cli/src/ui/hooks/shellCommandProcessor.test.ts | 512 | ||||
| -rw-r--r-- | packages/cli/src/ui/hooks/shellCommandProcessor.ts | 431 | ||||
| -rw-r--r-- | packages/cli/src/ui/utils/textUtils.test.ts | 41 | ||||
| -rw-r--r-- | packages/cli/src/ui/utils/textUtils.ts | 29 |
4 files changed, 561 insertions, 452 deletions
diff --git a/packages/cli/src/ui/hooks/shellCommandProcessor.test.ts b/packages/cli/src/ui/hooks/shellCommandProcessor.test.ts index 1b268502..d5270aba 100644 --- a/packages/cli/src/ui/hooks/shellCommandProcessor.test.ts +++ b/packages/cli/src/ui/hooks/shellCommandProcessor.test.ts @@ -5,70 +5,86 @@ */ import { act, renderHook } from '@testing-library/react'; -import { vi } from 'vitest'; -import { spawn } from 'child_process'; -import type { ChildProcessWithoutNullStreams } from 'child_process'; -import { useShellCommandProcessor } from './shellCommandProcessor'; -import { Config, GeminiClient } from '@google/gemini-cli-core'; -import * as fs from 'fs'; -import EventEmitter from 'events'; -import { ToolCallStatus } from '../types'; +import { + vi, + describe, + it, + expect, + beforeEach, + afterEach, + type Mock, +} from 'vitest'; -// Mock dependencies -vi.mock('child_process'); +const mockIsBinary = vi.hoisted(() => vi.fn()); +const mockShellExecutionService = vi.hoisted(() => vi.fn()); +vi.mock('@google/gemini-cli-core', async (importOriginal) => { + const original = + await importOriginal<typeof import('@google/gemini-cli-core')>(); + return { + ...original, + ShellExecutionService: { execute: mockShellExecutionService }, + isBinary: mockIsBinary, + }; +}); vi.mock('fs'); -vi.mock('os', () => ({ - default: { - platform: () => 'linux', - tmpdir: () => '/tmp', - homedir: () => '/home/user', - }, - platform: () => 'linux', - tmpdir: () => '/tmp', - homedir: () => '/home/user', -})); -vi.mock('@google/gemini-cli-core'); -vi.mock('../utils/textUtils.js', () => ({ - isBinary: vi.fn(), -})); +vi.mock('os'); +vi.mock('crypto'); +vi.mock('../utils/textUtils.js'); + +import { + useShellCommandProcessor, + OUTPUT_UPDATE_INTERVAL_MS, +} from './shellCommandProcessor.js'; +import { + type Config, + type GeminiClient, + type ShellExecutionResult, + type ShellOutputEvent, +} from '@google/gemini-cli-core'; +import * as fs from 'fs'; +import * as os from 'os'; +import * as path from 'path'; +import * as crypto from 'crypto'; +import { ToolCallStatus } from '../types.js'; describe('useShellCommandProcessor', () => { - let spawnEmitter: EventEmitter; - let addItemToHistoryMock: vi.Mock; - let setPendingHistoryItemMock: vi.Mock; - let onExecMock: vi.Mock; - let onDebugMessageMock: vi.Mock; - let configMock: Config; - let geminiClientMock: GeminiClient; + let addItemToHistoryMock: Mock; + let setPendingHistoryItemMock: Mock; + let onExecMock: Mock; + let onDebugMessageMock: Mock; + let mockConfig: Config; + let mockGeminiClient: GeminiClient; - beforeEach(() => { - spawnEmitter = new EventEmitter(); - spawnEmitter.stdout = new EventEmitter(); - spawnEmitter.stderr = new EventEmitter(); - vi.mocked(spawn).mockReturnValue( - spawnEmitter as ChildProcessWithoutNullStreams, - ); + let mockShellOutputCallback: (event: ShellOutputEvent) => void; + let resolveExecutionPromise: (result: ShellExecutionResult) => void; - vi.spyOn(fs, 'existsSync').mockReturnValue(false); - vi.spyOn(fs, 'readFileSync').mockReturnValue(''); - vi.spyOn(fs, 'unlinkSync').mockReturnValue(undefined); + beforeEach(() => { + vi.clearAllMocks(); addItemToHistoryMock = vi.fn(); setPendingHistoryItemMock = vi.fn(); onExecMock = vi.fn(); onDebugMessageMock = vi.fn(); + mockConfig = { getTargetDir: () => '/test/dir' } as Config; + mockGeminiClient = { addHistory: vi.fn() } as unknown as GeminiClient; - configMock = { - getTargetDir: () => '/test/dir', - } as unknown as Config; - - geminiClientMock = { - addHistory: vi.fn(), - } as unknown as GeminiClient; - }); + vi.mocked(os.platform).mockReturnValue('linux'); + vi.mocked(os.tmpdir).mockReturnValue('/tmp'); + (vi.mocked(crypto.randomBytes) as Mock).mockReturnValue( + Buffer.from('abcdef', 'hex'), + ); + mockIsBinary.mockReturnValue(false); + vi.mocked(fs.existsSync).mockReturnValue(false); - afterEach(() => { - vi.restoreAllMocks(); + mockShellExecutionService.mockImplementation((_cmd, _cwd, callback) => { + mockShellOutputCallback = callback; + return { + pid: 12345, + result: new Promise((resolve) => { + resolveExecutionPromise = resolve; + }), + }; + }); }); const renderProcessorHook = () => @@ -78,140 +94,386 @@ describe('useShellCommandProcessor', () => { setPendingHistoryItemMock, onExecMock, onDebugMessageMock, - configMock, - geminiClientMock, + mockConfig, + mockGeminiClient, ), ); - it('should execute a command and update history on success', async () => { + const createMockServiceResult = ( + overrides: Partial<ShellExecutionResult> = {}, + ): ShellExecutionResult => ({ + rawOutput: Buffer.from(overrides.output || ''), + output: 'Success', + stdout: 'Success', + stderr: '', + exitCode: 0, + signal: null, + error: null, + aborted: false, + pid: 12345, + ...overrides, + }); + + it('should initiate command execution and set pending state', async () => { const { result } = renderProcessorHook(); - const abortController = new AbortController(); act(() => { - result.current.handleShellCommand('ls -l', abortController.signal); + result.current.handleShellCommand('ls -l', new AbortController().signal); }); - expect(spawn).toHaveBeenCalledWith( - 'bash', - ['-c', expect.any(String)], - expect.objectContaining({ - env: expect.objectContaining({ - GEMINI_CLI: '1', + expect(addItemToHistoryMock).toHaveBeenCalledWith( + { type: 'user_shell', text: 'ls -l' }, + expect.any(Number), + ); + expect(setPendingHistoryItemMock).toHaveBeenCalledWith({ + type: 'tool_group', + tools: [ + expect.objectContaining({ + name: 'Shell Command', + status: ToolCallStatus.Executing, }), - }), + ], + }); + const tmpFile = path.join(os.tmpdir(), 'shell_pwd_abcdef.tmp'); + const wrappedCommand = `{ ls -l; }; __code=$?; pwd > "${tmpFile}"; exit $__code`; + expect(mockShellExecutionService).toHaveBeenCalledWith( + wrappedCommand, + '/test/dir', + expect.any(Function), + expect.any(Object), ); + expect(onExecMock).toHaveBeenCalledWith(expect.any(Promise)); + }); - expect(onExecMock).toHaveBeenCalledTimes(1); + it('should handle successful execution and update history correctly', async () => { + const { result } = renderProcessorHook(); + + act(() => { + result.current.handleShellCommand( + 'echo "ok"', + new AbortController().signal, + ); + }); const execPromise = onExecMock.mock.calls[0][0]; - // Simulate stdout act(() => { - spawnEmitter.stdout.emit('data', Buffer.from('file1.txt\nfile2.txt')); + resolveExecutionPromise(createMockServiceResult({ output: 'ok' })); + }); + await act(async () => await execPromise); + + expect(setPendingHistoryItemMock).toHaveBeenCalledWith(null); + expect(addItemToHistoryMock).toHaveBeenCalledTimes(2); // Initial + final + expect(addItemToHistoryMock.mock.calls[1][0]).toEqual( + expect.objectContaining({ + tools: [ + expect.objectContaining({ + status: ToolCallStatus.Success, + resultDisplay: 'ok', + }), + ], + }), + ); + expect(mockGeminiClient.addHistory).toHaveBeenCalled(); + }); + + it('should handle command failure and display error status', async () => { + const { result } = renderProcessorHook(); + + act(() => { + result.current.handleShellCommand( + 'bad-cmd', + new AbortController().signal, + ); }); + const execPromise = onExecMock.mock.calls[0][0]; - // Simulate process exit act(() => { - spawnEmitter.emit('exit', 0, null); + resolveExecutionPromise( + createMockServiceResult({ exitCode: 127, output: 'not found' }), + ); }); + await act(async () => await execPromise); + + const finalHistoryItem = addItemToHistoryMock.mock.calls[1][0]; + expect(finalHistoryItem.tools[0].status).toBe(ToolCallStatus.Error); + expect(finalHistoryItem.tools[0].resultDisplay).toContain( + 'Command exited with code 127', + ); + expect(finalHistoryItem.tools[0].resultDisplay).toContain('not found'); + }); - await act(async () => { - await execPromise; + describe('UI Streaming and Throttling', () => { + beforeEach(() => { + vi.useFakeTimers({ toFake: ['Date'] }); + }); + afterEach(() => { + vi.useRealTimers(); }); - expect(addItemToHistoryMock).toHaveBeenCalledTimes(2); - expect(addItemToHistoryMock.mock.calls[1][0]).toEqual({ - type: 'tool_group', - tools: [ + it('should throttle pending UI updates for text streams', async () => { + const { result } = renderProcessorHook(); + act(() => { + result.current.handleShellCommand( + 'stream', + new AbortController().signal, + ); + }); + + // Simulate rapid output + act(() => { + mockShellOutputCallback({ + type: 'data', + stream: 'stdout', + chunk: 'hello', + }); + }); + + // Should not have updated the UI yet + expect(setPendingHistoryItemMock).toHaveBeenCalledTimes(1); // Only the initial call + + // Advance time and send another event to trigger the throttled update + await act(async () => { + await vi.advanceTimersByTimeAsync(OUTPUT_UPDATE_INTERVAL_MS + 1); + }); + act(() => { + mockShellOutputCallback({ + type: 'data', + stream: 'stdout', + chunk: ' world', + }); + }); + + // Should now have been called with the cumulative output + expect(setPendingHistoryItemMock).toHaveBeenCalledTimes(2); + expect(setPendingHistoryItemMock).toHaveBeenLastCalledWith( expect.objectContaining({ - name: 'Shell Command', - description: 'ls -l', - status: ToolCallStatus.Success, - resultDisplay: 'file1.txt\nfile2.txt', + tools: [expect.objectContaining({ resultDisplay: 'hello world' })], }), - ], + ); + }); + + it('should show binary progress messages correctly', async () => { + const { result } = renderProcessorHook(); + act(() => { + result.current.handleShellCommand( + 'cat img', + new AbortController().signal, + ); + }); + + // Should immediately show the detection message + act(() => { + mockShellOutputCallback({ type: 'binary_detected' }); + }); + await act(async () => { + await vi.advanceTimersByTimeAsync(OUTPUT_UPDATE_INTERVAL_MS + 1); + }); + // Send another event to trigger the update + act(() => { + mockShellOutputCallback({ type: 'binary_progress', bytesReceived: 0 }); + }); + + expect(setPendingHistoryItemMock).toHaveBeenLastCalledWith( + expect.objectContaining({ + tools: [ + expect.objectContaining({ + resultDisplay: '[Binary output detected. Halting stream...]', + }), + ], + }), + ); + + // Now test progress updates + await act(async () => { + await vi.advanceTimersByTimeAsync(OUTPUT_UPDATE_INTERVAL_MS + 1); + }); + act(() => { + mockShellOutputCallback({ + type: 'binary_progress', + bytesReceived: 2048, + }); + }); + + expect(setPendingHistoryItemMock).toHaveBeenLastCalledWith( + expect.objectContaining({ + tools: [ + expect.objectContaining({ + resultDisplay: '[Receiving binary output... 2.0 KB received]', + }), + ], + }), + ); }); - expect(geminiClientMock.addHistory).toHaveBeenCalledTimes(1); }); - it('should handle binary output', async () => { + it('should not wrap the command on Windows', async () => { + vi.mocked(os.platform).mockReturnValue('win32'); + const { result } = renderProcessorHook(); + + act(() => { + result.current.handleShellCommand('dir', new AbortController().signal); + }); + + expect(mockShellExecutionService).toHaveBeenCalledWith( + 'dir', + '/test/dir', + expect.any(Function), + expect.any(Object), + ); + }); + + it('should handle command abort and display cancelled status', async () => { const { result } = renderProcessorHook(); const abortController = new AbortController(); - const { isBinary } = await import('../utils/textUtils.js'); - (isBinary as vi.Mock).mockReturnValue(true); act(() => { - result.current.handleShellCommand( - 'cat myimage.png', - abortController.signal, + result.current.handleShellCommand('sleep 5', abortController.signal); + }); + const execPromise = onExecMock.mock.calls[0][0]; + + act(() => { + abortController.abort(); + resolveExecutionPromise( + createMockServiceResult({ aborted: true, output: 'Canceled' }), ); }); + await act(async () => await execPromise); - expect(onExecMock).toHaveBeenCalledTimes(1); - const execPromise = onExecMock.mock.calls[0][0]; + const finalHistoryItem = addItemToHistoryMock.mock.calls[1][0]; + expect(finalHistoryItem.tools[0].status).toBe(ToolCallStatus.Canceled); + expect(finalHistoryItem.tools[0].resultDisplay).toContain( + 'Command was cancelled.', + ); + }); + + it('should handle binary output result correctly', async () => { + const { result } = renderProcessorHook(); + const binaryBuffer = Buffer.from([0x89, 0x50, 0x4e, 0x47]); + mockIsBinary.mockReturnValue(true); act(() => { - spawnEmitter.stdout.emit('data', Buffer.from([0x89, 0x50, 0x4e, 0x47])); + result.current.handleShellCommand( + 'cat image.png', + new AbortController().signal, + ); }); + const execPromise = onExecMock.mock.calls[0][0]; act(() => { - spawnEmitter.emit('exit', 0, null); + resolveExecutionPromise( + createMockServiceResult({ rawOutput: binaryBuffer }), + ); }); + await act(async () => await execPromise); - await act(async () => { - await execPromise; + const finalHistoryItem = addItemToHistoryMock.mock.calls[1][0]; + expect(finalHistoryItem.tools[0].status).toBe(ToolCallStatus.Success); + expect(finalHistoryItem.tools[0].resultDisplay).toBe( + '[Command produced binary output, which is not shown.]', + ); + }); + + it('should handle promise rejection and show an error', async () => { + const { result } = renderProcessorHook(); + const testError = new Error('Unexpected failure'); + mockShellExecutionService.mockImplementation(() => ({ + pid: 12345, + result: Promise.reject(testError), + })); + + act(() => { + result.current.handleShellCommand( + 'a-command', + new AbortController().signal, + ); }); + const execPromise = onExecMock.mock.calls[0][0]; + + await act(async () => await execPromise); + expect(setPendingHistoryItemMock).toHaveBeenCalledWith(null); expect(addItemToHistoryMock).toHaveBeenCalledTimes(2); expect(addItemToHistoryMock.mock.calls[1][0]).toEqual({ - type: 'tool_group', - tools: [ - expect.objectContaining({ - name: 'Shell Command', - description: 'cat myimage.png', - status: ToolCallStatus.Success, - resultDisplay: - '[Command produced binary output, which is not shown.]', - }), - ], + type: 'error', + text: 'An unexpected error occurred: Unexpected failure', }); }); - it('should handle command failure', async () => { + it('should handle synchronous errors during execution and clean up resources', async () => { + const testError = new Error('Synchronous spawn error'); + mockShellExecutionService.mockImplementation(() => { + throw testError; + }); + // Mock that the temp file was created before the error was thrown + vi.mocked(fs.existsSync).mockReturnValue(true); + const { result } = renderProcessorHook(); - const abortController = new AbortController(); act(() => { result.current.handleShellCommand( - 'a-bad-command', - abortController.signal, + 'a-command', + new AbortController().signal, ); }); - const execPromise = onExecMock.mock.calls[0][0]; - act(() => { - spawnEmitter.stderr.emit('data', Buffer.from('command not found')); - }); + await act(async () => await execPromise); - act(() => { - spawnEmitter.emit('exit', 127, null); + expect(setPendingHistoryItemMock).toHaveBeenCalledWith(null); + expect(addItemToHistoryMock).toHaveBeenCalledTimes(2); + expect(addItemToHistoryMock.mock.calls[1][0]).toEqual({ + type: 'error', + text: 'An unexpected error occurred: Synchronous spawn error', }); + const tmpFile = path.join(os.tmpdir(), 'shell_pwd_abcdef.tmp'); + // Verify that the temporary file was cleaned up + expect(vi.mocked(fs.unlinkSync)).toHaveBeenCalledWith(tmpFile); + }); + + describe('Directory Change Warning', () => { + it('should show a warning if the working directory changes', async () => { + const tmpFile = path.join(os.tmpdir(), 'shell_pwd_abcdef.tmp'); + vi.mocked(fs.existsSync).mockReturnValue(true); + vi.mocked(fs.readFileSync).mockReturnValue('/test/dir/new'); // A different directory + + const { result } = renderProcessorHook(); + act(() => { + result.current.handleShellCommand( + 'cd new', + new AbortController().signal, + ); + }); + const execPromise = onExecMock.mock.calls[0][0]; + + act(() => { + resolveExecutionPromise(createMockServiceResult()); + }); + await act(async () => await execPromise); - await act(async () => { - await execPromise; + const finalHistoryItem = addItemToHistoryMock.mock.calls[1][0]; + expect(finalHistoryItem.tools[0].resultDisplay).toContain( + "WARNING: shell mode is stateless; the directory change to '/test/dir/new' will not persist.", + ); + expect(vi.mocked(fs.unlinkSync)).toHaveBeenCalledWith(tmpFile); }); - expect(addItemToHistoryMock).toHaveBeenCalledTimes(2); - expect(addItemToHistoryMock.mock.calls[1][0]).toEqual({ - type: 'tool_group', - tools: [ - expect.objectContaining({ - name: 'Shell Command', - description: 'a-bad-command', - status: ToolCallStatus.Error, - resultDisplay: 'Command exited with code 127.\ncommand not found', - }), - ], + it('should NOT show a warning if the directory does not change', async () => { + vi.mocked(fs.existsSync).mockReturnValue(true); + vi.mocked(fs.readFileSync).mockReturnValue('/test/dir'); // The same directory + + const { result } = renderProcessorHook(); + act(() => { + result.current.handleShellCommand('ls', new AbortController().signal); + }); + const execPromise = onExecMock.mock.calls[0][0]; + + act(() => { + resolveExecutionPromise(createMockServiceResult()); + }); + await act(async () => await execPromise); + + const finalHistoryItem = addItemToHistoryMock.mock.calls[1][0]; + expect(finalHistoryItem.tools[0].resultDisplay).not.toContain('WARNING'); }); }); }); diff --git a/packages/cli/src/ui/hooks/shellCommandProcessor.ts b/packages/cli/src/ui/hooks/shellCommandProcessor.ts index 9e343f90..08df0a74 100644 --- a/packages/cli/src/ui/hooks/shellCommandProcessor.ts +++ b/packages/cli/src/ui/hooks/shellCommandProcessor.ts @@ -4,8 +4,6 @@ * SPDX-License-Identifier: Apache-2.0 */ -import { spawn } from 'child_process'; -import { TextDecoder } from 'util'; import { HistoryItemWithoutId, IndividualToolCallDisplay, @@ -15,186 +13,22 @@ import { useCallback } from 'react'; import { Config, GeminiClient, - getCachedEncodingForBuffer, + isBinary, + ShellExecutionResult, + ShellExecutionService, } from '@google/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 { SHELL_COMMAND_NAME } from '../constants.js'; +import { formatMemoryUsage } from '../utils/formatters.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; +export 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 - env: { - ...process.env, - GEMINI_CLI: '1', - }, - }); - - // 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 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') => { - if (!stdoutDecoder || !stderrDecoder) { - const encoding = getCachedEncodingForBuffer(data); - stdoutDecoder = new TextDecoder(encoding); - stderrDecoder = new TextDecoder(encoding); - } - - 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.decode(data, { stream: true }) - : stderrDecoder.decode(data, { stream: true }); - 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) { - // Fall back 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 - if (stdoutDecoder) { - stdout += stdoutDecoder.decode(); - } - if (stderrDecoder) { - stderr += stderrDecoder.decode(); - } - - 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, @@ -227,7 +61,6 @@ ${modelContent} * Hook to process shell commands. * Orchestrates command execution and updates history and agent context. */ - export const useShellCommandProcessor = ( addItemToHistory: UseHistoryManagerReturn['addItem'], setPendingHistoryItem: React.Dispatch< @@ -269,7 +102,11 @@ export const useShellCommandProcessor = ( } const execPromise = new Promise<void>((resolve) => { - let lastUpdateTime = 0; + let lastUpdateTime = Date.now(); + let cumulativeStdout = ''; + let cumulativeStderr = ''; + let isBinaryStream = false; + let binaryBytesReceived = 0; const initialToolDisplay: IndividualToolCallDisplay = { callId, @@ -285,103 +122,183 @@ export const useShellCommandProcessor = ( tools: [initialToolDisplay], }); + let executionPid: number | undefined; + + const abortHandler = () => { + onDebugMessage( + `Aborting shell command (PID: ${executionPid ?? 'unknown'})`, + ); + }; + abortSignal.addEventListener('abort', abortHandler, { once: true }); + 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: 'tool_group', - tools: [ - { ...initialToolDisplay, resultDisplay: streamedOutput }, - ], - }); - lastUpdateTime = Date.now(); - } - }, - onDebugMessage, - ) - .then((result) => { - setPendingHistoryItem(null); - let mainContent: string; + try { + const { pid, result } = ShellExecutionService.execute( + commandToExecute, + targetDir, + (event) => { + switch (event.type) { + case 'data': + // Do not process text data if we've already switched to binary mode. + if (isBinaryStream) break; + if (event.stream === 'stdout') { + cumulativeStdout += event.chunk; + } else { + cumulativeStderr += event.chunk; + } + break; + case 'binary_detected': + isBinaryStream = true; + break; + case 'binary_progress': + isBinaryStream = true; + binaryBytesReceived = event.bytesReceived; + break; + default: { + throw new Error('An unhandled ShellOutputEvent was found.'); + } + } + + // Compute the display string based on the *current* state. + let currentDisplayOutput: string; + if (isBinaryStream) { + if (binaryBytesReceived > 0) { + currentDisplayOutput = `[Receiving binary output... ${formatMemoryUsage( + binaryBytesReceived, + )} received]`; + } else { + currentDisplayOutput = + '[Binary output detected. Halting stream...]'; + } + } else { + currentDisplayOutput = + cumulativeStdout + + (cumulativeStderr ? `\n${cumulativeStderr}` : ''); + } + + // Throttle pending UI updates to avoid excessive re-renders. + if (Date.now() - lastUpdateTime > OUTPUT_UPDATE_INTERVAL_MS) { + setPendingHistoryItem({ + type: 'tool_group', + tools: [ + { + ...initialToolDisplay, + resultDisplay: currentDisplayOutput, + }, + ], + }); + lastUpdateTime = Date.now(); + } + }, + abortSignal, + ); - if (isBinary(result.rawOutput)) { - mainContent = - '[Command produced binary output, which is not shown.]'; - } else { - mainContent = - result.output.trim() || '(Command produced no output)'; - } + executionPid = pid; - let finalOutput = mainContent; - let finalStatus = ToolCallStatus.Success; + result + .then((result: ShellExecutionResult) => { + setPendingHistoryItem(null); - if (result.error) { - finalStatus = ToolCallStatus.Error; - finalOutput = `${result.error.message}\n${finalOutput}`; - } else if (result.aborted) { - finalStatus = ToolCallStatus.Canceled; - finalOutput = `Command was cancelled.\n${finalOutput}`; - } else if (result.signal) { - finalStatus = ToolCallStatus.Error; - finalOutput = `Command terminated by signal: ${result.signal}.\n${finalOutput}`; - } else if (result.exitCode !== 0) { - finalStatus = ToolCallStatus.Error; - finalOutput = `Command exited with code ${result.exitCode}.\n${finalOutput}`; - } + let mainContent: string; - 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}`; + if (isBinary(result.rawOutput)) { + mainContent = + '[Command produced binary output, which is not shown.]'; + } else { + mainContent = + result.output.trim() || '(Command produced no output)'; } - } - const finalToolDisplay: IndividualToolCallDisplay = { - ...initialToolDisplay, - status: finalStatus, - resultDisplay: finalOutput, - }; + let finalOutput = mainContent; + let finalStatus = ToolCallStatus.Success; + + if (result.error) { + finalStatus = ToolCallStatus.Error; + finalOutput = `${result.error.message}\n${finalOutput}`; + } else if (result.aborted) { + finalStatus = ToolCallStatus.Canceled; + finalOutput = `Command was cancelled.\n${finalOutput}`; + } else if (result.signal) { + finalStatus = ToolCallStatus.Error; + finalOutput = `Command terminated by signal: ${result.signal}.\n${finalOutput}`; + } else if (result.exitCode !== 0) { + finalStatus = ToolCallStatus.Error; + finalOutput = `Command exited with code ${result.exitCode}.\n${finalOutput}`; + } - // Add the complete, contextual result to the local UI history. - addItemToHistory( - { - type: 'tool_group', - tools: [finalToolDisplay], - } as HistoryItemWithoutId, - userMessageTimestamp, - ); + 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 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(); - }); + const finalToolDisplay: IndividualToolCallDisplay = { + ...initialToolDisplay, + status: finalStatus, + resultDisplay: finalOutput, + }; + + // Add the complete, contextual result to the local UI history. + addItemToHistory( + { + type: 'tool_group', + tools: [finalToolDisplay], + } as HistoryItemWithoutId, + 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(() => { + abortSignal.removeEventListener('abort', abortHandler); + if (pwdFilePath && fs.existsSync(pwdFilePath)) { + fs.unlinkSync(pwdFilePath); + } + resolve(); + }); + } catch (err) { + // This block handles synchronous errors from `execute` + setPendingHistoryItem(null); + const errorMessage = err instanceof Error ? err.message : String(err); + addItemToHistory( + { + type: 'error', + text: `An unexpected error occurred: ${errorMessage}`, + }, + userMessageTimestamp, + ); + + // Perform cleanup here as well + if (pwdFilePath && fs.existsSync(pwdFilePath)) { + fs.unlinkSync(pwdFilePath); + } + + resolve(); // Resolve the promise to unblock `onExec` + } }); onExec(execPromise); - return true; // Command was initiated + return true; }, [ config, diff --git a/packages/cli/src/ui/utils/textUtils.test.ts b/packages/cli/src/ui/utils/textUtils.test.ts deleted file mode 100644 index 5dd08875..00000000 --- a/packages/cli/src/ui/utils/textUtils.test.ts +++ /dev/null @@ -1,41 +0,0 @@ -/** - * @license - * Copyright 2025 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import { isBinary } from './textUtils'; - -describe('textUtils', () => { - describe('isBinary', () => { - it('should return true for a buffer containing a null byte', () => { - const buffer = Buffer.from([ - 0x89, 0x50, 0x4e, 0x47, 0x0d, 0x1a, 0x0a, 0x00, - ]); - expect(isBinary(buffer)).toBe(true); - }); - - it('should return false for a buffer containing only text', () => { - const buffer = Buffer.from('This is a test string.'); - expect(isBinary(buffer)).toBe(false); - }); - - it('should return false for an empty buffer', () => { - const buffer = Buffer.from([]); - expect(isBinary(buffer)).toBe(false); - }); - - it('should return false for a null or undefined buffer', () => { - expect(isBinary(null)).toBe(false); - expect(isBinary(undefined)).toBe(false); - }); - - it('should only check the sample size', () => { - const longBufferWithNullByteAtEnd = Buffer.concat([ - Buffer.from('a'.repeat(1024)), - Buffer.from([0x00]), - ]); - expect(isBinary(longBufferWithNullByteAtEnd, 512)).toBe(false); - }); - }); -}); diff --git a/packages/cli/src/ui/utils/textUtils.ts b/packages/cli/src/ui/utils/textUtils.ts index fa0abe9a..e4d8ea58 100644 --- a/packages/cli/src/ui/utils/textUtils.ts +++ b/packages/cli/src/ui/utils/textUtils.ts @@ -17,35 +17,6 @@ export const getAsciiArtWidth = (asciiArt: string): number => { return Math.max(...lines.map((line) => line.length)); }; -/** - * Checks if a Buffer is likely binary by testing for the presence of a NULL byte. - * The presence of a NULL byte is a strong indicator that the data is not plain text. - * @param data The Buffer to check. - * @param sampleSize The number of bytes from the start of the buffer to test. - * @returns True if a NULL byte is found, false otherwise. - */ -export function isBinary( - data: Buffer | null | undefined, - sampleSize = 512, -): boolean { - if (!data) { - return false; - } - - const sample = data.length > sampleSize ? data.subarray(0, sampleSize) : data; - - for (const byte of sample) { - // The presence of a NULL byte (0x00) is one of the most reliable - // indicators of a binary file. Text files should not contain them. - if (byte === 0) { - return true; - } - } - - // If no NULL bytes were found in the sample, we assume it's text. - return false; -} - /* * ------------------------------------------------------------------------- * Unicode‑aware helpers (work at the code‑point level rather than UTF‑16 |