diff options
Diffstat (limited to 'packages/cli/src')
7 files changed, 768 insertions, 197 deletions
diff --git a/packages/cli/src/services/FileCommandLoader.test.ts b/packages/cli/src/services/FileCommandLoader.test.ts index f4f8ac2c..42d93074 100644 --- a/packages/cli/src/services/FileCommandLoader.test.ts +++ b/packages/cli/src/services/FileCommandLoader.test.ts @@ -22,7 +22,8 @@ import { ConfirmationRequiredError, ShellProcessor, } from './prompt-processors/shellProcessor.js'; -import { ShorthandArgumentProcessor } from './prompt-processors/argumentProcessor.js'; +import { DefaultArgumentProcessor } from './prompt-processors/argumentProcessor.js'; +import { CommandContext } from '../ui/commands/types.js'; const mockShellProcess = vi.hoisted(() => vi.fn()); vi.mock('./prompt-processors/shellProcessor.js', () => ({ @@ -46,9 +47,6 @@ vi.mock('./prompt-processors/argumentProcessor.js', async (importOriginal) => { typeof import('./prompt-processors/argumentProcessor.js') >(); return { - ShorthandArgumentProcessor: vi - .fn() - .mockImplementation(() => new original.ShorthandArgumentProcessor()), DefaultArgumentProcessor: vi .fn() .mockImplementation(() => new original.DefaultArgumentProcessor()), @@ -71,7 +69,16 @@ describe('FileCommandLoader', () => { beforeEach(() => { vi.clearAllMocks(); - mockShellProcess.mockImplementation((prompt) => Promise.resolve(prompt)); + mockShellProcess.mockImplementation( + (prompt: string, context: CommandContext) => { + const userArgsRaw = context?.invocation?.args || ''; + const processedPrompt = prompt.replaceAll( + SHORTHAND_ARGS_PLACEHOLDER, + userArgsRaw, + ); + return Promise.resolve(processedPrompt); + }, + ); }); afterEach(() => { @@ -198,7 +205,7 @@ describe('FileCommandLoader', () => { const mockConfig = { getProjectRoot: vi.fn(() => '/path/to/project'), getExtensions: vi.fn(() => []), - } as Config; + } as unknown as Config; const loader = new FileCommandLoader(mockConfig); const commands = await loader.loadCommands(signal); expect(commands).toHaveLength(1); @@ -239,7 +246,7 @@ describe('FileCommandLoader', () => { const mockConfig = { getProjectRoot: vi.fn(() => process.cwd()), getExtensions: vi.fn(() => []), - } as Config; + } as unknown as Config; const loader = new FileCommandLoader(mockConfig); const commands = await loader.loadCommands(signal); @@ -379,6 +386,68 @@ describe('FileCommandLoader', () => { expect(command.name).toBe('legacy_command'); }); + describe('Processor Instantiation Logic', () => { + it('instantiates only DefaultArgumentProcessor if no {{args}} or !{} are present', async () => { + const userCommandsDir = getUserCommandsDir(); + mock({ + [userCommandsDir]: { + 'simple.toml': `prompt = "Just a regular prompt"`, + }, + }); + + const loader = new FileCommandLoader(null as unknown as Config); + await loader.loadCommands(signal); + + expect(ShellProcessor).not.toHaveBeenCalled(); + expect(DefaultArgumentProcessor).toHaveBeenCalledTimes(1); + }); + + it('instantiates only ShellProcessor if {{args}} is present (but not !{})', async () => { + const userCommandsDir = getUserCommandsDir(); + mock({ + [userCommandsDir]: { + 'args.toml': `prompt = "Prompt with {{args}}"`, + }, + }); + + const loader = new FileCommandLoader(null as unknown as Config); + await loader.loadCommands(signal); + + expect(ShellProcessor).toHaveBeenCalledTimes(1); + expect(DefaultArgumentProcessor).not.toHaveBeenCalled(); + }); + + it('instantiates ShellProcessor and DefaultArgumentProcessor if !{} is present (but not {{args}})', async () => { + const userCommandsDir = getUserCommandsDir(); + mock({ + [userCommandsDir]: { + 'shell.toml': `prompt = "Prompt with !{cmd}"`, + }, + }); + + const loader = new FileCommandLoader(null as unknown as Config); + await loader.loadCommands(signal); + + expect(ShellProcessor).toHaveBeenCalledTimes(1); + expect(DefaultArgumentProcessor).toHaveBeenCalledTimes(1); + }); + + it('instantiates only ShellProcessor if both {{args}} and !{} are present', async () => { + const userCommandsDir = getUserCommandsDir(); + mock({ + [userCommandsDir]: { + 'both.toml': `prompt = "Prompt with {{args}} and !{cmd}"`, + }, + }); + + const loader = new FileCommandLoader(null as unknown as Config); + await loader.loadCommands(signal); + + expect(ShellProcessor).toHaveBeenCalledTimes(1); + expect(DefaultArgumentProcessor).not.toHaveBeenCalled(); + }); + }); + describe('Extension Command Loading', () => { it('loads commands from active extensions', async () => { const userCommandsDir = getUserCommandsDir(); @@ -416,7 +485,7 @@ describe('FileCommandLoader', () => { path: extensionDir, }, ]), - } as Config; + } as unknown as Config; const loader = new FileCommandLoader(mockConfig); const commands = await loader.loadCommands(signal); @@ -465,7 +534,7 @@ describe('FileCommandLoader', () => { path: extensionDir, }, ]), - } as Config; + } as unknown as Config; const loader = new FileCommandLoader(mockConfig); const commands = await loader.loadCommands(signal); @@ -572,7 +641,7 @@ describe('FileCommandLoader', () => { path: extensionDir2, }, ]), - } as Config; + } as unknown as Config; const loader = new FileCommandLoader(mockConfig); const commands = await loader.loadCommands(signal); @@ -608,7 +677,7 @@ describe('FileCommandLoader', () => { path: extensionDir, }, ]), - } as Config; + } as unknown as Config; const loader = new FileCommandLoader(mockConfig); const commands = await loader.loadCommands(signal); expect(commands).toHaveLength(0); @@ -640,7 +709,7 @@ describe('FileCommandLoader', () => { getExtensions: vi.fn(() => [ { name: 'a', version: '1.0.0', isActive: true, path: extensionDir }, ]), - } as Config; + } as unknown as Config; const loader = new FileCommandLoader(mockConfig); const commands = await loader.loadCommands(signal); @@ -671,7 +740,7 @@ describe('FileCommandLoader', () => { }); }); - describe('Shorthand Argument Processor Integration', () => { + describe('Argument Handling Integration (via ShellProcessor)', () => { it('correctly processes a command with {{args}}', async () => { const userCommandsDir = getUserCommandsDir(); mock({ @@ -738,6 +807,19 @@ describe('FileCommandLoader', () => { }); describe('Shell Processor Integration', () => { + it('instantiates ShellProcessor if {{args}} is present (even without shell trigger)', async () => { + const userCommandsDir = getUserCommandsDir(); + mock({ + [userCommandsDir]: { + 'args_only.toml': `prompt = "Hello {{args}}"`, + }, + }); + + const loader = new FileCommandLoader(null as unknown as Config); + await loader.loadCommands(signal); + + expect(ShellProcessor).toHaveBeenCalledWith('args_only'); + }); it('instantiates ShellProcessor if the trigger is present', async () => { const userCommandsDir = getUserCommandsDir(); mock({ @@ -752,7 +834,7 @@ describe('FileCommandLoader', () => { expect(ShellProcessor).toHaveBeenCalledWith('shell'); }); - it('does not instantiate ShellProcessor if trigger is missing', async () => { + it('does not instantiate ShellProcessor if no triggers ({{args}} or !{}) are present', async () => { const userCommandsDir = getUserCommandsDir(); mock({ [userCommandsDir]: { @@ -852,32 +934,30 @@ describe('FileCommandLoader', () => { ), ).rejects.toThrow('Something else went wrong'); }); - - it('assembles the processor pipeline in the correct order (Shell -> Argument)', async () => { + it('assembles the processor pipeline in the correct order (Shell -> Default)', async () => { const userCommandsDir = getUserCommandsDir(); mock({ [userCommandsDir]: { + // This prompt uses !{} but NOT {{args}}, so both processors should be active. 'pipeline.toml': ` - prompt = "Shell says: ${SHELL_INJECTION_TRIGGER}echo foo} and user says: ${SHORTHAND_ARGS_PLACEHOLDER}" - `, + prompt = "Shell says: ${SHELL_INJECTION_TRIGGER}echo foo}." + `, }, }); - // Mock the process methods to track call order - const argProcessMock = vi + const defaultProcessMock = vi .fn() - .mockImplementation((p) => `${p}-arg-processed`); + .mockImplementation((p) => Promise.resolve(`${p}-default-processed`)); - // Redefine the mock for this specific test mockShellProcess.mockImplementation((p) => Promise.resolve(`${p}-shell-processed`), ); - vi.mocked(ShorthandArgumentProcessor).mockImplementation( + vi.mocked(DefaultArgumentProcessor).mockImplementation( () => ({ - process: argProcessMock, - }) as unknown as ShorthandArgumentProcessor, + process: defaultProcessMock, + }) as unknown as DefaultArgumentProcessor, ); const loader = new FileCommandLoader(null as unknown as Config); @@ -885,7 +965,7 @@ describe('FileCommandLoader', () => { const command = commands.find((c) => c.name === 'pipeline'); expect(command).toBeDefined(); - await command!.action!( + const result = await command!.action!( createMockCommandContext({ invocation: { raw: '/pipeline bar', @@ -896,20 +976,27 @@ describe('FileCommandLoader', () => { 'bar', ); - // Verify that the shell processor was called before the argument processor expect(mockShellProcess.mock.invocationCallOrder[0]).toBeLessThan( - argProcessMock.mock.invocationCallOrder[0], + defaultProcessMock.mock.invocationCallOrder[0], ); - // Also verify the flow of the prompt through the processors + // Verify the flow of the prompt through the processors + // 1. Shell processor runs first expect(mockShellProcess).toHaveBeenCalledWith( - expect.any(String), + expect.stringContaining(SHELL_INJECTION_TRIGGER), expect.any(Object), ); - expect(argProcessMock).toHaveBeenCalledWith( - expect.stringContaining('-shell-processed'), // It receives the output of the shell processor + // 2. Default processor runs second + expect(defaultProcessMock).toHaveBeenCalledWith( + expect.stringContaining('-shell-processed'), expect.any(Object), ); + + if (result?.type === 'submit_prompt') { + expect(result.content).toContain('-shell-processed-default-processed'); + } else { + assert.fail('Incorrect action type'); + } }); }); }); diff --git a/packages/cli/src/services/FileCommandLoader.ts b/packages/cli/src/services/FileCommandLoader.ts index 5b8d8c42..2942a8b9 100644 --- a/packages/cli/src/services/FileCommandLoader.ts +++ b/packages/cli/src/services/FileCommandLoader.ts @@ -21,10 +21,7 @@ import { SlashCommand, SlashCommandActionReturn, } from '../ui/commands/types.js'; -import { - DefaultArgumentProcessor, - ShorthandArgumentProcessor, -} from './prompt-processors/argumentProcessor.js'; +import { DefaultArgumentProcessor } from './prompt-processors/argumentProcessor.js'; import { IPromptProcessor, SHORTHAND_ARGS_PLACEHOLDER, @@ -224,16 +221,21 @@ export class FileCommandLoader implements ICommandLoader { } const processors: IPromptProcessor[] = []; + const usesArgs = validDef.prompt.includes(SHORTHAND_ARGS_PLACEHOLDER); + const usesShellInjection = validDef.prompt.includes( + SHELL_INJECTION_TRIGGER, + ); - // Add the Shell Processor if needed. - if (validDef.prompt.includes(SHELL_INJECTION_TRIGGER)) { + // Interpolation (Shell Execution and Argument Injection) + // If the prompt uses either shell injection OR argument placeholders, + // we must use the ShellProcessor. + if (usesShellInjection || usesArgs) { processors.push(new ShellProcessor(baseCommandName)); } - // The presence of '{{args}}' is the switch that determines the behavior. - if (validDef.prompt.includes(SHORTHAND_ARGS_PLACEHOLDER)) { - processors.push(new ShorthandArgumentProcessor()); - } else { + // Default Argument Handling + // If NO explicit argument injection ({{args}}) was used, we append the raw invocation. + if (!usesArgs) { processors.push(new DefaultArgumentProcessor()); } diff --git a/packages/cli/src/services/prompt-processors/argumentProcessor.test.ts b/packages/cli/src/services/prompt-processors/argumentProcessor.test.ts index 6af578a9..1a4c0c6b 100644 --- a/packages/cli/src/services/prompt-processors/argumentProcessor.test.ts +++ b/packages/cli/src/services/prompt-processors/argumentProcessor.test.ts @@ -4,69 +4,11 @@ * SPDX-License-Identifier: Apache-2.0 */ -import { - DefaultArgumentProcessor, - ShorthandArgumentProcessor, -} from './argumentProcessor.js'; +import { DefaultArgumentProcessor } from './argumentProcessor.js'; import { createMockCommandContext } from '../../test-utils/mockCommandContext.js'; +import { describe, it, expect } from 'vitest'; describe('Argument Processors', () => { - describe('ShorthandArgumentProcessor', () => { - const processor = new ShorthandArgumentProcessor(); - - it('should replace a single {{args}} instance', async () => { - const prompt = 'Refactor the following code: {{args}}'; - const context = createMockCommandContext({ - invocation: { - raw: '/refactor make it faster', - name: 'refactor', - args: 'make it faster', - }, - }); - const result = await processor.process(prompt, context); - expect(result).toBe('Refactor the following code: make it faster'); - }); - - it('should replace multiple {{args}} instances', async () => { - const prompt = 'User said: {{args}}. I repeat: {{args}}!'; - const context = createMockCommandContext({ - invocation: { - raw: '/repeat hello world', - name: 'repeat', - args: 'hello world', - }, - }); - const result = await processor.process(prompt, context); - expect(result).toBe('User said: hello world. I repeat: hello world!'); - }); - - it('should handle an empty args string', async () => { - const prompt = 'The user provided no input: {{args}}.'; - const context = createMockCommandContext({ - invocation: { - raw: '/input', - name: 'input', - args: '', - }, - }); - const result = await processor.process(prompt, context); - expect(result).toBe('The user provided no input: .'); - }); - - it('should not change the prompt if {{args}} is not present', async () => { - const prompt = 'This is a static prompt.'; - const context = createMockCommandContext({ - invocation: { - raw: '/static some arguments', - name: 'static', - args: 'some arguments', - }, - }); - const result = await processor.process(prompt, context); - expect(result).toBe('This is a static prompt.'); - }); - }); - describe('DefaultArgumentProcessor', () => { const processor = new DefaultArgumentProcessor(); diff --git a/packages/cli/src/services/prompt-processors/argumentProcessor.ts b/packages/cli/src/services/prompt-processors/argumentProcessor.ts index a7efeea9..9d5fc369 100644 --- a/packages/cli/src/services/prompt-processors/argumentProcessor.ts +++ b/packages/cli/src/services/prompt-processors/argumentProcessor.ts @@ -4,25 +4,14 @@ * SPDX-License-Identifier: Apache-2.0 */ -import { IPromptProcessor, SHORTHAND_ARGS_PLACEHOLDER } from './types.js'; +import { IPromptProcessor } from './types.js'; import { CommandContext } from '../../ui/commands/types.js'; /** - * Replaces all instances of `{{args}}` in a prompt with the user-provided - * argument string. - */ -export class ShorthandArgumentProcessor implements IPromptProcessor { - async process(prompt: string, context: CommandContext): Promise<string> { - return prompt.replaceAll( - SHORTHAND_ARGS_PLACEHOLDER, - context.invocation!.args, - ); - } -} - -/** * Appends the user's full command invocation to the prompt if arguments are * provided, allowing the model to perform its own argument parsing. + * + * This processor is only used if the prompt does NOT contain {{args}}. */ export class DefaultArgumentProcessor implements IPromptProcessor { async process(prompt: string, context: CommandContext): Promise<string> { diff --git a/packages/cli/src/services/prompt-processors/shellProcessor.test.ts b/packages/cli/src/services/prompt-processors/shellProcessor.test.ts index a2883923..fe4ae7a2 100644 --- a/packages/cli/src/services/prompt-processors/shellProcessor.test.ts +++ b/packages/cli/src/services/prompt-processors/shellProcessor.test.ts @@ -4,11 +4,32 @@ * SPDX-License-Identifier: Apache-2.0 */ -import { vi, describe, it, expect, beforeEach } from 'vitest'; +import { describe, it, expect, beforeEach, vi } from 'vitest'; import { ConfirmationRequiredError, ShellProcessor } from './shellProcessor.js'; import { createMockCommandContext } from '../../test-utils/mockCommandContext.js'; import { CommandContext } from '../../ui/commands/types.js'; import { Config } from '@google/gemini-cli-core'; +import os from 'os'; +import { quote } from 'shell-quote'; + +// Helper function to determine the expected escaped string based on the current OS, +// mirroring the logic in the actual `escapeShellArg` implementation. This makes +// our tests robust and platform-agnostic. +function getExpectedEscapedArgForPlatform(arg: string): string { + if (os.platform() === 'win32') { + const comSpec = (process.env.ComSpec || 'cmd.exe').toLowerCase(); + const isPowerShell = + comSpec.endsWith('powershell.exe') || comSpec.endsWith('pwsh.exe'); + + if (isPowerShell) { + return `'${arg.replace(/'/g, "''")}'`; + } else { + return `"${arg.replace(/"/g, '""')}"`; + } + } else { + return quote([arg]); + } +} const mockCheckCommandPermissions = vi.hoisted(() => vi.fn()); const mockShellExecute = vi.hoisted(() => vi.fn()); @@ -24,6 +45,15 @@ vi.mock('@google/gemini-cli-core', async (importOriginal) => { }; }); +const SUCCESS_RESULT = { + stdout: 'default shell output', + stderr: '', + exitCode: 0, + error: null, + aborted: false, + signal: null, +}; + describe('ShellProcessor', () => { let context: CommandContext; let mockConfig: Partial<Config>; @@ -36,6 +66,11 @@ describe('ShellProcessor', () => { }; context = createMockCommandContext({ + invocation: { + raw: '/cmd default args', + name: 'cmd', + args: 'default args', + }, services: { config: mockConfig as Config, }, @@ -45,16 +80,29 @@ describe('ShellProcessor', () => { }); mockShellExecute.mockReturnValue({ - result: Promise.resolve({ - output: 'default shell output', - }), + result: Promise.resolve(SUCCESS_RESULT), }); + mockCheckCommandPermissions.mockReturnValue({ allAllowed: true, disallowedCommands: [], }); }); + it('should throw an error if config is missing', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = '!{ls}'; + const contextWithoutConfig = createMockCommandContext({ + services: { + config: null, + }, + }); + + await expect( + processor.process(prompt, contextWithoutConfig), + ).rejects.toThrow(/Security configuration not loaded/); + }); + it('should not change the prompt if no shell injections are present', async () => { const processor = new ShellProcessor('test-command'); const prompt = 'This is a simple prompt with no injections.'; @@ -71,7 +119,7 @@ describe('ShellProcessor', () => { disallowedCommands: [], }); mockShellExecute.mockReturnValue({ - result: Promise.resolve({ output: 'On branch main' }), + result: Promise.resolve({ ...SUCCESS_RESULT, stdout: 'On branch main' }), }); const result = await processor.process(prompt, context); @@ -100,10 +148,13 @@ describe('ShellProcessor', () => { mockShellExecute .mockReturnValueOnce({ - result: Promise.resolve({ output: 'On branch main' }), + result: Promise.resolve({ + ...SUCCESS_RESULT, + stdout: 'On branch main', + }), }) .mockReturnValueOnce({ - result: Promise.resolve({ output: '/usr/home' }), + result: Promise.resolve({ ...SUCCESS_RESULT, stdout: '/usr/home' }), }); const result = await processor.process(prompt, context); @@ -226,8 +277,12 @@ describe('ShellProcessor', () => { }); mockShellExecute - .mockReturnValueOnce({ result: Promise.resolve({ output: 'output1' }) }) - .mockReturnValueOnce({ result: Promise.resolve({ output: 'output2' }) }); + .mockReturnValueOnce({ + result: Promise.resolve({ ...SUCCESS_RESULT, stdout: 'output1' }), + }) + .mockReturnValueOnce({ + result: Promise.resolve({ ...SUCCESS_RESULT, stdout: 'output2' }), + }); const result = await processor.process(prompt, context); @@ -245,56 +300,406 @@ describe('ShellProcessor', () => { expect(result).toBe('Run output1 and output2'); }); - it('should trim whitespace from the command inside the injection', async () => { + it('should trim whitespace from the command inside the injection before interpolation', async () => { const processor = new ShellProcessor('test-command'); - const prompt = 'Files: !{ ls -l }'; + const prompt = 'Files: !{ ls {{args}} -l }'; + + const rawArgs = context.invocation!.args; + + const expectedEscapedArgs = getExpectedEscapedArgForPlatform(rawArgs); + + const expectedCommand = `ls ${expectedEscapedArgs} -l`; + mockCheckCommandPermissions.mockReturnValue({ allAllowed: true, disallowedCommands: [], }); mockShellExecute.mockReturnValue({ - result: Promise.resolve({ output: 'total 0' }), + result: Promise.resolve({ ...SUCCESS_RESULT, stdout: 'total 0' }), }); await processor.process(prompt, context); expect(mockCheckCommandPermissions).toHaveBeenCalledWith( - 'ls -l', // Verifies that the command was trimmed + expectedCommand, expect.any(Object), context.session.sessionShellAllowlist, ); expect(mockShellExecute).toHaveBeenCalledWith( - 'ls -l', + expectedCommand, expect.any(String), expect.any(Function), expect.any(Object), ); }); - it('should handle an empty command inside the injection gracefully', async () => { + it('should handle an empty command inside the injection gracefully (skips execution)', async () => { const processor = new ShellProcessor('test-command'); const prompt = 'This is weird: !{}'; - mockCheckCommandPermissions.mockReturnValue({ - allAllowed: true, - disallowedCommands: [], + + const result = await processor.process(prompt, context); + + expect(mockCheckCommandPermissions).not.toHaveBeenCalled(); + expect(mockShellExecute).not.toHaveBeenCalled(); + + // It replaces !{} with an empty string. + expect(result).toBe('This is weird: '); + }); + + describe('Robust Parsing (Balanced Braces)', () => { + it('should correctly parse commands containing nested braces (e.g., awk)', async () => { + const processor = new ShellProcessor('test-command'); + const command = "awk '{print $1}' file.txt"; + const prompt = `Output: !{${command}}`; + mockShellExecute.mockReturnValue({ + result: Promise.resolve({ ...SUCCESS_RESULT, stdout: 'result' }), + }); + + const result = await processor.process(prompt, context); + + expect(mockCheckCommandPermissions).toHaveBeenCalledWith( + command, + expect.any(Object), + context.session.sessionShellAllowlist, + ); + expect(mockShellExecute).toHaveBeenCalledWith( + command, + expect.any(String), + expect.any(Function), + expect.any(Object), + ); + expect(result).toBe('Output: result'); }); - mockShellExecute.mockReturnValue({ - result: Promise.resolve({ output: 'empty output' }), + + it('should handle deeply nested braces correctly', async () => { + const processor = new ShellProcessor('test-command'); + const command = "echo '{{a},{b}}'"; + const prompt = `!{${command}}`; + mockShellExecute.mockReturnValue({ + result: Promise.resolve({ ...SUCCESS_RESULT, stdout: '{{a},{b}}' }), + }); + + const result = await processor.process(prompt, context); + expect(mockShellExecute).toHaveBeenCalledWith( + command, + expect.any(String), + expect.any(Function), + expect.any(Object), + ); + expect(result).toBe('{{a},{b}}'); }); - const result = await processor.process(prompt, context); + it('should throw an error for unclosed shell injections', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = 'This prompt is broken: !{ls -l'; - expect(mockCheckCommandPermissions).toHaveBeenCalledWith( - '', - expect.any(Object), - context.session.sessionShellAllowlist, - ); - expect(mockShellExecute).toHaveBeenCalledWith( - '', - expect.any(String), - expect.any(Function), - expect.any(Object), - ); - expect(result).toBe('This is weird: empty output'); + await expect(processor.process(prompt, context)).rejects.toThrow( + /Unclosed shell injection/, + ); + }); + + it('should throw an error for unclosed nested braces', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = 'Broken: !{echo {a}'; + + await expect(processor.process(prompt, context)).rejects.toThrow( + /Unclosed shell injection/, + ); + }); + }); + + describe('Error Reporting', () => { + it('should append stderr information if the command produces it', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = '!{cmd}'; + mockShellExecute.mockReturnValue({ + result: Promise.resolve({ + ...SUCCESS_RESULT, + stdout: 'some output', + stderr: 'some error', + }), + }); + + const result = await processor.process(prompt, context); + + expect(result).toBe('some output\n--- STDERR ---\nsome error'); + }); + + it('should handle stderr-only output correctly', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = '!{cmd}'; + mockShellExecute.mockReturnValue({ + result: Promise.resolve({ + ...SUCCESS_RESULT, + stdout: '', + stderr: 'error only', + }), + }); + + const result = await processor.process(prompt, context); + + expect(result).toBe('--- STDERR ---\nerror only'); + }); + + it('should append exit code and command name if the command fails', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = 'Run a failing command: !{exit 1}'; + mockShellExecute.mockReturnValue({ + result: Promise.resolve({ + ...SUCCESS_RESULT, + stdout: 'some error output', + stderr: '', + exitCode: 1, + }), + }); + + const result = await processor.process(prompt, context); + + expect(result).toBe( + "Run a failing command: some error output\n[Shell command 'exit 1' exited with code 1]", + ); + }); + + it('should append signal info and command name if terminated by signal', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = '!{cmd}'; + mockShellExecute.mockReturnValue({ + result: Promise.resolve({ + ...SUCCESS_RESULT, + stdout: 'output', + stderr: '', + exitCode: null, + signal: 'SIGTERM', + }), + }); + + const result = await processor.process(prompt, context); + + expect(result).toBe( + "output\n[Shell command 'cmd' terminated by signal SIGTERM]", + ); + }); + + it('should append stderr and exit code information correctly', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = '!{cmd}'; + mockShellExecute.mockReturnValue({ + result: Promise.resolve({ + ...SUCCESS_RESULT, + stdout: 'out', + stderr: 'err', + exitCode: 127, + }), + }); + + const result = await processor.process(prompt, context); + + expect(result).toBe( + "out\n--- STDERR ---\nerr\n[Shell command 'cmd' exited with code 127]", + ); + }); + + it('should throw a detailed error if the shell fails to spawn', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = '!{bad-command}'; + const spawnError = new Error('spawn EACCES'); + mockShellExecute.mockReturnValue({ + result: Promise.resolve({ + ...SUCCESS_RESULT, + stdout: '', + stderr: '', + exitCode: null, + error: spawnError, + aborted: false, + }), + }); + + await expect(processor.process(prompt, context)).rejects.toThrow( + "Failed to start shell command in 'test-command': spawn EACCES. Command: bad-command", + ); + }); + + it('should report abort status with command name if aborted', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = '!{long-running-command}'; + const spawnError = new Error('Aborted'); + mockShellExecute.mockReturnValue({ + result: Promise.resolve({ + ...SUCCESS_RESULT, + stdout: 'partial output', + stderr: '', + exitCode: null, + error: spawnError, + aborted: true, // Key difference + }), + }); + + const result = await processor.process(prompt, context); + expect(result).toBe( + "partial output\n[Shell command 'long-running-command' aborted]", + ); + }); + }); + + describe('Context-Aware Argument Interpolation ({{args}})', () => { + const rawArgs = 'user input'; + + beforeEach(() => { + // Update context for these tests to use specific arguments + context.invocation!.args = rawArgs; + }); + + it('should perform raw replacement if no shell injections are present (optimization path)', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = 'The user said: {{args}}'; + + const result = await processor.process(prompt, context); + + expect(result).toBe(`The user said: ${rawArgs}`); + expect(mockShellExecute).not.toHaveBeenCalled(); + }); + + it('should perform raw replacement outside !{} blocks', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = 'Outside: {{args}}. Inside: !{echo "hello"}'; + mockShellExecute.mockReturnValue({ + result: Promise.resolve({ ...SUCCESS_RESULT, stdout: 'hello' }), + }); + + const result = await processor.process(prompt, context); + + expect(result).toBe(`Outside: ${rawArgs}. Inside: hello`); + }); + + it('should perform escaped replacement inside !{} blocks', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = 'Command: !{grep {{args}} file.txt}'; + mockShellExecute.mockReturnValue({ + result: Promise.resolve({ ...SUCCESS_RESULT, stdout: 'match found' }), + }); + + const result = await processor.process(prompt, context); + + const expectedEscapedArgs = getExpectedEscapedArgForPlatform(rawArgs); + const expectedCommand = `grep ${expectedEscapedArgs} file.txt`; + + expect(mockShellExecute).toHaveBeenCalledWith( + expectedCommand, + expect.any(String), + expect.any(Function), + expect.any(Object), + ); + + expect(result).toBe('Command: match found'); + }); + + it('should handle both raw (outside) and escaped (inside) injection simultaneously', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = 'User "({{args}})" requested search: !{search {{args}}}'; + mockShellExecute.mockReturnValue({ + result: Promise.resolve({ ...SUCCESS_RESULT, stdout: 'results' }), + }); + + const result = await processor.process(prompt, context); + + const expectedEscapedArgs = getExpectedEscapedArgForPlatform(rawArgs); + const expectedCommand = `search ${expectedEscapedArgs}`; + expect(mockShellExecute).toHaveBeenCalledWith( + expectedCommand, + expect.any(String), + expect.any(Function), + expect.any(Object), + ); + + expect(result).toBe(`User "(${rawArgs})" requested search: results`); + }); + + it('should perform security checks on the final, resolved (escaped) command', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = '!{rm {{args}}}'; + + const expectedEscapedArgs = getExpectedEscapedArgForPlatform(rawArgs); + const expectedResolvedCommand = `rm ${expectedEscapedArgs}`; + mockCheckCommandPermissions.mockReturnValue({ + allAllowed: false, + disallowedCommands: [expectedResolvedCommand], + isHardDenial: false, + }); + + await expect(processor.process(prompt, context)).rejects.toThrow( + ConfirmationRequiredError, + ); + + expect(mockCheckCommandPermissions).toHaveBeenCalledWith( + expectedResolvedCommand, + expect.any(Object), + context.session.sessionShellAllowlist, + ); + }); + + it('should report the resolved command if a hard denial occurs', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = '!{rm {{args}}}'; + const expectedEscapedArgs = getExpectedEscapedArgForPlatform(rawArgs); + const expectedResolvedCommand = `rm ${expectedEscapedArgs}`; + mockCheckCommandPermissions.mockReturnValue({ + allAllowed: false, + disallowedCommands: [expectedResolvedCommand], + isHardDenial: true, + blockReason: 'It is forbidden.', + }); + + await expect(processor.process(prompt, context)).rejects.toThrow( + `Blocked command: "${expectedResolvedCommand}". Reason: It is forbidden.`, + ); + }); + }); + describe('Real-World Escaping Scenarios', () => { + it('should correctly handle multiline arguments', async () => { + const processor = new ShellProcessor('test-command'); + const multilineArgs = 'first line\nsecond line'; + context.invocation!.args = multilineArgs; + const prompt = 'Commit message: !{git commit -m {{args}}}'; + + const expectedEscapedArgs = + getExpectedEscapedArgForPlatform(multilineArgs); + const expectedCommand = `git commit -m ${expectedEscapedArgs}`; + + await processor.process(prompt, context); + + expect(mockShellExecute).toHaveBeenCalledWith( + expectedCommand, + expect.any(String), + expect.any(Function), + expect.any(Object), + ); + }); + + it.each([ + { name: 'spaces', input: 'file with spaces.txt' }, + { name: 'double quotes', input: 'a "quoted" string' }, + { name: 'single quotes', input: "it's a string" }, + { name: 'command substitution (backticks)', input: '`reboot`' }, + { name: 'command substitution (dollar)', input: '$(reboot)' }, + { name: 'variable expansion', input: '$HOME' }, + { name: 'command chaining (semicolon)', input: 'a; reboot' }, + { name: 'command chaining (ampersand)', input: 'a && reboot' }, + ])('should safely escape args containing $name', async ({ input }) => { + const processor = new ShellProcessor('test-command'); + context.invocation!.args = input; + const prompt = '!{echo {{args}}}'; + + const expectedEscapedArgs = getExpectedEscapedArgForPlatform(input); + const expectedCommand = `echo ${expectedEscapedArgs}`; + + await processor.process(prompt, context); + + expect(mockShellExecute).toHaveBeenCalledWith( + expectedCommand, + expect.any(String), + expect.any(Function), + expect.any(Object), + ); + }); }); }); diff --git a/packages/cli/src/services/prompt-processors/shellProcessor.ts b/packages/cli/src/services/prompt-processors/shellProcessor.ts index bf811d66..9531d57e 100644 --- a/packages/cli/src/services/prompt-processors/shellProcessor.ts +++ b/packages/cli/src/services/prompt-processors/shellProcessor.ts @@ -6,11 +6,17 @@ import { checkCommandPermissions, + escapeShellArg, + getShellConfiguration, ShellExecutionService, } from '@google/gemini-cli-core'; import { CommandContext } from '../../ui/commands/types.js'; -import { IPromptProcessor } from './types.js'; +import { + IPromptProcessor, + SHELL_INJECTION_TRIGGER, + SHORTHAND_ARGS_PLACEHOLDER, +} from './types.js'; export class ConfirmationRequiredError extends Error { constructor( @@ -23,60 +29,89 @@ export class ConfirmationRequiredError extends Error { } /** - * Finds all instances of shell command injections (`!{...}`) in a prompt, - * executes them, and replaces the injection site with the command's output. + * Represents a single detected shell injection site in the prompt. + */ +interface ShellInjection { + /** The shell command extracted from within !{...}, trimmed. */ + command: string; + /** The starting index of the injection (inclusive, points to '!'). */ + startIndex: number; + /** The ending index of the injection (exclusive, points after '}'). */ + endIndex: number; + /** The command after {{args}} has been escaped and substituted. */ + resolvedCommand?: string; +} + +/** + * Handles prompt interpolation, including shell command execution (`!{...}`) + * and context-aware argument injection (`{{args}}`). * - * This processor ensures that only allowlisted commands are executed. If a - * disallowed command is found, it halts execution and reports an error. + * This processor ensures that: + * 1. `{{args}}` outside `!{...}` are replaced with raw input. + * 2. `{{args}}` inside `!{...}` are replaced with shell-escaped input. + * 3. Shell commands are executed securely after argument substitution. + * 4. Parsing correctly handles nested braces. */ export class ShellProcessor implements IPromptProcessor { - /** - * A regular expression to find all instances of `!{...}`. The inner - * capture group extracts the command itself. - */ - private static readonly SHELL_INJECTION_REGEX = /!\{([^}]*)\}/g; - - /** - * @param commandName The name of the custom command being executed, used - * for logging and error messages. - */ constructor(private readonly commandName: string) {} async process(prompt: string, context: CommandContext): Promise<string> { - const { config, sessionShellAllowlist } = { - ...context.services, - ...context.session, - }; - const commandsToExecute: Array<{ fullMatch: string; command: string }> = []; - const commandsToConfirm = new Set<string>(); + const userArgsRaw = context.invocation?.args || ''; - const matches = [...prompt.matchAll(ShellProcessor.SHELL_INJECTION_REGEX)]; - if (matches.length === 0) { - return prompt; // No shell commands, nothing to do. + if (!prompt.includes(SHELL_INJECTION_TRIGGER)) { + return prompt.replaceAll(SHORTHAND_ARGS_PLACEHOLDER, userArgsRaw); } - // Discover all commands and check permissions. - for (const match of matches) { - const command = match[1].trim(); + const config = context.services.config; + if (!config) { + throw new Error( + `Security configuration not loaded. Cannot verify shell command permissions for '${this.commandName}'. Aborting.`, + ); + } + const { sessionShellAllowlist } = context.session; + + const injections = this.extractInjections(prompt); + // If extractInjections found no closed blocks (and didn't throw), treat as raw. + if (injections.length === 0) { + return prompt.replaceAll(SHORTHAND_ARGS_PLACEHOLDER, userArgsRaw); + } + + const { shell } = getShellConfiguration(); + const userArgsEscaped = escapeShellArg(userArgsRaw, shell); + + const resolvedInjections = injections.map((injection) => { + if (injection.command === '') { + return injection; + } + // Replace {{args}} inside the command string with the escaped version. + const resolvedCommand = injection.command.replaceAll( + SHORTHAND_ARGS_PLACEHOLDER, + userArgsEscaped, + ); + return { ...injection, resolvedCommand }; + }); + + const commandsToConfirm = new Set<string>(); + for (const injection of resolvedInjections) { + const command = injection.resolvedCommand; + + if (!command) continue; + + // Security check on the final, escaped command string. const { allAllowed, disallowedCommands, blockReason, isHardDenial } = - checkCommandPermissions(command, config!, sessionShellAllowlist); + checkCommandPermissions(command, config, sessionShellAllowlist); if (!allAllowed) { - // If it's a hard denial, this is a non-recoverable security error. if (isHardDenial) { throw new Error( - `${this.commandName} cannot be run. ${blockReason || 'A shell command in this custom command is explicitly blocked in your config settings.'}`, + `${this.commandName} cannot be run. Blocked command: "${command}". Reason: ${blockReason || 'Blocked by configuration.'}`, ); } - - // Add each soft denial disallowed command to the set for confirmation. disallowedCommands.forEach((uc) => commandsToConfirm.add(uc)); } - commandsToExecute.push({ fullMatch: match[0], command }); } - // If any commands require confirmation, throw a special error to halt the - // pipeline and trigger the UI flow. + // Handle confirmation requirements. if (commandsToConfirm.size > 0) { throw new ConfirmationRequiredError( 'Shell command confirmation required', @@ -84,23 +119,132 @@ export class ShellProcessor implements IPromptProcessor { ); } - // Execute all commands (only runs if no confirmation was needed). - let processedPrompt = prompt; - for (const { fullMatch, command } of commandsToExecute) { - const { result } = ShellExecutionService.execute( - command, - config!.getTargetDir(), - () => {}, // No streaming needed. - new AbortController().signal, // For now, we don't support cancellation from here. - ); + let processedPrompt = ''; + let lastIndex = 0; - const executionResult = await result; - processedPrompt = processedPrompt.replace( - fullMatch, - executionResult.output, + for (const injection of resolvedInjections) { + // Append the text segment BEFORE the injection, substituting {{args}} with RAW input. + const segment = prompt.substring(lastIndex, injection.startIndex); + processedPrompt += segment.replaceAll( + SHORTHAND_ARGS_PLACEHOLDER, + userArgsRaw, ); + + // Execute the resolved command (which already has ESCAPED input). + if (injection.resolvedCommand) { + const { result } = ShellExecutionService.execute( + injection.resolvedCommand, + config.getTargetDir(), + () => {}, + new AbortController().signal, + ); + + const executionResult = await result; + + // Handle Spawn Errors + if (executionResult.error && !executionResult.aborted) { + throw new Error( + `Failed to start shell command in '${this.commandName}': ${executionResult.error.message}. Command: ${injection.resolvedCommand}`, + ); + } + + // Append the output, making stderr explicit for the model. + if (executionResult.stdout) { + processedPrompt += executionResult.stdout; + } + if (executionResult.stderr) { + if (executionResult.stdout) { + processedPrompt += '\n'; + } + processedPrompt += `--- STDERR ---\n${executionResult.stderr}`; + } + + // Append a status message if the command did not succeed. + if (executionResult.aborted) { + processedPrompt += `\n[Shell command '${injection.resolvedCommand}' aborted]`; + } else if ( + executionResult.exitCode !== 0 && + executionResult.exitCode !== null + ) { + processedPrompt += `\n[Shell command '${injection.resolvedCommand}' exited with code ${executionResult.exitCode}]`; + } else if (executionResult.signal !== null) { + processedPrompt += `\n[Shell command '${injection.resolvedCommand}' terminated by signal ${executionResult.signal}]`; + } + } + + lastIndex = injection.endIndex; } + // Append the remaining text AFTER the last injection, substituting {{args}} with RAW input. + const finalSegment = prompt.substring(lastIndex); + processedPrompt += finalSegment.replaceAll( + SHORTHAND_ARGS_PLACEHOLDER, + userArgsRaw, + ); + return processedPrompt; } + + /** + * Iteratively parses the prompt string to extract shell injections (!{...}), + * correctly handling nested braces within the command. + * + * @param prompt The prompt string to parse. + * @returns An array of extracted ShellInjection objects. + * @throws Error if an unclosed injection (`!{`) is found. + */ + private extractInjections(prompt: string): ShellInjection[] { + const injections: ShellInjection[] = []; + let index = 0; + + while (index < prompt.length) { + const startIndex = prompt.indexOf(SHELL_INJECTION_TRIGGER, index); + + if (startIndex === -1) { + break; + } + + let currentIndex = startIndex + SHELL_INJECTION_TRIGGER.length; + let braceCount = 1; + let foundEnd = false; + + while (currentIndex < prompt.length) { + const char = prompt[currentIndex]; + + // We count literal braces. This parser does not interpret shell quoting/escaping. + if (char === '{') { + braceCount++; + } else if (char === '}') { + braceCount--; + if (braceCount === 0) { + const commandContent = prompt.substring( + startIndex + SHELL_INJECTION_TRIGGER.length, + currentIndex, + ); + const endIndex = currentIndex + 1; + + injections.push({ + command: commandContent.trim(), + startIndex, + endIndex, + }); + + index = endIndex; + foundEnd = true; + break; + } + } + currentIndex++; + } + + // Check if the inner loop finished without finding the closing brace. + if (!foundEnd) { + throw new Error( + `Invalid syntax in command '${this.commandName}': Unclosed shell injection starting at index ${startIndex} ('!{'). Ensure braces are balanced.`, + ); + } + } + + return injections; + } } diff --git a/packages/cli/src/services/prompt-processors/types.ts b/packages/cli/src/services/prompt-processors/types.ts index 2653d2b7..956bb432 100644 --- a/packages/cli/src/services/prompt-processors/types.ts +++ b/packages/cli/src/services/prompt-processors/types.ts @@ -33,6 +33,8 @@ export interface IPromptProcessor { /** * The placeholder string for shorthand argument injection in custom commands. + * When used outside of !{...}, arguments are injected raw. + * When used inside !{...}, arguments are shell-escaped. */ export const SHORTHAND_ARGS_PLACEHOLDER = '{{args}}'; |