mirror of
https://github.com/QwenLM/qwen-code.git
synced 2025-12-21 09:17:53 +00:00
Merge tag 'v0.1.15' into feature/yiheng/sync-gemini-cli-0.1.15
This commit is contained in:
@@ -0,0 +1,99 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import {
|
||||
DefaultArgumentProcessor,
|
||||
ShorthandArgumentProcessor,
|
||||
} from './argumentProcessor.js';
|
||||
import { createMockCommandContext } from '../../test-utils/mockCommandContext.js';
|
||||
|
||||
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();
|
||||
|
||||
it('should append the full command if args are provided', async () => {
|
||||
const prompt = 'Parse the command.';
|
||||
const context = createMockCommandContext({
|
||||
invocation: {
|
||||
raw: '/mycommand arg1 "arg two"',
|
||||
name: 'mycommand',
|
||||
args: 'arg1 "arg two"',
|
||||
},
|
||||
});
|
||||
const result = await processor.process(prompt, context);
|
||||
expect(result).toBe('Parse the command.\n\n/mycommand arg1 "arg two"');
|
||||
});
|
||||
|
||||
it('should NOT append the full command if no args are provided', async () => {
|
||||
const prompt = 'Parse the command.';
|
||||
const context = createMockCommandContext({
|
||||
invocation: {
|
||||
raw: '/mycommand',
|
||||
name: 'mycommand',
|
||||
args: '',
|
||||
},
|
||||
});
|
||||
const result = await processor.process(prompt, context);
|
||||
expect(result).toBe('Parse the command.');
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,34 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { IPromptProcessor, SHORTHAND_ARGS_PLACEHOLDER } 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.
|
||||
*/
|
||||
export class DefaultArgumentProcessor implements IPromptProcessor {
|
||||
async process(prompt: string, context: CommandContext): Promise<string> {
|
||||
if (context.invocation!.args) {
|
||||
return `${prompt}\n\n${context.invocation!.raw}`;
|
||||
}
|
||||
return prompt;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,300 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { vi, describe, it, expect, beforeEach } 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';
|
||||
|
||||
const mockCheckCommandPermissions = vi.hoisted(() => vi.fn());
|
||||
const mockShellExecute = vi.hoisted(() => vi.fn());
|
||||
|
||||
vi.mock('@qwen-code/qwen-code-core', async (importOriginal) => {
|
||||
const original = await importOriginal<object>();
|
||||
return {
|
||||
...original,
|
||||
checkCommandPermissions: mockCheckCommandPermissions,
|
||||
ShellExecutionService: {
|
||||
execute: mockShellExecute,
|
||||
},
|
||||
};
|
||||
});
|
||||
|
||||
describe('ShellProcessor', () => {
|
||||
let context: CommandContext;
|
||||
let mockConfig: Partial<Config>;
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
|
||||
mockConfig = {
|
||||
getTargetDir: vi.fn().mockReturnValue('/test/dir'),
|
||||
};
|
||||
|
||||
context = createMockCommandContext({
|
||||
services: {
|
||||
config: mockConfig as Config,
|
||||
},
|
||||
session: {
|
||||
sessionShellAllowlist: new Set(),
|
||||
},
|
||||
});
|
||||
|
||||
mockShellExecute.mockReturnValue({
|
||||
result: Promise.resolve({
|
||||
output: 'default shell output',
|
||||
}),
|
||||
});
|
||||
mockCheckCommandPermissions.mockReturnValue({
|
||||
allAllowed: true,
|
||||
disallowedCommands: [],
|
||||
});
|
||||
});
|
||||
|
||||
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.';
|
||||
const result = await processor.process(prompt, context);
|
||||
expect(result).toBe(prompt);
|
||||
expect(mockShellExecute).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should process a single valid shell injection if allowed', async () => {
|
||||
const processor = new ShellProcessor('test-command');
|
||||
const prompt = 'The current status is: !{git status}';
|
||||
mockCheckCommandPermissions.mockReturnValue({
|
||||
allAllowed: true,
|
||||
disallowedCommands: [],
|
||||
});
|
||||
mockShellExecute.mockReturnValue({
|
||||
result: Promise.resolve({ output: 'On branch main' }),
|
||||
});
|
||||
|
||||
const result = await processor.process(prompt, context);
|
||||
|
||||
expect(mockCheckCommandPermissions).toHaveBeenCalledWith(
|
||||
'git status',
|
||||
expect.any(Object),
|
||||
context.session.sessionShellAllowlist,
|
||||
);
|
||||
expect(mockShellExecute).toHaveBeenCalledWith(
|
||||
'git status',
|
||||
expect.any(String),
|
||||
expect.any(Function),
|
||||
expect.any(Object),
|
||||
);
|
||||
expect(result).toBe('The current status is: On branch main');
|
||||
});
|
||||
|
||||
it('should process multiple valid shell injections if all are allowed', async () => {
|
||||
const processor = new ShellProcessor('test-command');
|
||||
const prompt = '!{git status} in !{pwd}';
|
||||
mockCheckCommandPermissions.mockReturnValue({
|
||||
allAllowed: true,
|
||||
disallowedCommands: [],
|
||||
});
|
||||
|
||||
mockShellExecute
|
||||
.mockReturnValueOnce({
|
||||
result: Promise.resolve({ output: 'On branch main' }),
|
||||
})
|
||||
.mockReturnValueOnce({
|
||||
result: Promise.resolve({ output: '/usr/home' }),
|
||||
});
|
||||
|
||||
const result = await processor.process(prompt, context);
|
||||
|
||||
expect(mockCheckCommandPermissions).toHaveBeenCalledTimes(2);
|
||||
expect(mockShellExecute).toHaveBeenCalledTimes(2);
|
||||
expect(result).toBe('On branch main in /usr/home');
|
||||
});
|
||||
|
||||
it('should throw ConfirmationRequiredError if a command is not allowed', async () => {
|
||||
const processor = new ShellProcessor('test-command');
|
||||
const prompt = 'Do something dangerous: !{rm -rf /}';
|
||||
mockCheckCommandPermissions.mockReturnValue({
|
||||
allAllowed: false,
|
||||
disallowedCommands: ['rm -rf /'],
|
||||
});
|
||||
|
||||
await expect(processor.process(prompt, context)).rejects.toThrow(
|
||||
ConfirmationRequiredError,
|
||||
);
|
||||
});
|
||||
|
||||
it('should throw ConfirmationRequiredError with the correct command', async () => {
|
||||
const processor = new ShellProcessor('test-command');
|
||||
const prompt = 'Do something dangerous: !{rm -rf /}';
|
||||
mockCheckCommandPermissions.mockReturnValue({
|
||||
allAllowed: false,
|
||||
disallowedCommands: ['rm -rf /'],
|
||||
});
|
||||
|
||||
try {
|
||||
await processor.process(prompt, context);
|
||||
// Fail if it doesn't throw
|
||||
expect(true).toBe(false);
|
||||
} catch (e) {
|
||||
expect(e).toBeInstanceOf(ConfirmationRequiredError);
|
||||
if (e instanceof ConfirmationRequiredError) {
|
||||
expect(e.commandsToConfirm).toEqual(['rm -rf /']);
|
||||
}
|
||||
}
|
||||
|
||||
expect(mockShellExecute).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should throw ConfirmationRequiredError with multiple commands if multiple are disallowed', async () => {
|
||||
const processor = new ShellProcessor('test-command');
|
||||
const prompt = '!{cmd1} and !{cmd2}';
|
||||
mockCheckCommandPermissions.mockImplementation((cmd) => {
|
||||
if (cmd === 'cmd1') {
|
||||
return { allAllowed: false, disallowedCommands: ['cmd1'] };
|
||||
}
|
||||
if (cmd === 'cmd2') {
|
||||
return { allAllowed: false, disallowedCommands: ['cmd2'] };
|
||||
}
|
||||
return { allAllowed: true, disallowedCommands: [] };
|
||||
});
|
||||
|
||||
try {
|
||||
await processor.process(prompt, context);
|
||||
// Fail if it doesn't throw
|
||||
expect(true).toBe(false);
|
||||
} catch (e) {
|
||||
expect(e).toBeInstanceOf(ConfirmationRequiredError);
|
||||
if (e instanceof ConfirmationRequiredError) {
|
||||
expect(e.commandsToConfirm).toEqual(['cmd1', 'cmd2']);
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
it('should not execute any commands if at least one requires confirmation', async () => {
|
||||
const processor = new ShellProcessor('test-command');
|
||||
const prompt = 'First: !{echo "hello"}, Second: !{rm -rf /}';
|
||||
|
||||
mockCheckCommandPermissions.mockImplementation((cmd) => {
|
||||
if (cmd.includes('rm')) {
|
||||
return { allAllowed: false, disallowedCommands: [cmd] };
|
||||
}
|
||||
return { allAllowed: true, disallowedCommands: [] };
|
||||
});
|
||||
|
||||
await expect(processor.process(prompt, context)).rejects.toThrow(
|
||||
ConfirmationRequiredError,
|
||||
);
|
||||
|
||||
// Ensure no commands were executed because the pipeline was halted.
|
||||
expect(mockShellExecute).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should only request confirmation for disallowed commands in a mixed prompt', async () => {
|
||||
const processor = new ShellProcessor('test-command');
|
||||
const prompt = 'Allowed: !{ls -l}, Disallowed: !{rm -rf /}';
|
||||
|
||||
mockCheckCommandPermissions.mockImplementation((cmd) => ({
|
||||
allAllowed: !cmd.includes('rm'),
|
||||
disallowedCommands: cmd.includes('rm') ? [cmd] : [],
|
||||
}));
|
||||
|
||||
try {
|
||||
await processor.process(prompt, context);
|
||||
expect.fail('Should have thrown ConfirmationRequiredError');
|
||||
} catch (e) {
|
||||
expect(e).toBeInstanceOf(ConfirmationRequiredError);
|
||||
if (e instanceof ConfirmationRequiredError) {
|
||||
expect(e.commandsToConfirm).toEqual(['rm -rf /']);
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
it('should execute all commands if they are on the session allowlist', async () => {
|
||||
const processor = new ShellProcessor('test-command');
|
||||
const prompt = 'Run !{cmd1} and !{cmd2}';
|
||||
|
||||
// Add commands to the session allowlist
|
||||
context.session.sessionShellAllowlist = new Set(['cmd1', 'cmd2']);
|
||||
|
||||
// checkCommandPermissions should now pass for these
|
||||
mockCheckCommandPermissions.mockReturnValue({
|
||||
allAllowed: true,
|
||||
disallowedCommands: [],
|
||||
});
|
||||
|
||||
mockShellExecute
|
||||
.mockReturnValueOnce({ result: Promise.resolve({ output: 'output1' }) })
|
||||
.mockReturnValueOnce({ result: Promise.resolve({ output: 'output2' }) });
|
||||
|
||||
const result = await processor.process(prompt, context);
|
||||
|
||||
expect(mockCheckCommandPermissions).toHaveBeenCalledWith(
|
||||
'cmd1',
|
||||
expect.any(Object),
|
||||
context.session.sessionShellAllowlist,
|
||||
);
|
||||
expect(mockCheckCommandPermissions).toHaveBeenCalledWith(
|
||||
'cmd2',
|
||||
expect.any(Object),
|
||||
context.session.sessionShellAllowlist,
|
||||
);
|
||||
expect(mockShellExecute).toHaveBeenCalledTimes(2);
|
||||
expect(result).toBe('Run output1 and output2');
|
||||
});
|
||||
|
||||
it('should trim whitespace from the command inside the injection', async () => {
|
||||
const processor = new ShellProcessor('test-command');
|
||||
const prompt = 'Files: !{ ls -l }';
|
||||
mockCheckCommandPermissions.mockReturnValue({
|
||||
allAllowed: true,
|
||||
disallowedCommands: [],
|
||||
});
|
||||
mockShellExecute.mockReturnValue({
|
||||
result: Promise.resolve({ output: 'total 0' }),
|
||||
});
|
||||
|
||||
await processor.process(prompt, context);
|
||||
|
||||
expect(mockCheckCommandPermissions).toHaveBeenCalledWith(
|
||||
'ls -l', // Verifies that the command was trimmed
|
||||
expect.any(Object),
|
||||
context.session.sessionShellAllowlist,
|
||||
);
|
||||
expect(mockShellExecute).toHaveBeenCalledWith(
|
||||
'ls -l',
|
||||
expect.any(String),
|
||||
expect.any(Function),
|
||||
expect.any(Object),
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle an empty command inside the injection gracefully', 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');
|
||||
});
|
||||
});
|
||||
106
packages/cli/src/services/prompt-processors/shellProcessor.ts
Normal file
106
packages/cli/src/services/prompt-processors/shellProcessor.ts
Normal file
@@ -0,0 +1,106 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import {
|
||||
checkCommandPermissions,
|
||||
ShellExecutionService,
|
||||
} from '@qwen-code/qwen-code-core';
|
||||
|
||||
import { CommandContext } from '../../ui/commands/types.js';
|
||||
import { IPromptProcessor } from './types.js';
|
||||
|
||||
export class ConfirmationRequiredError extends Error {
|
||||
constructor(
|
||||
message: string,
|
||||
public commandsToConfirm: string[],
|
||||
) {
|
||||
super(message);
|
||||
this.name = 'ConfirmationRequiredError';
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Finds all instances of shell command injections (`!{...}`) in a prompt,
|
||||
* executes them, and replaces the injection site with the command's output.
|
||||
*
|
||||
* This processor ensures that only allowlisted commands are executed. If a
|
||||
* disallowed command is found, it halts execution and reports an error.
|
||||
*/
|
||||
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 matches = [...prompt.matchAll(ShellProcessor.SHELL_INJECTION_REGEX)];
|
||||
if (matches.length === 0) {
|
||||
return prompt; // No shell commands, nothing to do.
|
||||
}
|
||||
|
||||
// Discover all commands and check permissions.
|
||||
for (const match of matches) {
|
||||
const command = match[1].trim();
|
||||
const { allAllowed, disallowedCommands, blockReason, isHardDenial } =
|
||||
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.'}`,
|
||||
);
|
||||
}
|
||||
|
||||
// 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.
|
||||
if (commandsToConfirm.size > 0) {
|
||||
throw new ConfirmationRequiredError(
|
||||
'Shell command confirmation required',
|
||||
Array.from(commandsToConfirm),
|
||||
);
|
||||
}
|
||||
|
||||
// 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.
|
||||
);
|
||||
|
||||
const executionResult = await result;
|
||||
processedPrompt = processedPrompt.replace(
|
||||
fullMatch,
|
||||
executionResult.output,
|
||||
);
|
||||
}
|
||||
|
||||
return processedPrompt;
|
||||
}
|
||||
}
|
||||
42
packages/cli/src/services/prompt-processors/types.ts
Normal file
42
packages/cli/src/services/prompt-processors/types.ts
Normal file
@@ -0,0 +1,42 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { CommandContext } from '../../ui/commands/types.js';
|
||||
|
||||
/**
|
||||
* Defines the interface for a prompt processor, a module that can transform
|
||||
* a prompt string before it is sent to the model. Processors are chained
|
||||
* together to create a processing pipeline.
|
||||
*/
|
||||
export interface IPromptProcessor {
|
||||
/**
|
||||
* Processes a prompt string, applying a specific transformation as part of a pipeline.
|
||||
*
|
||||
* Each processor in a command's pipeline receives the output of the previous
|
||||
* processor. This method provides the full command context, allowing for
|
||||
* complex transformations that may require access to invocation details,
|
||||
* application services, or UI state.
|
||||
*
|
||||
* @param prompt The current state of the prompt string. This may have been
|
||||
* modified by previous processors in the pipeline.
|
||||
* @param context The full command context, providing access to invocation
|
||||
* details (like `context.invocation.raw` and `context.invocation.args`),
|
||||
* application services, and UI handlers.
|
||||
* @returns A promise that resolves to the transformed prompt string, which
|
||||
* will be passed to the next processor or, if it's the last one, sent to the model.
|
||||
*/
|
||||
process(prompt: string, context: CommandContext): Promise<string>;
|
||||
}
|
||||
|
||||
/**
|
||||
* The placeholder string for shorthand argument injection in custom commands.
|
||||
*/
|
||||
export const SHORTHAND_ARGS_PLACEHOLDER = '{{args}}';
|
||||
|
||||
/**
|
||||
* The trigger string for shell command injection in custom commands.
|
||||
*/
|
||||
export const SHELL_INJECTION_TRIGGER = '!{';
|
||||
Reference in New Issue
Block a user