This commit is contained in:
Coco Sheng 2026-04-21 16:41:54 +08:00 committed by GitHub
commit 589eb9bee7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 140 additions and 73 deletions

View file

@ -16,7 +16,7 @@ import {
afterEach,
type Mock,
} from 'vitest';
import { NoopSandboxManager } from '@google/gemini-cli-core';
import { NoopSandboxManager, escapeShellArg } from '@google/gemini-cli-core';
const mockIsBinary = vi.hoisted(() => vi.fn());
const mockShellExecutionService = vi.hoisted(() => vi.fn());
@ -76,7 +76,21 @@ vi.mock('@google/gemini-cli-core', async (importOriginal) => {
isBinary: mockIsBinary,
};
});
vi.mock('node:fs');
vi.mock('node:fs', async (importOriginal) => {
const actual = await importOriginal<typeof import('node:fs')>();
const mockFs = {
...actual,
existsSync: vi.fn(),
mkdtempSync: vi.fn(),
unlinkSync: vi.fn(),
readFileSync: vi.fn(),
rmSync: vi.fn(),
};
return {
...mockFs,
default: mockFs,
};
});
vi.mock('node:os', async (importOriginal) => {
const actual = await importOriginal<typeof import('node:os')>();
const mocked = {
@ -154,6 +168,7 @@ describe('useExecutionLifecycle', () => {
);
mockIsBinary.mockReturnValue(false);
vi.mocked(fs.existsSync).mockReturnValue(false);
vi.mocked(fs.mkdtempSync).mockReturnValue('/tmp/gemini-shell-abcdef');
mockShellExecutionService.mockImplementation((_cmd, _cwd, callback) => {
mockShellOutputCallback = callback;
@ -239,8 +254,9 @@ describe('useExecutionLifecycle', () => {
}),
],
});
const tmpFile = path.join(os.tmpdir(), 'shell_pwd_abcdef.tmp');
const wrappedCommand = `{ ls -l; }; __code=$?; pwd > "${tmpFile}"; exit $__code`;
const tmpFile = path.join('/tmp/gemini-shell-abcdef', 'pwd.tmp');
const escapedTmpFile = escapeShellArg(tmpFile, 'bash');
const wrappedCommand = `{\nls -l\n}\n__code=$?; pwd > ${escapedTmpFile}; exit $__code`;
expect(mockShellExecutionService).toHaveBeenCalledWith(
wrappedCommand,
'/test/dir',
@ -349,11 +365,9 @@ describe('useExecutionLifecycle', () => {
);
});
// Verify it's using the non-pty shell
const wrappedCommand = `{ stream; }; __code=$?; pwd > "${path.join(
os.tmpdir(),
'shell_pwd_abcdef.tmp',
)}"; exit $__code`;
const tmpFile = path.join('/tmp/gemini-shell-abcdef', 'pwd.tmp');
const escapedTmpFile = escapeShellArg(tmpFile, 'bash');
const wrappedCommand = `{\nstream\n}\n__code=$?; pwd > ${escapedTmpFile}; exit $__code`;
expect(mockShellExecutionService).toHaveBeenCalledWith(
wrappedCommand,
'/test/dir',
@ -644,7 +658,7 @@ describe('useExecutionLifecycle', () => {
type: 'error',
text: 'An unexpected error occurred: Synchronous spawn error',
});
const tmpFile = path.join(os.tmpdir(), 'shell_pwd_abcdef.tmp');
const tmpFile = path.join('/tmp/gemini-shell-abcdef', 'pwd.tmp');
// Verify that the temporary file was cleaned up
expect(vi.mocked(fs.unlinkSync)).toHaveBeenCalledWith(tmpFile);
expect(setShellInputFocusedMock).toHaveBeenCalledWith(false);
@ -652,7 +666,7 @@ describe('useExecutionLifecycle', () => {
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');
const tmpFile = path.join('/tmp/gemini-shell-abcdef', 'pwd.tmp');
vi.mocked(fs.existsSync).mockReturnValue(true);
vi.mocked(fs.readFileSync).mockReturnValue('/test/dir/new'); // A different directory

View file

@ -20,12 +20,12 @@ import {
ShellExecutionService,
ExecutionLifecycleService,
CoreToolCallStatus,
escapeShellArg,
} from '@google/gemini-cli-core';
import { type PartListUnion } from '@google/genai';
import type { UseHistoryManagerReturn } from './useHistoryManager.js';
import { SHELL_COMMAND_NAME } from '../constants.js';
import { formatBytes } from '../utils/formatters.js';
import crypto from 'node:crypto';
import path from 'node:path';
import os from 'node:os';
import fs from 'node:fs';
@ -362,18 +362,6 @@ export const useExecutionLifecycle = (
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 executeCommand = async () => {
let cumulativeStdout: string | AnsiOutput = '';
let isBinaryStream = false;
@ -403,9 +391,23 @@ export const useExecutionLifecycle = (
};
abortSignal.addEventListener('abort', abortHandler, { once: true });
onDebugMessage(`Executing in ${targetDir}: ${commandToExecute}`);
try {
// On non-windows, wrap the command to capture the final working directory.
if (!isWindows) {
let command = rawQuery.trim();
if (command.endsWith('\\')) {
command += ' ';
}
const tmpDir = fs.mkdtempSync(
path.join(os.tmpdir(), 'gemini-shell-'),
);
pwdFilePath = path.join(tmpDir, 'pwd.tmp');
const escapedPwdFilePath = escapeShellArg(pwdFilePath, 'bash');
commandToExecute = `{\n${command}\n}\n__code=$?; pwd > ${escapedPwdFilePath}; exit $__code`;
}
onDebugMessage(`Executing in ${targetDir}: ${commandToExecute}`);
const activeTheme = themeManager.getActiveTheme();
const shellExecutionConfig = {
...config.getShellExecutionConfig(),
@ -630,8 +632,18 @@ export const useExecutionLifecycle = (
);
} finally {
abortSignal.removeEventListener('abort', abortHandler);
if (pwdFilePath && fs.existsSync(pwdFilePath)) {
fs.unlinkSync(pwdFilePath);
if (pwdFilePath) {
const tmpDir = path.dirname(pwdFilePath);
try {
if (fs.existsSync(pwdFilePath)) {
fs.unlinkSync(pwdFilePath);
}
if (fs.existsSync(tmpDir)) {
fs.rmSync(tmpDir, { recursive: true, force: true });
}
} catch {
// Ignore cleanup errors
}
}
dispatch({ type: 'SET_ACTIVE_PTY', pid: null });

View file

@ -96,6 +96,7 @@ describe('ShellTool', () => {
let mockShellOutputCallback: (event: ShellOutputEvent) => void;
let resolveExecutionPromise: (result: ShellExecutionResult) => void;
let tempRootDir: string;
let extractedTmpFile: string;
beforeEach(() => {
vi.clearAllMocks();
@ -197,16 +198,28 @@ describe('ShellTool', () => {
process.env['ComSpec'] =
'C:\\Windows\\System32\\WindowsPowerShell\\v1.0\\powershell.exe';
extractedTmpFile = '';
// Capture the output callback to simulate streaming events from the service
mockShellExecutionService.mockImplementation((_cmd, _cwd, callback) => {
mockShellOutputCallback = callback;
return {
pid: 12345,
result: new Promise((resolve) => {
resolveExecutionPromise = resolve;
}),
};
});
mockShellExecutionService.mockImplementation(
(
cmd: string,
_cwd: string,
callback: (event: ShellOutputEvent) => void,
) => {
mockShellOutputCallback = callback;
const match = cmd.match(/pgrep -g 0 >([^ ]+)/);
if (match) {
extractedTmpFile = match[1].replace(/['"]/g, ''); // remove any quotes if present
}
return {
pid: 12345,
result: new Promise((resolve) => {
resolveExecutionPromise = resolve;
}),
};
},
);
mockShellBackground.mockImplementation(() => {
resolveExecutionPromise({
@ -293,17 +306,16 @@ describe('ShellTool', () => {
it('should wrap command on linux and parse pgrep output', async () => {
const invocation = shellTool.build({ command: 'my-command &' });
const promise = invocation.execute({ abortSignal: mockAbortSignal });
resolveShellExecution({ pid: 54321 });
// Simulate pgrep output file creation by the shell command
const tmpFile = path.join(os.tmpdir(), 'shell_pgrep_abcdef.tmp');
fs.writeFileSync(tmpFile, `54321${os.EOL}54322${os.EOL}`);
fs.writeFileSync(extractedTmpFile, `54321${os.EOL}54322${os.EOL}`);
resolveShellExecution({ pid: 54321 });
const result = await promise;
const wrappedCommand = `(\n${'my-command &'}\n); __code=$?; pgrep -g 0 >${tmpFile} 2>&1; exit $__code;`;
expect(mockShellExecutionService).toHaveBeenCalledWith(
wrappedCommand,
expect.stringMatching(/pgrep -g 0 >.*gemini-shell-.*[/\\]pgrep\.tmp/),
tempRootDir,
expect.any(Function),
expect.any(AbortSignal),
@ -316,7 +328,7 @@ describe('ShellTool', () => {
);
expect(result.llmContent).toContain('Background PIDs: 54322');
// The file should be deleted by the tool
expect(fs.existsSync(tmpFile)).toBe(false);
expect(fs.existsSync(extractedTmpFile)).toBe(false);
});
it('should add a space when command ends with a backslash to prevent escaping newline', async () => {
@ -325,10 +337,8 @@ describe('ShellTool', () => {
resolveShellExecution();
await promise;
const tmpFile = path.join(os.tmpdir(), 'shell_pgrep_abcdef.tmp');
const wrappedCommand = `(\nls\\ \n); __code=$?; pgrep -g 0 >${tmpFile} 2>&1; exit $__code;`;
expect(mockShellExecutionService).toHaveBeenCalledWith(
wrappedCommand,
expect.stringMatching(/pgrep -g 0 >.*gemini-shell-.*[/\\]pgrep\.tmp/),
tempRootDir,
expect.any(Function),
expect.any(AbortSignal),
@ -343,10 +353,8 @@ describe('ShellTool', () => {
resolveShellExecution();
await promise;
const tmpFile = path.join(os.tmpdir(), 'shell_pgrep_abcdef.tmp');
const wrappedCommand = `(\nls # comment\n); __code=$?; pgrep -g 0 >${tmpFile} 2>&1; exit $__code;`;
expect(mockShellExecutionService).toHaveBeenCalledWith(
wrappedCommand,
expect.stringMatching(/pgrep -g 0 >.*gemini-shell-.*[/\\]pgrep\.tmp/),
tempRootDir,
expect.any(Function),
expect.any(AbortSignal),
@ -365,10 +373,8 @@ describe('ShellTool', () => {
resolveShellExecution();
await promise;
const tmpFile = path.join(os.tmpdir(), 'shell_pgrep_abcdef.tmp');
const wrappedCommand = `(\n${'ls'}\n); __code=$?; pgrep -g 0 >${tmpFile} 2>&1; exit $__code;`;
expect(mockShellExecutionService).toHaveBeenCalledWith(
wrappedCommand,
expect.stringMatching(/pgrep -g 0 >.*gemini-shell-.*[/\\]pgrep\.tmp/),
subdir,
expect.any(Function),
expect.any(AbortSignal),
@ -390,10 +396,8 @@ describe('ShellTool', () => {
resolveShellExecution();
await promise;
const tmpFile = path.join(os.tmpdir(), 'shell_pgrep_abcdef.tmp');
const wrappedCommand = `(\n${'ls'}\n); __code=$?; pgrep -g 0 >${tmpFile} 2>&1; exit $__code;`;
expect(mockShellExecutionService).toHaveBeenCalledWith(
wrappedCommand,
expect.stringMatching(/pgrep -g 0 >.*gemini-shell-.*[/\\]pgrep\.tmp/),
path.join(tempRootDir, 'subdir'),
expect.any(Function),
expect.any(AbortSignal),
@ -462,6 +466,26 @@ describe('ShellTool', () => {
20000,
);
it('should correctly wrap heredoc commands', async () => {
const command = `cat << 'EOF'
hello world
EOF`;
const invocation = shellTool.build({ command });
const promise = invocation.execute({ abortSignal: mockAbortSignal });
resolveShellExecution();
await promise;
expect(mockShellExecutionService).toHaveBeenCalledWith(
expect.stringMatching(/pgrep -g 0 >.*gemini-shell-.*[/\\]pgrep\.tmp/),
tempRootDir,
expect.any(Function),
expect.any(AbortSignal),
false,
expect.any(Object),
);
expect(mockShellExecutionService.mock.calls[0][0]).toMatch(/\nEOF\n\)\n/);
});
it('should format error messages correctly', async () => {
const error = new Error('wrapped command failed');
const invocation = shellTool.build({ command: 'user-command' });
@ -562,10 +586,13 @@ describe('ShellTool', () => {
it('should clean up the temp file on synchronous execution error', async () => {
const error = new Error('sync spawn error');
mockShellExecutionService.mockImplementation(() => {
// Create the temp file before throwing to simulate it being left behind
const tmpFile = path.join(os.tmpdir(), 'shell_pgrep_abcdef.tmp');
fs.writeFileSync(tmpFile, '');
mockShellExecutionService.mockImplementation((cmd: string) => {
const match = cmd.match(/pgrep -g 0 >([^ ]+)/);
if (match) {
extractedTmpFile = match[1].replace(/['"]/g, ''); // remove any quotes if present
// Create the temp file before throwing to simulate it being left behind
fs.writeFileSync(extractedTmpFile, '');
}
throw error;
});
@ -574,8 +601,7 @@ describe('ShellTool', () => {
invocation.execute({ abortSignal: mockAbortSignal }),
).rejects.toThrow(error);
const tmpFile = path.join(os.tmpdir(), 'shell_pgrep_abcdef.tmp');
expect(fs.existsSync(tmpFile)).toBe(false);
expect(fs.existsSync(extractedTmpFile)).toBe(false);
});
it('should not log "missing pgrep output" when process is backgrounded', async () => {

View file

@ -8,7 +8,6 @@ import fsPromises from 'node:fs/promises';
import fs from 'node:fs';
import path from 'node:path';
import os from 'node:os';
import crypto from 'node:crypto';
import { debugLogger } from '../index.js';
import { type SandboxPermissions } from '../services/sandboxManager.js';
import { ToolErrorType } from './tool-error.js';
@ -42,6 +41,7 @@ import {
parseCommandDetails,
hasRedirection,
normalizeCommand,
escapeShellArg,
} from '../utils/shell-utils.js';
import { SHELL_TOOL_NAME } from './tool-names.js';
import { PARAM_ADDITIONAL_PERMISSIONS } from './definitions/base-declarations.js';
@ -111,7 +111,8 @@ export class ShellToolInvocation extends BaseToolInvocation<
if (trimmed.endsWith('\\')) {
trimmed += ' ';
}
return `(\n${trimmed}\n); __code=$?; pgrep -g 0 >${tempFilePath} 2>&1; exit $__code;`;
const escapedTempFilePath = escapeShellArg(tempFilePath, 'bash');
return `(\n${trimmed}\n)\n__code=$?; pgrep -g 0 >${escapedTempFilePath} 2>&1; exit $__code;`;
}
private getContextualDetails(): string {
@ -450,10 +451,8 @@ export class ShellToolInvocation extends BaseToolInvocation<
}
const isWindows = os.platform() === 'win32';
const tempFileName = `shell_pgrep_${crypto
.randomBytes(6)
.toString('hex')}.tmp`;
const tempFilePath = path.join(os.tmpdir(), tempFileName);
let tempFilePath = '';
let tempDir = '';
const timeoutMs = this.context.config.getShellToolInactivityTimeout();
const timeoutController = new AbortController();
@ -463,8 +462,12 @@ export class ShellToolInvocation extends BaseToolInvocation<
const combinedController = new AbortController();
const onAbort = () => combinedController.abort();
try {
if (!isWindows) {
tempDir = fs.mkdtempSync(path.join(os.tmpdir(), 'gemini-shell-'));
tempFilePath = path.join(tempDir, 'pgrep.tmp');
}
// pgrep is not available on Windows, so we can't get background PIDs
const commandToExecute = this.wrapCommandForPgrep(
strippedCommand,
@ -638,7 +641,10 @@ export class ShellToolInvocation extends BaseToolInvocation<
if (tempFileExists) {
const pgrepContent = await fsPromises.readFile(tempFilePath, 'utf8');
const pgrepLines = pgrepContent.split(os.EOL).filter(Boolean);
const pgrepLines = pgrepContent
.split('\n')
.map((line) => line.trim())
.filter(Boolean);
for (const line of pgrepLines) {
if (!/^\d+$/.test(line)) {
if (
@ -935,10 +941,19 @@ export class ShellToolInvocation extends BaseToolInvocation<
if (timeoutTimer) clearTimeout(timeoutTimer);
signal.removeEventListener('abort', onAbort);
timeoutController.signal.removeEventListener('abort', onAbort);
try {
await fsPromises.unlink(tempFilePath);
} catch {
// Ignore errors during unlink
if (tempFilePath) {
try {
await fsPromises.unlink(tempFilePath);
} catch {
// Ignore errors during unlink
}
}
if (tempDir) {
try {
await fsPromises.rm(tempDir, { recursive: true, force: true });
} catch {
// Ignore errors during rm
}
}
}
}