feat: Adds shell command context to gemini history (#1076)
This commit is contained in:
Abhi
GitHub
parent
7f06ad40c5
commit
bedff2ca79
|
|
@ -4,290 +4,176 @@
|
|||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
/* eslint-disable @typescript-eslint/no-explicit-any */
|
||||
import { act, renderHook } from '@testing-library/react';
|
||||
import { useShellCommandProcessor } from './shellCommandProcessor.js';
|
||||
import { type Config } from '@gemini-cli/core';
|
||||
import { type PartListUnion } from '@google/genai';
|
||||
import { existsSync, readFileSync, unlinkSync } from 'fs';
|
||||
import type * as FsMod from 'fs';
|
||||
import type { exec as ExecType } from 'child_process'; // For typing the injected mock
|
||||
|
||||
// Mocks
|
||||
const mockAddItemToHistory = vi.fn();
|
||||
const mockSetPendingHistoryItem = vi.fn();
|
||||
const mockOnExec = vi.fn(async (promise) => await promise);
|
||||
const mockOnDebugMessage = vi.fn();
|
||||
const mockGetTargetDir = vi.fn();
|
||||
let mockExecuteCommand: ReturnType<typeof vi.fn>; // This will be our injected mock for child_process.exec
|
||||
|
||||
const mockConfig = {
|
||||
getTargetDir: mockGetTargetDir,
|
||||
} as unknown as Config;
|
||||
|
||||
vi.mock('crypto', () => ({
|
||||
default: {
|
||||
randomBytes: vi.fn(() => ({ toString: vi.fn(() => 'randomBytes') })),
|
||||
},
|
||||
randomBytes: vi.fn(() => ({ toString: vi.fn(() => 'randomBytes') })),
|
||||
}));
|
||||
|
||||
vi.mock('path', () => ({
|
||||
default: {
|
||||
join: vi.fn((...args) => args.join('/')),
|
||||
sep: '/',
|
||||
},
|
||||
join: vi.fn((...args) => args.join('/')),
|
||||
sep: '/',
|
||||
}));
|
||||
import { vi } from 'vitest';
|
||||
import { useShellCommandProcessor } from './shellCommandProcessor';
|
||||
import { Config, GeminiClient } from '@gemini-cli/core';
|
||||
import * as fs from 'fs';
|
||||
import EventEmitter from 'events';
|
||||
|
||||
// Mock dependencies
|
||||
vi.mock('child_process');
|
||||
vi.mock('fs');
|
||||
vi.mock('os', () => ({
|
||||
default: {
|
||||
tmpdir: vi.fn(() => '/tmp'),
|
||||
platform: vi.fn(() => 'linux'),
|
||||
platform: () => 'linux',
|
||||
tmpdir: () => '/tmp',
|
||||
},
|
||||
tmpdir: vi.fn(() => '/tmp'),
|
||||
platform: vi.fn(() => 'linux'),
|
||||
platform: () => 'linux',
|
||||
tmpdir: () => '/tmp',
|
||||
}));
|
||||
vi.mock('@gemini-cli/core');
|
||||
vi.mock('../utils/textUtils.js', () => ({
|
||||
isBinary: vi.fn(),
|
||||
}));
|
||||
|
||||
// Configure the fs mock to use new vi.fn() instances created within the factory
|
||||
vi.mock('fs', async (importOriginal) => {
|
||||
const original = (await importOriginal()) as typeof FsMod;
|
||||
const internalMockFsExistsSync = vi.fn();
|
||||
const internalMockFsReadFileSync = vi.fn();
|
||||
const internalMockFsUnlinkSync = vi.fn();
|
||||
return {
|
||||
...original,
|
||||
existsSync: internalMockFsExistsSync,
|
||||
readFileSync: internalMockFsReadFileSync,
|
||||
unlinkSync: internalMockFsUnlinkSync,
|
||||
};
|
||||
});
|
||||
|
||||
describe('useShellCommandProcessor', () => {
|
||||
// These spies will point to the vi.fn() instances created by the vi.mock('fs') factory.
|
||||
let existsSyncSpy: ReturnType<typeof vi.mocked<typeof existsSync>>;
|
||||
let readFileSyncSpy: ReturnType<typeof vi.mocked<typeof readFileSync>>;
|
||||
let unlinkSyncSpy: ReturnType<typeof vi.mocked<typeof unlinkSync>>;
|
||||
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;
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
mockExecuteCommand = vi.fn(); // Injected exec command mock
|
||||
beforeEach(async () => {
|
||||
const { spawn } = await import('child_process');
|
||||
spawnEmitter = new EventEmitter();
|
||||
spawnEmitter.stdout = new EventEmitter();
|
||||
spawnEmitter.stderr = new EventEmitter();
|
||||
(spawn as vi.Mock).mockReturnValue(spawnEmitter);
|
||||
|
||||
// Assign the imported (and mocked by vi.mock factory) fs functions to spies
|
||||
existsSyncSpy = existsSync as unknown as ReturnType<
|
||||
typeof vi.mocked<typeof existsSync>
|
||||
>;
|
||||
readFileSyncSpy = readFileSync as unknown as ReturnType<
|
||||
typeof vi.mocked<typeof readFileSync>
|
||||
>;
|
||||
unlinkSyncSpy = unlinkSync as unknown as ReturnType<
|
||||
typeof vi.mocked<typeof unlinkSync>
|
||||
>;
|
||||
vi.spyOn(fs, 'existsSync').mockReturnValue(false);
|
||||
vi.spyOn(fs, 'readFileSync').mockReturnValue('');
|
||||
vi.spyOn(fs, 'unlinkSync').mockReturnValue(undefined);
|
||||
|
||||
// It's important to reset these spies if they are to be configured per test
|
||||
existsSyncSpy.mockReset();
|
||||
readFileSyncSpy.mockReset();
|
||||
unlinkSyncSpy.mockReset();
|
||||
addItemToHistoryMock = vi.fn();
|
||||
setPendingHistoryItemMock = vi.fn();
|
||||
onExecMock = vi.fn();
|
||||
onDebugMessageMock = vi.fn();
|
||||
|
||||
mockGetTargetDir.mockReturnValue('/current/dir');
|
||||
configMock = {
|
||||
getTargetDir: () => '/test/dir',
|
||||
} as unknown as Config;
|
||||
|
||||
geminiClientMock = {
|
||||
addHistory: vi.fn(),
|
||||
} as unknown as GeminiClient;
|
||||
});
|
||||
|
||||
const setupHook = () =>
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks();
|
||||
});
|
||||
|
||||
const renderProcessorHook = () =>
|
||||
renderHook(() =>
|
||||
useShellCommandProcessor(
|
||||
mockAddItemToHistory,
|
||||
mockSetPendingHistoryItem,
|
||||
mockOnExec,
|
||||
mockOnDebugMessage,
|
||||
mockConfig,
|
||||
mockExecuteCommand as unknown as typeof ExecType, // Cast to satisfy TypeScript
|
||||
addItemToHistoryMock,
|
||||
setPendingHistoryItemMock,
|
||||
onExecMock,
|
||||
onDebugMessageMock,
|
||||
configMock,
|
||||
geminiClientMock,
|
||||
),
|
||||
);
|
||||
|
||||
it('should return false for non-string input', () => {
|
||||
const { result } = setupHook();
|
||||
const handled = result.current.handleShellCommand(
|
||||
[{ text: 'not a string' }] as PartListUnion,
|
||||
new AbortController().signal,
|
||||
);
|
||||
expect(handled).toBe(false);
|
||||
expect(mockAddItemToHistory).not.toHaveBeenCalled();
|
||||
});
|
||||
it('should execute a command and update history on success', async () => {
|
||||
const { result } = renderProcessorHook();
|
||||
const abortController = new AbortController();
|
||||
|
||||
it('should handle empty shell command', () => {
|
||||
const { result } = setupHook();
|
||||
act(() => {
|
||||
const handled = result.current.handleShellCommand(
|
||||
'',
|
||||
new AbortController().signal,
|
||||
result.current.handleShellCommand('ls -l', abortController.signal);
|
||||
});
|
||||
|
||||
expect(onExecMock).toHaveBeenCalledTimes(1);
|
||||
const execPromise = onExecMock.mock.calls[0][0];
|
||||
|
||||
// Simulate stdout
|
||||
act(() => {
|
||||
spawnEmitter.stdout.emit('data', Buffer.from('file1.txt\nfile2.txt'));
|
||||
});
|
||||
|
||||
// Simulate process exit
|
||||
act(() => {
|
||||
spawnEmitter.emit('exit', 0, null);
|
||||
});
|
||||
|
||||
await act(async () => {
|
||||
await execPromise;
|
||||
});
|
||||
|
||||
expect(addItemToHistoryMock).toHaveBeenCalledTimes(2);
|
||||
expect(addItemToHistoryMock.mock.calls[1][0]).toEqual({
|
||||
type: 'info',
|
||||
text: 'file1.txt\nfile2.txt',
|
||||
});
|
||||
expect(geminiClientMock.addHistory).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it('should handle binary output', 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,
|
||||
);
|
||||
expect(handled).toBe(true);
|
||||
});
|
||||
expect(mockAddItemToHistory).toHaveBeenNthCalledWith(
|
||||
1,
|
||||
{ type: 'user_shell', text: '' },
|
||||
expect.any(Number),
|
||||
);
|
||||
expect(mockAddItemToHistory).toHaveBeenNthCalledWith(
|
||||
2,
|
||||
{ type: 'error', text: 'Empty shell command.' },
|
||||
expect.any(Number),
|
||||
);
|
||||
expect(mockExecuteCommand).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should execute a successful command and add output to history', async () => {
|
||||
const { result } = setupHook();
|
||||
const command = '!ls -l';
|
||||
const stdout = 'total 0';
|
||||
const stderr = '';
|
||||
expect(onExecMock).toHaveBeenCalledTimes(1);
|
||||
const execPromise = onExecMock.mock.calls[0][0];
|
||||
|
||||
mockExecuteCommand.mockImplementation((_cmd, _options, callback) => {
|
||||
if (callback) callback(null, stdout, stderr);
|
||||
return {} as any;
|
||||
act(() => {
|
||||
spawnEmitter.stdout.emit('data', Buffer.from([0x89, 0x50, 0x4e, 0x47]));
|
||||
});
|
||||
|
||||
act(() => {
|
||||
spawnEmitter.emit('exit', 0, null);
|
||||
});
|
||||
existsSyncSpy.mockReturnValue(false);
|
||||
|
||||
await act(async () => {
|
||||
result.current.handleShellCommand(command, new AbortController().signal);
|
||||
await new Promise(process.nextTick);
|
||||
await execPromise;
|
||||
});
|
||||
|
||||
expect(mockAddItemToHistory).toHaveBeenNthCalledWith(
|
||||
1,
|
||||
{ type: 'user_shell', text: command },
|
||||
expect.any(Number),
|
||||
);
|
||||
expect(mockOnDebugMessage).toHaveBeenCalledWith(
|
||||
expect.stringContaining('Executing shell command in /current/dir:'),
|
||||
);
|
||||
expect(mockExecuteCommand).toHaveBeenCalledWith(
|
||||
'{ !ls -l; }; __code=$?; pwd >/tmp/shell_pwd_randomBytes.tmp; exit $__code',
|
||||
{ cwd: '/current/dir' },
|
||||
expect.any(Function),
|
||||
);
|
||||
expect(mockOnExec).toHaveBeenCalled();
|
||||
expect(mockAddItemToHistory).toHaveBeenNthCalledWith(
|
||||
2,
|
||||
{ type: 'info', text: stdout },
|
||||
expect.any(Number),
|
||||
);
|
||||
expect(addItemToHistoryMock).toHaveBeenCalledTimes(2);
|
||||
expect(addItemToHistoryMock.mock.calls[1][0]).toEqual({
|
||||
type: 'info',
|
||||
text: '[Command produced binary output, which is not shown.]',
|
||||
});
|
||||
});
|
||||
|
||||
it('should handle command with stderr output', async () => {
|
||||
const { result } = setupHook();
|
||||
const command = '!some_command_with_stderr';
|
||||
const stdout = 'some output';
|
||||
const stderr = 'some error output';
|
||||
it('should handle command failure', async () => {
|
||||
const { result } = renderProcessorHook();
|
||||
const abortController = new AbortController();
|
||||
|
||||
mockExecuteCommand.mockImplementation((_cmd, _options, callback) => {
|
||||
if (callback) callback(null, stdout, stderr);
|
||||
return {} as any;
|
||||
act(() => {
|
||||
result.current.handleShellCommand(
|
||||
'a-bad-command',
|
||||
abortController.signal,
|
||||
);
|
||||
});
|
||||
|
||||
const execPromise = onExecMock.mock.calls[0][0];
|
||||
|
||||
act(() => {
|
||||
spawnEmitter.stderr.emit('data', Buffer.from('command not found'));
|
||||
});
|
||||
|
||||
act(() => {
|
||||
spawnEmitter.emit('exit', 127, null);
|
||||
});
|
||||
existsSyncSpy.mockReturnValue(false);
|
||||
|
||||
await act(async () => {
|
||||
result.current.handleShellCommand(command, new AbortController().signal);
|
||||
await new Promise(process.nextTick);
|
||||
});
|
||||
expect(mockAddItemToHistory).toHaveBeenNthCalledWith(
|
||||
1,
|
||||
{ type: 'user_shell', text: command },
|
||||
expect.any(Number),
|
||||
);
|
||||
expect(mockAddItemToHistory).toHaveBeenNthCalledWith(
|
||||
2,
|
||||
{ type: 'info', text: `${stdout}\n${stderr}` },
|
||||
expect.any(Number),
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle command with only stderr output', async () => {
|
||||
const { result } = setupHook();
|
||||
const command = '!command_only_stderr';
|
||||
const stdout = '';
|
||||
const stderr = 'just stderr';
|
||||
|
||||
mockExecuteCommand.mockImplementation((_cmd, _options, callback) => {
|
||||
if (callback) callback(null, stdout, stderr);
|
||||
return {} as any;
|
||||
});
|
||||
existsSyncSpy.mockReturnValue(false);
|
||||
|
||||
await act(async () => {
|
||||
result.current.handleShellCommand(command, new AbortController().signal);
|
||||
await new Promise(process.nextTick);
|
||||
});
|
||||
expect(mockAddItemToHistory).toHaveBeenNthCalledWith(
|
||||
1,
|
||||
{ type: 'user_shell', text: command },
|
||||
expect.any(Number),
|
||||
);
|
||||
expect(mockAddItemToHistory).toHaveBeenNthCalledWith(
|
||||
2,
|
||||
{ type: 'info', text: stderr },
|
||||
expect.any(Number),
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle command error', async () => {
|
||||
const { result } = setupHook();
|
||||
const command = '!failing_command';
|
||||
const error = new Error('Command failed');
|
||||
|
||||
mockExecuteCommand.mockImplementation((_cmd, _options, callback) => {
|
||||
if (callback) callback(error, '', '');
|
||||
return {} as any;
|
||||
});
|
||||
existsSyncSpy.mockReturnValue(false);
|
||||
|
||||
await act(async () => {
|
||||
result.current.handleShellCommand(command, new AbortController().signal);
|
||||
await new Promise(process.nextTick);
|
||||
});
|
||||
expect(mockAddItemToHistory).toHaveBeenNthCalledWith(
|
||||
1,
|
||||
{ type: 'user_shell', text: command },
|
||||
expect.any(Number),
|
||||
);
|
||||
expect(mockAddItemToHistory).toHaveBeenNthCalledWith(
|
||||
2,
|
||||
{ type: 'error', text: error.message },
|
||||
expect.any(Number),
|
||||
);
|
||||
});
|
||||
|
||||
it('should correctly handle commands ending with &', async () => {
|
||||
const { result } = setupHook();
|
||||
const command = '!sleep 5 &';
|
||||
mockGetTargetDir.mockReturnValue('/current/dir');
|
||||
|
||||
mockExecuteCommand.mockImplementation((_cmd, _options, callback) => {
|
||||
if (callback) callback(null, '', '');
|
||||
return {} as any;
|
||||
});
|
||||
existsSyncSpy.mockReturnValue(false);
|
||||
|
||||
await act(async () => {
|
||||
result.current.handleShellCommand(command, new AbortController().signal);
|
||||
await new Promise(process.nextTick);
|
||||
await execPromise;
|
||||
});
|
||||
|
||||
expect(mockAddItemToHistory).toHaveBeenNthCalledWith(
|
||||
1,
|
||||
{ type: 'user_shell', text: command },
|
||||
expect.any(Number),
|
||||
);
|
||||
expect(mockExecuteCommand).toHaveBeenCalledWith(
|
||||
'{ !sleep 5 & }; __code=$?; pwd >/tmp/shell_pwd_randomBytes.tmp; exit $__code',
|
||||
{ cwd: '/current/dir' },
|
||||
expect.any(Function),
|
||||
);
|
||||
expect(mockAddItemToHistory).toHaveBeenNthCalledWith(
|
||||
2,
|
||||
{ type: 'info', text: '(Command produced no output)' },
|
||||
expect.any(Number),
|
||||
);
|
||||
expect(addItemToHistoryMock).toHaveBeenCalledTimes(2);
|
||||
expect(addItemToHistoryMock.mock.calls[1][0]).toEqual({
|
||||
type: 'error',
|
||||
text: 'Command exited with code 127.\ncommand not found',
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
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 targetDir = config.getTargetDir();
|
||||
onDebugMessage(
|
||||
`Executing shell command in ${targetDir}: ${commandToExecute}`,
|
||||
);
|
||||
const execOptions = {
|
||||
cwd: targetDir,
|
||||
};
|
||||
|
||||
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
|
||||
});
|
||||
|
||||
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();
|
||||
}
|
||||
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();
|
||||
}
|
||||
};
|
||||
child.stdout.on('data', handleOutput);
|
||||
child.stderr.on('data', handleOutput);
|
||||
|
||||
let error: Error | null = null;
|
||||
child.on('error', (err: Error) => {
|
||||
error = err;
|
||||
});
|
||||
|
||||
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}`,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
abortSignal.addEventListener('abort', abortHandler, { once: true });
|
||||
|
||||
child.on('exit', (code, signal) => {
|
||||
exited = true;
|
||||
abortSignal.removeEventListener('abort', abortHandler);
|
||||
},
|
||||
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);
|
||||
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);
|
||||
|
||||
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 {
|
||||
addItemToHistory(
|
||||
{ type: 'info', text: output + '\n' },
|
||||
userMessageTimestamp,
|
||||
);
|
||||
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 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,
|
||||
);
|
||||
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();
|
||||
});
|
||||
}
|
||||
});
|
||||
|
||||
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,
|
||||
],
|
||||
);
|
||||
|
||||
|
|
|
|||
|
|
@ -74,7 +74,7 @@ enum StreamProcessingStatus {
|
|||
* API interaction, and tool call lifecycle.
|
||||
*/
|
||||
export const useGeminiStream = (
|
||||
geminiClient: GeminiClient | null,
|
||||
geminiClient: GeminiClient,
|
||||
history: HistoryItem[],
|
||||
addItem: UseHistoryManagerReturn['addItem'],
|
||||
setShowHelp: React.Dispatch<React.SetStateAction<boolean>>,
|
||||
|
|
@ -140,6 +140,7 @@ export const useGeminiStream = (
|
|||
onExec,
|
||||
onDebugMessage,
|
||||
config,
|
||||
geminiClient,
|
||||
);
|
||||
|
||||
const streamingState = useMemo(() => {
|
||||
|
|
@ -474,13 +475,6 @@ export const useGeminiStream = (
|
|||
startNewTurn();
|
||||
}
|
||||
|
||||
if (!geminiClient) {
|
||||
const errorMsg = 'Gemini client is not available.';
|
||||
setInitError(errorMsg);
|
||||
addItem({ type: MessageType.ERROR, text: errorMsg }, Date.now());
|
||||
return;
|
||||
}
|
||||
|
||||
setIsResponding(true);
|
||||
setInitError(null);
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,41 @@
|
|||
/**
|
||||
* @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);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
@ -16,3 +16,32 @@ export const getAsciiArtWidth = (asciiArt: string): number => {
|
|||
const lines = asciiArt.split('\n');
|
||||
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 (let i = 0; i < sample.length; i++) {
|
||||
// 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 (sample[i] === 0) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
|
||||
// If no NULL bytes were found in the sample, we assume it's text.
|
||||
return false;
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in New Issue