mirror of
https://github.com/QwenLM/qwen-code.git
synced 2025-12-21 09:17:53 +00:00
# 🚀 Sync Gemini CLI v0.2.1 - Major Feature Update (#483)
This commit is contained in:
@@ -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();
|
||||
|
||||
|
||||
@@ -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> {
|
||||
|
||||
@@ -4,11 +4,32 @@
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { vi, describe, it, expect, beforeEach } from 'vitest';
|
||||
import { describe, it, expect, beforeEach, vi, type Mock } 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 '@qwen-code/qwen-code-core';
|
||||
import { ApprovalMode, Config } from '@qwen-code/qwen-code-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,14 @@ vi.mock('@qwen-code/qwen-code-core', async (importOriginal) => {
|
||||
};
|
||||
});
|
||||
|
||||
const SUCCESS_RESULT = {
|
||||
output: 'default shell output',
|
||||
exitCode: 0,
|
||||
error: null,
|
||||
aborted: false,
|
||||
signal: null,
|
||||
};
|
||||
|
||||
describe('ShellProcessor', () => {
|
||||
let context: CommandContext;
|
||||
let mockConfig: Partial<Config>;
|
||||
@@ -33,9 +62,16 @@ describe('ShellProcessor', () => {
|
||||
|
||||
mockConfig = {
|
||||
getTargetDir: vi.fn().mockReturnValue('/test/dir'),
|
||||
getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
|
||||
getShouldUseNodePtyShell: vi.fn().mockReturnValue(false),
|
||||
};
|
||||
|
||||
context = createMockCommandContext({
|
||||
invocation: {
|
||||
raw: '/cmd default args',
|
||||
name: 'cmd',
|
||||
args: 'default args',
|
||||
},
|
||||
services: {
|
||||
config: mockConfig as Config,
|
||||
},
|
||||
@@ -45,16 +81,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 +120,7 @@ describe('ShellProcessor', () => {
|
||||
disallowedCommands: [],
|
||||
});
|
||||
mockShellExecute.mockReturnValue({
|
||||
result: Promise.resolve({ output: 'On branch main' }),
|
||||
result: Promise.resolve({ ...SUCCESS_RESULT, output: 'On branch main' }),
|
||||
});
|
||||
|
||||
const result = await processor.process(prompt, context);
|
||||
@@ -86,6 +135,7 @@ describe('ShellProcessor', () => {
|
||||
expect.any(String),
|
||||
expect.any(Function),
|
||||
expect.any(Object),
|
||||
false,
|
||||
);
|
||||
expect(result).toBe('The current status is: On branch main');
|
||||
});
|
||||
@@ -100,10 +150,13 @@ describe('ShellProcessor', () => {
|
||||
|
||||
mockShellExecute
|
||||
.mockReturnValueOnce({
|
||||
result: Promise.resolve({ output: 'On branch main' }),
|
||||
result: Promise.resolve({
|
||||
...SUCCESS_RESULT,
|
||||
output: 'On branch main',
|
||||
}),
|
||||
})
|
||||
.mockReturnValueOnce({
|
||||
result: Promise.resolve({ output: '/usr/home' }),
|
||||
result: Promise.resolve({ ...SUCCESS_RESULT, output: '/usr/home' }),
|
||||
});
|
||||
|
||||
const result = await processor.process(prompt, context);
|
||||
@@ -113,7 +166,7 @@ describe('ShellProcessor', () => {
|
||||
expect(result).toBe('On branch main in /usr/home');
|
||||
});
|
||||
|
||||
it('should throw ConfirmationRequiredError if a command is not allowed', async () => {
|
||||
it('should throw ConfirmationRequiredError if a command is not allowed in default mode', async () => {
|
||||
const processor = new ShellProcessor('test-command');
|
||||
const prompt = 'Do something dangerous: !{rm -rf /}';
|
||||
mockCheckCommandPermissions.mockReturnValue({
|
||||
@@ -126,6 +179,52 @@ describe('ShellProcessor', () => {
|
||||
);
|
||||
});
|
||||
|
||||
it('should NOT throw ConfirmationRequiredError if a command is not allowed but approval mode is YOLO', async () => {
|
||||
const processor = new ShellProcessor('test-command');
|
||||
const prompt = 'Do something dangerous: !{rm -rf /}';
|
||||
mockCheckCommandPermissions.mockReturnValue({
|
||||
allAllowed: false,
|
||||
disallowedCommands: ['rm -rf /'],
|
||||
});
|
||||
// Override the approval mode for this test
|
||||
(mockConfig.getApprovalMode as Mock).mockReturnValue(ApprovalMode.YOLO);
|
||||
mockShellExecute.mockReturnValue({
|
||||
result: Promise.resolve({ ...SUCCESS_RESULT, output: 'deleted' }),
|
||||
});
|
||||
|
||||
const result = await processor.process(prompt, context);
|
||||
|
||||
// It should proceed with execution
|
||||
expect(mockShellExecute).toHaveBeenCalledWith(
|
||||
'rm -rf /',
|
||||
expect.any(String),
|
||||
expect.any(Function),
|
||||
expect.any(Object),
|
||||
false,
|
||||
);
|
||||
expect(result).toBe('Do something dangerous: deleted');
|
||||
});
|
||||
|
||||
it('should still throw an error for a hard-denied command even in YOLO mode', async () => {
|
||||
const processor = new ShellProcessor('test-command');
|
||||
const prompt = 'Do something forbidden: !{reboot}';
|
||||
mockCheckCommandPermissions.mockReturnValue({
|
||||
allAllowed: false,
|
||||
disallowedCommands: ['reboot'],
|
||||
isHardDenial: true, // This is the key difference
|
||||
blockReason: 'System commands are blocked',
|
||||
});
|
||||
// Set approval mode to YOLO
|
||||
(mockConfig.getApprovalMode as Mock).mockReturnValue(ApprovalMode.YOLO);
|
||||
|
||||
await expect(processor.process(prompt, context)).rejects.toThrow(
|
||||
/Blocked command: "reboot". Reason: System commands are blocked/,
|
||||
);
|
||||
|
||||
// Ensure it never tried to execute
|
||||
expect(mockShellExecute).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should throw ConfirmationRequiredError with the correct command', async () => {
|
||||
const processor = new ShellProcessor('test-command');
|
||||
const prompt = 'Do something dangerous: !{rm -rf /}';
|
||||
@@ -226,8 +325,12 @@ describe('ShellProcessor', () => {
|
||||
});
|
||||
|
||||
mockShellExecute
|
||||
.mockReturnValueOnce({ result: Promise.resolve({ output: 'output1' }) })
|
||||
.mockReturnValueOnce({ result: Promise.resolve({ output: 'output2' }) });
|
||||
.mockReturnValueOnce({
|
||||
result: Promise.resolve({ ...SUCCESS_RESULT, output: 'output1' }),
|
||||
})
|
||||
.mockReturnValueOnce({
|
||||
result: Promise.resolve({ ...SUCCESS_RESULT, output: 'output2' }),
|
||||
});
|
||||
|
||||
const result = await processor.process(prompt, context);
|
||||
|
||||
@@ -245,56 +348,362 @@ 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, output: '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),
|
||||
false,
|
||||
);
|
||||
});
|
||||
|
||||
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: [],
|
||||
});
|
||||
mockShellExecute.mockReturnValue({
|
||||
result: Promise.resolve({ output: 'empty output' }),
|
||||
});
|
||||
|
||||
const result = await processor.process(prompt, context);
|
||||
|
||||
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');
|
||||
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, output: '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),
|
||||
false,
|
||||
);
|
||||
expect(result).toBe('Output: result');
|
||||
});
|
||||
|
||||
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, output: '{{a},{b}}' }),
|
||||
});
|
||||
|
||||
const result = await processor.process(prompt, context);
|
||||
expect(mockShellExecute).toHaveBeenCalledWith(
|
||||
command,
|
||||
expect.any(String),
|
||||
expect.any(Function),
|
||||
expect.any(Object),
|
||||
false,
|
||||
);
|
||||
expect(result).toBe('{{a},{b}}');
|
||||
});
|
||||
|
||||
it('should throw an error for unclosed shell injections', async () => {
|
||||
const processor = new ShellProcessor('test-command');
|
||||
const prompt = 'This prompt is broken: !{ls -l';
|
||||
|
||||
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 exit code and command name on failure', async () => {
|
||||
const processor = new ShellProcessor('test-command');
|
||||
const prompt = '!{cmd}';
|
||||
mockShellExecute.mockReturnValue({
|
||||
result: Promise.resolve({
|
||||
...SUCCESS_RESULT,
|
||||
output: 'some error output',
|
||||
stderr: '',
|
||||
exitCode: 1,
|
||||
}),
|
||||
});
|
||||
|
||||
const result = await processor.process(prompt, context);
|
||||
|
||||
expect(result).toBe(
|
||||
"some error output\n[Shell command 'cmd' 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,
|
||||
output: '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 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,
|
||||
output: '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, output: '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, output: '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),
|
||||
false,
|
||||
);
|
||||
|
||||
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, output: '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),
|
||||
false,
|
||||
);
|
||||
|
||||
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),
|
||||
false,
|
||||
);
|
||||
});
|
||||
|
||||
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),
|
||||
false,
|
||||
);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
@@ -5,12 +5,19 @@
|
||||
*/
|
||||
|
||||
import {
|
||||
ApprovalMode,
|
||||
checkCommandPermissions,
|
||||
escapeShellArg,
|
||||
getShellConfiguration,
|
||||
ShellExecutionService,
|
||||
} from '@qwen-code/qwen-code-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 +30,93 @@ 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));
|
||||
// If not a hard denial, respect YOLO mode and auto-approve.
|
||||
if (config.getApprovalMode() !== ApprovalMode.YOLO) {
|
||||
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 +124,125 @@ 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;
|
||||
|
||||
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,
|
||||
);
|
||||
|
||||
const executionResult = await result;
|
||||
processedPrompt = processedPrompt.replace(
|
||||
fullMatch,
|
||||
executionResult.output,
|
||||
);
|
||||
// Execute the resolved command (which already has ESCAPED input).
|
||||
if (injection.resolvedCommand) {
|
||||
const { result } = await ShellExecutionService.execute(
|
||||
injection.resolvedCommand,
|
||||
config.getTargetDir(),
|
||||
() => {},
|
||||
new AbortController().signal,
|
||||
config.getShouldUseNodePtyShell(),
|
||||
);
|
||||
|
||||
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.
|
||||
processedPrompt += executionResult.output;
|
||||
|
||||
// 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;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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}}';
|
||||
|
||||
|
||||
Reference in New Issue
Block a user