mirror of
https://github.com/QwenLM/qwen-code.git
synced 2025-12-21 17:27:54 +00:00
Merge tag 'v0.3.0' into chore/sync-gemini-cli-v0.3.0
This commit is contained in:
@@ -22,7 +22,7 @@ vi.mock('../ui/commands/restoreCommand.js', () => ({
|
||||
|
||||
import { describe, it, expect, vi, beforeEach, type Mock } from 'vitest';
|
||||
import { BuiltinCommandLoader } from './BuiltinCommandLoader.js';
|
||||
import { Config } from '@qwen-code/qwen-code-core';
|
||||
import type { Config } from '@qwen-code/qwen-code-core';
|
||||
import { CommandKind } from '../ui/commands/types.js';
|
||||
|
||||
import { ideCommand } from '../ui/commands/ideCommand.js';
|
||||
|
||||
@@ -4,9 +4,9 @@
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { ICommandLoader } from './types.js';
|
||||
import { SlashCommand } from '../ui/commands/types.js';
|
||||
import { Config } from '@qwen-code/qwen-code-core';
|
||||
import type { ICommandLoader } from './types.js';
|
||||
import type { SlashCommand } from '../ui/commands/types.js';
|
||||
import type { Config } from '@qwen-code/qwen-code-core';
|
||||
import { aboutCommand } from '../ui/commands/aboutCommand.js';
|
||||
import { authCommand } from '../ui/commands/authCommand.js';
|
||||
import { bugCommand } from '../ui/commands/bugCommand.js';
|
||||
|
||||
@@ -4,8 +4,8 @@
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { SlashCommand } from '../ui/commands/types.js';
|
||||
import { ICommandLoader } from './types.js';
|
||||
import type { SlashCommand } from '../ui/commands/types.js';
|
||||
import type { ICommandLoader } from './types.js';
|
||||
|
||||
/**
|
||||
* Orchestrates the discovery and loading of all slash commands for the CLI.
|
||||
|
||||
@@ -5,11 +5,8 @@
|
||||
*/
|
||||
|
||||
import * as path from 'node:path';
|
||||
import {
|
||||
Config,
|
||||
getProjectCommandsDir,
|
||||
getUserCommandsDir,
|
||||
} from '@qwen-code/qwen-code-core';
|
||||
import type { Config } from '@qwen-code/qwen-code-core';
|
||||
import { Storage } from '@qwen-code/qwen-code-core';
|
||||
import mock from 'mock-fs';
|
||||
import { FileCommandLoader } from './FileCommandLoader.js';
|
||||
import { assert, vi } from 'vitest';
|
||||
@@ -17,15 +14,23 @@ import { createMockCommandContext } from '../test-utils/mockCommandContext.js';
|
||||
import {
|
||||
SHELL_INJECTION_TRIGGER,
|
||||
SHORTHAND_ARGS_PLACEHOLDER,
|
||||
type PromptPipelineContent,
|
||||
} from './prompt-processors/types.js';
|
||||
import {
|
||||
ConfirmationRequiredError,
|
||||
ShellProcessor,
|
||||
} from './prompt-processors/shellProcessor.js';
|
||||
import { DefaultArgumentProcessor } from './prompt-processors/argumentProcessor.js';
|
||||
import { CommandContext } from '../ui/commands/types.js';
|
||||
import type { CommandContext } from '../ui/commands/types.js';
|
||||
import { AtFileProcessor } from './prompt-processors/atFileProcessor.js';
|
||||
|
||||
const mockShellProcess = vi.hoisted(() => vi.fn());
|
||||
const mockAtFileProcess = vi.hoisted(() => vi.fn());
|
||||
vi.mock('./prompt-processors/atFileProcessor.js', () => ({
|
||||
AtFileProcessor: vi.fn().mockImplementation(() => ({
|
||||
process: mockAtFileProcess,
|
||||
})),
|
||||
}));
|
||||
vi.mock('./prompt-processors/shellProcessor.js', () => ({
|
||||
ShellProcessor: vi.fn().mockImplementation(() => ({
|
||||
process: mockShellProcess,
|
||||
@@ -57,6 +62,7 @@ vi.mock('@qwen-code/qwen-code-core', async (importOriginal) => {
|
||||
await importOriginal<typeof import('@qwen-code/qwen-code-core')>();
|
||||
return {
|
||||
...original,
|
||||
Storage: original.Storage,
|
||||
isCommandAllowed: vi.fn(),
|
||||
ShellExecutionService: {
|
||||
execute: vi.fn(),
|
||||
@@ -70,15 +76,28 @@ describe('FileCommandLoader', () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
mockShellProcess.mockImplementation(
|
||||
(prompt: string, context: CommandContext) => {
|
||||
(prompt: PromptPipelineContent, context: CommandContext) => {
|
||||
const userArgsRaw = context?.invocation?.args || '';
|
||||
const processedPrompt = prompt.replaceAll(
|
||||
// This is a simplified mock. A real implementation would need to iterate
|
||||
// through all parts and process only the text parts.
|
||||
const firstTextPart = prompt.find(
|
||||
(p) => typeof p === 'string' || 'text' in p,
|
||||
);
|
||||
let textContent = '';
|
||||
if (typeof firstTextPart === 'string') {
|
||||
textContent = firstTextPart;
|
||||
} else if (firstTextPart && 'text' in firstTextPart) {
|
||||
textContent = firstTextPart.text ?? '';
|
||||
}
|
||||
|
||||
const processedText = textContent.replaceAll(
|
||||
SHORTHAND_ARGS_PLACEHOLDER,
|
||||
userArgsRaw,
|
||||
);
|
||||
return Promise.resolve(processedPrompt);
|
||||
return Promise.resolve([{ text: processedText }]);
|
||||
},
|
||||
);
|
||||
mockAtFileProcess.mockImplementation(async (prompt: string) => prompt);
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
@@ -86,7 +105,7 @@ describe('FileCommandLoader', () => {
|
||||
});
|
||||
|
||||
it('loads a single command from a file', async () => {
|
||||
const userCommandsDir = getUserCommandsDir();
|
||||
const userCommandsDir = Storage.getUserCommandsDir();
|
||||
mock({
|
||||
[userCommandsDir]: {
|
||||
'test.toml': 'prompt = "This is a test prompt"',
|
||||
@@ -112,7 +131,7 @@ describe('FileCommandLoader', () => {
|
||||
'',
|
||||
);
|
||||
if (result?.type === 'submit_prompt') {
|
||||
expect(result.content).toBe('This is a test prompt');
|
||||
expect(result.content).toEqual([{ text: 'This is a test prompt' }]);
|
||||
} else {
|
||||
assert.fail('Incorrect action type');
|
||||
}
|
||||
@@ -127,7 +146,7 @@ describe('FileCommandLoader', () => {
|
||||
itif(process.platform !== 'win32')(
|
||||
'loads commands from a symlinked directory',
|
||||
async () => {
|
||||
const userCommandsDir = getUserCommandsDir();
|
||||
const userCommandsDir = Storage.getUserCommandsDir();
|
||||
const realCommandsDir = '/real/commands';
|
||||
mock({
|
||||
[realCommandsDir]: {
|
||||
@@ -152,7 +171,7 @@ describe('FileCommandLoader', () => {
|
||||
itif(process.platform !== 'win32')(
|
||||
'loads commands from a symlinked subdirectory',
|
||||
async () => {
|
||||
const userCommandsDir = getUserCommandsDir();
|
||||
const userCommandsDir = Storage.getUserCommandsDir();
|
||||
const realNamespacedDir = '/real/namespaced-commands';
|
||||
mock({
|
||||
[userCommandsDir]: {
|
||||
@@ -176,7 +195,7 @@ describe('FileCommandLoader', () => {
|
||||
);
|
||||
|
||||
it('loads multiple commands', async () => {
|
||||
const userCommandsDir = getUserCommandsDir();
|
||||
const userCommandsDir = Storage.getUserCommandsDir();
|
||||
mock({
|
||||
[userCommandsDir]: {
|
||||
'test1.toml': 'prompt = "Prompt 1"',
|
||||
@@ -191,7 +210,7 @@ describe('FileCommandLoader', () => {
|
||||
});
|
||||
|
||||
it('creates deeply nested namespaces correctly', async () => {
|
||||
const userCommandsDir = getUserCommandsDir();
|
||||
const userCommandsDir = Storage.getUserCommandsDir();
|
||||
|
||||
mock({
|
||||
[userCommandsDir]: {
|
||||
@@ -213,7 +232,7 @@ describe('FileCommandLoader', () => {
|
||||
});
|
||||
|
||||
it('creates namespaces from nested directories', async () => {
|
||||
const userCommandsDir = getUserCommandsDir();
|
||||
const userCommandsDir = Storage.getUserCommandsDir();
|
||||
mock({
|
||||
[userCommandsDir]: {
|
||||
git: {
|
||||
@@ -232,8 +251,10 @@ describe('FileCommandLoader', () => {
|
||||
});
|
||||
|
||||
it('returns both user and project commands in order', async () => {
|
||||
const userCommandsDir = getUserCommandsDir();
|
||||
const projectCommandsDir = getProjectCommandsDir(process.cwd());
|
||||
const userCommandsDir = Storage.getUserCommandsDir();
|
||||
const projectCommandsDir = new Storage(
|
||||
process.cwd(),
|
||||
).getProjectCommandsDir();
|
||||
mock({
|
||||
[userCommandsDir]: {
|
||||
'test.toml': 'prompt = "User prompt"',
|
||||
@@ -262,7 +283,7 @@ describe('FileCommandLoader', () => {
|
||||
'',
|
||||
);
|
||||
if (userResult?.type === 'submit_prompt') {
|
||||
expect(userResult.content).toBe('User prompt');
|
||||
expect(userResult.content).toEqual([{ text: 'User prompt' }]);
|
||||
} else {
|
||||
assert.fail('Incorrect action type for user command');
|
||||
}
|
||||
@@ -277,14 +298,14 @@ describe('FileCommandLoader', () => {
|
||||
'',
|
||||
);
|
||||
if (projectResult?.type === 'submit_prompt') {
|
||||
expect(projectResult.content).toBe('Project prompt');
|
||||
expect(projectResult.content).toEqual([{ text: 'Project prompt' }]);
|
||||
} else {
|
||||
assert.fail('Incorrect action type for project command');
|
||||
}
|
||||
});
|
||||
|
||||
it('ignores files with TOML syntax errors', async () => {
|
||||
const userCommandsDir = getUserCommandsDir();
|
||||
const userCommandsDir = Storage.getUserCommandsDir();
|
||||
mock({
|
||||
[userCommandsDir]: {
|
||||
'invalid.toml': 'this is not valid toml',
|
||||
@@ -300,7 +321,7 @@ describe('FileCommandLoader', () => {
|
||||
});
|
||||
|
||||
it('ignores files that are semantically invalid (missing prompt)', async () => {
|
||||
const userCommandsDir = getUserCommandsDir();
|
||||
const userCommandsDir = Storage.getUserCommandsDir();
|
||||
mock({
|
||||
[userCommandsDir]: {
|
||||
'no_prompt.toml': 'description = "This file is missing a prompt"',
|
||||
@@ -316,7 +337,7 @@ describe('FileCommandLoader', () => {
|
||||
});
|
||||
|
||||
it('handles filename edge cases correctly', async () => {
|
||||
const userCommandsDir = getUserCommandsDir();
|
||||
const userCommandsDir = Storage.getUserCommandsDir();
|
||||
mock({
|
||||
[userCommandsDir]: {
|
||||
'test.v1.toml': 'prompt = "Test prompt"',
|
||||
@@ -338,7 +359,7 @@ describe('FileCommandLoader', () => {
|
||||
});
|
||||
|
||||
it('uses a default description if not provided', async () => {
|
||||
const userCommandsDir = getUserCommandsDir();
|
||||
const userCommandsDir = Storage.getUserCommandsDir();
|
||||
mock({
|
||||
[userCommandsDir]: {
|
||||
'test.toml': 'prompt = "Test prompt"',
|
||||
@@ -353,7 +374,7 @@ describe('FileCommandLoader', () => {
|
||||
});
|
||||
|
||||
it('uses the provided description', async () => {
|
||||
const userCommandsDir = getUserCommandsDir();
|
||||
const userCommandsDir = Storage.getUserCommandsDir();
|
||||
mock({
|
||||
[userCommandsDir]: {
|
||||
'test.toml': 'prompt = "Test prompt"\ndescription = "My test command"',
|
||||
@@ -368,7 +389,7 @@ describe('FileCommandLoader', () => {
|
||||
});
|
||||
|
||||
it('should sanitize colons in filenames to prevent namespace conflicts', async () => {
|
||||
const userCommandsDir = getUserCommandsDir();
|
||||
const userCommandsDir = Storage.getUserCommandsDir();
|
||||
mock({
|
||||
[userCommandsDir]: {
|
||||
'legacy:command.toml': 'prompt = "This is a legacy command"',
|
||||
@@ -388,7 +409,7 @@ describe('FileCommandLoader', () => {
|
||||
|
||||
describe('Processor Instantiation Logic', () => {
|
||||
it('instantiates only DefaultArgumentProcessor if no {{args}} or !{} are present', async () => {
|
||||
const userCommandsDir = getUserCommandsDir();
|
||||
const userCommandsDir = Storage.getUserCommandsDir();
|
||||
mock({
|
||||
[userCommandsDir]: {
|
||||
'simple.toml': `prompt = "Just a regular prompt"`,
|
||||
@@ -403,7 +424,7 @@ describe('FileCommandLoader', () => {
|
||||
});
|
||||
|
||||
it('instantiates only ShellProcessor if {{args}} is present (but not !{})', async () => {
|
||||
const userCommandsDir = getUserCommandsDir();
|
||||
const userCommandsDir = Storage.getUserCommandsDir();
|
||||
mock({
|
||||
[userCommandsDir]: {
|
||||
'args.toml': `prompt = "Prompt with {{args}}"`,
|
||||
@@ -418,7 +439,7 @@ describe('FileCommandLoader', () => {
|
||||
});
|
||||
|
||||
it('instantiates ShellProcessor and DefaultArgumentProcessor if !{} is present (but not {{args}})', async () => {
|
||||
const userCommandsDir = getUserCommandsDir();
|
||||
const userCommandsDir = Storage.getUserCommandsDir();
|
||||
mock({
|
||||
[userCommandsDir]: {
|
||||
'shell.toml': `prompt = "Prompt with !{cmd}"`,
|
||||
@@ -433,7 +454,7 @@ describe('FileCommandLoader', () => {
|
||||
});
|
||||
|
||||
it('instantiates only ShellProcessor if both {{args}} and !{} are present', async () => {
|
||||
const userCommandsDir = getUserCommandsDir();
|
||||
const userCommandsDir = Storage.getUserCommandsDir();
|
||||
mock({
|
||||
[userCommandsDir]: {
|
||||
'both.toml': `prompt = "Prompt with {{args}} and !{cmd}"`,
|
||||
@@ -446,12 +467,62 @@ describe('FileCommandLoader', () => {
|
||||
expect(ShellProcessor).toHaveBeenCalledTimes(1);
|
||||
expect(DefaultArgumentProcessor).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('instantiates AtFileProcessor and DefaultArgumentProcessor if @{} is present', async () => {
|
||||
const userCommandsDir = Storage.getUserCommandsDir();
|
||||
mock({
|
||||
[userCommandsDir]: {
|
||||
'at-file.toml': `prompt = "Context: @{./my-file.txt}"`,
|
||||
},
|
||||
});
|
||||
|
||||
const loader = new FileCommandLoader(null as unknown as Config);
|
||||
await loader.loadCommands(signal);
|
||||
|
||||
expect(AtFileProcessor).toHaveBeenCalledTimes(1);
|
||||
expect(ShellProcessor).not.toHaveBeenCalled();
|
||||
expect(DefaultArgumentProcessor).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it('instantiates ShellProcessor and AtFileProcessor if !{} and @{} are present', async () => {
|
||||
const userCommandsDir = Storage.getUserCommandsDir();
|
||||
mock({
|
||||
[userCommandsDir]: {
|
||||
'shell-and-at.toml': `prompt = "Run !{cmd} with @{file.txt}"`,
|
||||
},
|
||||
});
|
||||
|
||||
const loader = new FileCommandLoader(null as unknown as Config);
|
||||
await loader.loadCommands(signal);
|
||||
|
||||
expect(ShellProcessor).toHaveBeenCalledTimes(1);
|
||||
expect(AtFileProcessor).toHaveBeenCalledTimes(1);
|
||||
expect(DefaultArgumentProcessor).toHaveBeenCalledTimes(1); // because no {{args}}
|
||||
});
|
||||
|
||||
it('instantiates only ShellProcessor and AtFileProcessor if {{args}} and @{} are present', async () => {
|
||||
const userCommandsDir = Storage.getUserCommandsDir();
|
||||
mock({
|
||||
[userCommandsDir]: {
|
||||
'args-and-at.toml': `prompt = "Run {{args}} with @{file.txt}"`,
|
||||
},
|
||||
});
|
||||
|
||||
const loader = new FileCommandLoader(null as unknown as Config);
|
||||
await loader.loadCommands(signal);
|
||||
|
||||
expect(ShellProcessor).toHaveBeenCalledTimes(1);
|
||||
expect(AtFileProcessor).toHaveBeenCalledTimes(1);
|
||||
expect(DefaultArgumentProcessor).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
describe('Extension Command Loading', () => {
|
||||
it('loads commands from active extensions', async () => {
|
||||
const userCommandsDir = getUserCommandsDir();
|
||||
const projectCommandsDir = getProjectCommandsDir(process.cwd());
|
||||
const userCommandsDir = Storage.getUserCommandsDir();
|
||||
const projectCommandsDir = new Storage(
|
||||
process.cwd(),
|
||||
).getProjectCommandsDir();
|
||||
const extensionDir = path.join(
|
||||
process.cwd(),
|
||||
'.gemini/extensions/test-ext',
|
||||
@@ -499,8 +570,10 @@ describe('FileCommandLoader', () => {
|
||||
});
|
||||
|
||||
it('extension commands have extensionName metadata for conflict resolution', async () => {
|
||||
const userCommandsDir = getUserCommandsDir();
|
||||
const projectCommandsDir = getProjectCommandsDir(process.cwd());
|
||||
const userCommandsDir = Storage.getUserCommandsDir();
|
||||
const projectCommandsDir = new Storage(
|
||||
process.cwd(),
|
||||
).getProjectCommandsDir();
|
||||
const extensionDir = path.join(
|
||||
process.cwd(),
|
||||
'.gemini/extensions/test-ext',
|
||||
@@ -555,7 +628,7 @@ describe('FileCommandLoader', () => {
|
||||
);
|
||||
expect(result0?.type).toBe('submit_prompt');
|
||||
if (result0?.type === 'submit_prompt') {
|
||||
expect(result0.content).toBe('User deploy command');
|
||||
expect(result0.content).toEqual([{ text: 'User deploy command' }]);
|
||||
}
|
||||
|
||||
expect(commands[1].name).toBe('deploy');
|
||||
@@ -572,7 +645,7 @@ describe('FileCommandLoader', () => {
|
||||
);
|
||||
expect(result1?.type).toBe('submit_prompt');
|
||||
if (result1?.type === 'submit_prompt') {
|
||||
expect(result1.content).toBe('Project deploy command');
|
||||
expect(result1.content).toEqual([{ text: 'Project deploy command' }]);
|
||||
}
|
||||
|
||||
expect(commands[2].name).toBe('deploy');
|
||||
@@ -590,7 +663,7 @@ describe('FileCommandLoader', () => {
|
||||
);
|
||||
expect(result2?.type).toBe('submit_prompt');
|
||||
if (result2?.type === 'submit_prompt') {
|
||||
expect(result2.content).toBe('Extension deploy command');
|
||||
expect(result2.content).toEqual([{ text: 'Extension deploy command' }]);
|
||||
}
|
||||
});
|
||||
|
||||
@@ -733,7 +806,9 @@ describe('FileCommandLoader', () => {
|
||||
'',
|
||||
);
|
||||
if (result?.type === 'submit_prompt') {
|
||||
expect(result.content).toBe('Nested command from extension a');
|
||||
expect(result.content).toEqual([
|
||||
{ text: 'Nested command from extension a' },
|
||||
]);
|
||||
} else {
|
||||
assert.fail('Incorrect action type');
|
||||
}
|
||||
@@ -742,7 +817,7 @@ describe('FileCommandLoader', () => {
|
||||
|
||||
describe('Argument Handling Integration (via ShellProcessor)', () => {
|
||||
it('correctly processes a command with {{args}}', async () => {
|
||||
const userCommandsDir = getUserCommandsDir();
|
||||
const userCommandsDir = Storage.getUserCommandsDir();
|
||||
mock({
|
||||
[userCommandsDir]: {
|
||||
'shorthand.toml':
|
||||
@@ -767,14 +842,16 @@ describe('FileCommandLoader', () => {
|
||||
);
|
||||
expect(result?.type).toBe('submit_prompt');
|
||||
if (result?.type === 'submit_prompt') {
|
||||
expect(result.content).toBe('The user wants to: do something cool');
|
||||
expect(result.content).toEqual([
|
||||
{ text: 'The user wants to: do something cool' },
|
||||
]);
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
describe('Default Argument Processor Integration', () => {
|
||||
it('correctly processes a command without {{args}}', async () => {
|
||||
const userCommandsDir = getUserCommandsDir();
|
||||
const userCommandsDir = Storage.getUserCommandsDir();
|
||||
mock({
|
||||
[userCommandsDir]: {
|
||||
'model_led.toml':
|
||||
@@ -801,14 +878,14 @@ describe('FileCommandLoader', () => {
|
||||
if (result?.type === 'submit_prompt') {
|
||||
const expectedContent =
|
||||
'This is the instruction.\n\n/model_led 1.2.0 added "a feature"';
|
||||
expect(result.content).toBe(expectedContent);
|
||||
expect(result.content).toEqual([{ text: expectedContent }]);
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
describe('Shell Processor Integration', () => {
|
||||
it('instantiates ShellProcessor if {{args}} is present (even without shell trigger)', async () => {
|
||||
const userCommandsDir = getUserCommandsDir();
|
||||
const userCommandsDir = Storage.getUserCommandsDir();
|
||||
mock({
|
||||
[userCommandsDir]: {
|
||||
'args_only.toml': `prompt = "Hello {{args}}"`,
|
||||
@@ -821,7 +898,7 @@ describe('FileCommandLoader', () => {
|
||||
expect(ShellProcessor).toHaveBeenCalledWith('args_only');
|
||||
});
|
||||
it('instantiates ShellProcessor if the trigger is present', async () => {
|
||||
const userCommandsDir = getUserCommandsDir();
|
||||
const userCommandsDir = Storage.getUserCommandsDir();
|
||||
mock({
|
||||
[userCommandsDir]: {
|
||||
'shell.toml': `prompt = "Run this: ${SHELL_INJECTION_TRIGGER}echo hello}"`,
|
||||
@@ -835,7 +912,7 @@ describe('FileCommandLoader', () => {
|
||||
});
|
||||
|
||||
it('does not instantiate ShellProcessor if no triggers ({{args}} or !{}) are present', async () => {
|
||||
const userCommandsDir = getUserCommandsDir();
|
||||
const userCommandsDir = Storage.getUserCommandsDir();
|
||||
mock({
|
||||
[userCommandsDir]: {
|
||||
'regular.toml': `prompt = "Just a regular prompt"`,
|
||||
@@ -849,13 +926,13 @@ describe('FileCommandLoader', () => {
|
||||
});
|
||||
|
||||
it('returns a "submit_prompt" action if shell processing succeeds', async () => {
|
||||
const userCommandsDir = getUserCommandsDir();
|
||||
const userCommandsDir = Storage.getUserCommandsDir();
|
||||
mock({
|
||||
[userCommandsDir]: {
|
||||
'shell.toml': `prompt = "Run !{echo 'hello'}"`,
|
||||
},
|
||||
});
|
||||
mockShellProcess.mockResolvedValue('Run hello');
|
||||
mockShellProcess.mockResolvedValue([{ text: 'Run hello' }]);
|
||||
|
||||
const loader = new FileCommandLoader(null as unknown as Config);
|
||||
const commands = await loader.loadCommands(signal);
|
||||
@@ -871,12 +948,12 @@ describe('FileCommandLoader', () => {
|
||||
|
||||
expect(result?.type).toBe('submit_prompt');
|
||||
if (result?.type === 'submit_prompt') {
|
||||
expect(result.content).toBe('Run hello');
|
||||
expect(result.content).toEqual([{ text: 'Run hello' }]);
|
||||
}
|
||||
});
|
||||
|
||||
it('returns a "confirm_shell_commands" action if shell processing requires it', async () => {
|
||||
const userCommandsDir = getUserCommandsDir();
|
||||
const userCommandsDir = Storage.getUserCommandsDir();
|
||||
const rawInvocation = '/shell rm -rf /';
|
||||
mock({
|
||||
[userCommandsDir]: {
|
||||
@@ -910,7 +987,7 @@ describe('FileCommandLoader', () => {
|
||||
});
|
||||
|
||||
it('re-throws other errors from the processor', async () => {
|
||||
const userCommandsDir = getUserCommandsDir();
|
||||
const userCommandsDir = Storage.getUserCommandsDir();
|
||||
mock({
|
||||
[userCommandsDir]: {
|
||||
'shell.toml': `prompt = "Run !{something}"`,
|
||||
@@ -934,23 +1011,36 @@ describe('FileCommandLoader', () => {
|
||||
),
|
||||
).rejects.toThrow('Something else went wrong');
|
||||
});
|
||||
it('assembles the processor pipeline in the correct order (Shell -> Default)', async () => {
|
||||
const userCommandsDir = getUserCommandsDir();
|
||||
it('assembles the processor pipeline in the correct order (AtFile -> Shell -> Default)', async () => {
|
||||
const userCommandsDir = Storage.getUserCommandsDir();
|
||||
mock({
|
||||
[userCommandsDir]: {
|
||||
// This prompt uses !{} but NOT {{args}}, so both processors should be active.
|
||||
// This prompt uses !{}, @{}, but NOT {{args}}, so all processors should be active.
|
||||
'pipeline.toml': `
|
||||
prompt = "Shell says: ${SHELL_INJECTION_TRIGGER}echo foo}."
|
||||
prompt = "Shell says: !{echo foo}. File says: @{./bar.txt}"
|
||||
`,
|
||||
},
|
||||
'./bar.txt': 'bar content',
|
||||
});
|
||||
|
||||
const defaultProcessMock = vi
|
||||
.fn()
|
||||
.mockImplementation((p) => Promise.resolve(`${p}-default-processed`));
|
||||
.mockImplementation((p: PromptPipelineContent) =>
|
||||
Promise.resolve([
|
||||
{ text: `${(p[0] as { text: string }).text}-default-processed` },
|
||||
]),
|
||||
);
|
||||
|
||||
mockShellProcess.mockImplementation((p) =>
|
||||
Promise.resolve(`${p}-shell-processed`),
|
||||
mockShellProcess.mockImplementation((p: PromptPipelineContent) =>
|
||||
Promise.resolve([
|
||||
{ text: `${(p[0] as { text: string }).text}-shell-processed` },
|
||||
]),
|
||||
);
|
||||
|
||||
mockAtFileProcess.mockImplementation((p: PromptPipelineContent) =>
|
||||
Promise.resolve([
|
||||
{ text: `${(p[0] as { text: string }).text}-at-file-processed` },
|
||||
]),
|
||||
);
|
||||
|
||||
vi.mocked(DefaultArgumentProcessor).mockImplementation(
|
||||
@@ -968,35 +1058,115 @@ describe('FileCommandLoader', () => {
|
||||
const result = await command!.action!(
|
||||
createMockCommandContext({
|
||||
invocation: {
|
||||
raw: '/pipeline bar',
|
||||
raw: '/pipeline baz',
|
||||
name: 'pipeline',
|
||||
args: 'bar',
|
||||
args: 'baz',
|
||||
},
|
||||
}),
|
||||
'bar',
|
||||
'baz',
|
||||
);
|
||||
|
||||
expect(mockAtFileProcess.mock.invocationCallOrder[0]).toBeLessThan(
|
||||
mockShellProcess.mock.invocationCallOrder[0],
|
||||
);
|
||||
expect(mockShellProcess.mock.invocationCallOrder[0]).toBeLessThan(
|
||||
defaultProcessMock.mock.invocationCallOrder[0],
|
||||
);
|
||||
|
||||
// Verify the flow of the prompt through the processors
|
||||
// 1. Shell processor runs first
|
||||
expect(mockShellProcess).toHaveBeenCalledWith(
|
||||
expect.stringContaining(SHELL_INJECTION_TRIGGER),
|
||||
// 1. AtFile processor runs first
|
||||
expect(mockAtFileProcess).toHaveBeenCalledWith(
|
||||
[{ text: expect.stringContaining('@{./bar.txt}') }],
|
||||
expect.any(Object),
|
||||
);
|
||||
// 2. Default processor runs second
|
||||
// 2. Shell processor runs second
|
||||
expect(mockShellProcess).toHaveBeenCalledWith(
|
||||
[{ text: expect.stringContaining('-at-file-processed') }],
|
||||
expect.any(Object),
|
||||
);
|
||||
// 3. Default processor runs third
|
||||
expect(defaultProcessMock).toHaveBeenCalledWith(
|
||||
expect.stringContaining('-shell-processed'),
|
||||
[{ text: expect.stringContaining('-shell-processed') }],
|
||||
expect.any(Object),
|
||||
);
|
||||
|
||||
if (result?.type === 'submit_prompt') {
|
||||
expect(result.content).toContain('-shell-processed-default-processed');
|
||||
const contentAsArray = Array.isArray(result.content)
|
||||
? result.content
|
||||
: [result.content];
|
||||
expect(contentAsArray.length).toBeGreaterThan(0);
|
||||
const firstPart = contentAsArray[0];
|
||||
|
||||
if (typeof firstPart === 'object' && firstPart && 'text' in firstPart) {
|
||||
expect(firstPart.text).toContain(
|
||||
'-at-file-processed-shell-processed-default-processed',
|
||||
);
|
||||
} else {
|
||||
assert.fail(
|
||||
'First part of content is not a text part or is a string',
|
||||
);
|
||||
}
|
||||
} else {
|
||||
assert.fail('Incorrect action type');
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
describe('@-file Processor Integration', () => {
|
||||
it('correctly processes a command with @{file}', async () => {
|
||||
const userCommandsDir = Storage.getUserCommandsDir();
|
||||
mock({
|
||||
[userCommandsDir]: {
|
||||
'at-file.toml':
|
||||
'prompt = "Context from file: @{./test.txt}"\ndescription = "@-file test"',
|
||||
},
|
||||
'./test.txt': 'file content',
|
||||
});
|
||||
|
||||
mockAtFileProcess.mockImplementation(
|
||||
async (prompt: PromptPipelineContent) => {
|
||||
// A simplified mock of AtFileProcessor's behavior
|
||||
const textContent = (prompt[0] as { text: string }).text;
|
||||
if (textContent.includes('@{./test.txt}')) {
|
||||
return [
|
||||
{
|
||||
text: textContent.replace('@{./test.txt}', 'file content'),
|
||||
},
|
||||
];
|
||||
}
|
||||
return prompt;
|
||||
},
|
||||
);
|
||||
|
||||
// Prevent default processor from interfering
|
||||
vi.mocked(DefaultArgumentProcessor).mockImplementation(
|
||||
() =>
|
||||
({
|
||||
process: (p: PromptPipelineContent) => Promise.resolve(p),
|
||||
}) as unknown as DefaultArgumentProcessor,
|
||||
);
|
||||
|
||||
const loader = new FileCommandLoader(null as unknown as Config);
|
||||
const commands = await loader.loadCommands(signal);
|
||||
const command = commands.find((c) => c.name === 'at-file');
|
||||
expect(command).toBeDefined();
|
||||
|
||||
const result = await command!.action?.(
|
||||
createMockCommandContext({
|
||||
invocation: {
|
||||
raw: '/at-file',
|
||||
name: 'at-file',
|
||||
args: '',
|
||||
},
|
||||
}),
|
||||
'',
|
||||
);
|
||||
expect(result?.type).toBe('submit_prompt');
|
||||
if (result?.type === 'submit_prompt') {
|
||||
expect(result.content).toEqual([
|
||||
{ text: 'Context from file: file content' },
|
||||
]);
|
||||
}
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
@@ -4,33 +4,35 @@
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { promises as fs } from 'fs';
|
||||
import path from 'path';
|
||||
import { promises as fs } from 'node:fs';
|
||||
import path from 'node:path';
|
||||
import toml from '@iarna/toml';
|
||||
import { glob } from 'glob';
|
||||
import { z } from 'zod';
|
||||
import {
|
||||
Config,
|
||||
getProjectCommandsDir,
|
||||
getUserCommandsDir,
|
||||
} from '@qwen-code/qwen-code-core';
|
||||
import { ICommandLoader } from './types.js';
|
||||
import {
|
||||
import type { Config } from '@qwen-code/qwen-code-core';
|
||||
import { Storage } from '@qwen-code/qwen-code-core';
|
||||
import type { ICommandLoader } from './types.js';
|
||||
import type {
|
||||
CommandContext,
|
||||
CommandKind,
|
||||
SlashCommand,
|
||||
SlashCommandActionReturn,
|
||||
} from '../ui/commands/types.js';
|
||||
import { CommandKind } from '../ui/commands/types.js';
|
||||
import { DefaultArgumentProcessor } from './prompt-processors/argumentProcessor.js';
|
||||
import {
|
||||
import type {
|
||||
IPromptProcessor,
|
||||
PromptPipelineContent,
|
||||
} from './prompt-processors/types.js';
|
||||
import {
|
||||
SHORTHAND_ARGS_PLACEHOLDER,
|
||||
SHELL_INJECTION_TRIGGER,
|
||||
AT_FILE_INJECTION_TRIGGER,
|
||||
} from './prompt-processors/types.js';
|
||||
import {
|
||||
ConfirmationRequiredError,
|
||||
ShellProcessor,
|
||||
} from './prompt-processors/shellProcessor.js';
|
||||
import { AtFileProcessor } from './prompt-processors/atFileProcessor.js';
|
||||
|
||||
interface CommandDirectory {
|
||||
path: string;
|
||||
@@ -130,11 +132,13 @@ export class FileCommandLoader implements ICommandLoader {
|
||||
private getCommandDirectories(): CommandDirectory[] {
|
||||
const dirs: CommandDirectory[] = [];
|
||||
|
||||
const storage = this.config?.storage ?? new Storage(this.projectRoot);
|
||||
|
||||
// 1. User commands
|
||||
dirs.push({ path: getUserCommandsDir() });
|
||||
dirs.push({ path: Storage.getUserCommandsDir() });
|
||||
|
||||
// 2. Project commands (override user commands)
|
||||
dirs.push({ path: getProjectCommandsDir(this.projectRoot) });
|
||||
dirs.push({ path: storage.getProjectCommandsDir() });
|
||||
|
||||
// 3. Extension commands (processed last to detect all conflicts)
|
||||
if (this.config) {
|
||||
@@ -225,16 +229,25 @@ export class FileCommandLoader implements ICommandLoader {
|
||||
const usesShellInjection = validDef.prompt.includes(
|
||||
SHELL_INJECTION_TRIGGER,
|
||||
);
|
||||
const usesAtFileInjection = validDef.prompt.includes(
|
||||
AT_FILE_INJECTION_TRIGGER,
|
||||
);
|
||||
|
||||
// Interpolation (Shell Execution and Argument Injection)
|
||||
// If the prompt uses either shell injection OR argument placeholders,
|
||||
// we must use the ShellProcessor.
|
||||
// 1. @-File Injection (Security First).
|
||||
// This runs first to ensure we're not executing shell commands that
|
||||
// could dynamically generate malicious @-paths.
|
||||
if (usesAtFileInjection) {
|
||||
processors.push(new AtFileProcessor(baseCommandName));
|
||||
}
|
||||
|
||||
// 2. Argument and Shell Injection.
|
||||
// This runs after file content has been safely injected.
|
||||
if (usesShellInjection || usesArgs) {
|
||||
processors.push(new ShellProcessor(baseCommandName));
|
||||
}
|
||||
|
||||
// Default Argument Handling
|
||||
// If NO explicit argument injection ({{args}}) was used, we append the raw invocation.
|
||||
// 3. Default Argument Handling.
|
||||
// Appends the raw invocation if no explicit {{args}} are used.
|
||||
if (!usesArgs) {
|
||||
processors.push(new DefaultArgumentProcessor());
|
||||
}
|
||||
@@ -254,19 +267,24 @@ export class FileCommandLoader implements ICommandLoader {
|
||||
);
|
||||
return {
|
||||
type: 'submit_prompt',
|
||||
content: validDef.prompt, // Fallback to unprocessed prompt
|
||||
content: [{ text: validDef.prompt }], // Fallback to unprocessed prompt
|
||||
};
|
||||
}
|
||||
|
||||
try {
|
||||
let processedPrompt = validDef.prompt;
|
||||
let processedContent: PromptPipelineContent = [
|
||||
{ text: validDef.prompt },
|
||||
];
|
||||
for (const processor of processors) {
|
||||
processedPrompt = await processor.process(processedPrompt, context);
|
||||
processedContent = await processor.process(
|
||||
processedContent,
|
||||
context,
|
||||
);
|
||||
}
|
||||
|
||||
return {
|
||||
type: 'submit_prompt',
|
||||
content: processedPrompt,
|
||||
content: processedContent,
|
||||
};
|
||||
} catch (e) {
|
||||
// Check if it's our specific error type
|
||||
|
||||
128
packages/cli/src/services/McpPromptLoader.test.ts
Normal file
128
packages/cli/src/services/McpPromptLoader.test.ts
Normal file
@@ -0,0 +1,128 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { McpPromptLoader } from './McpPromptLoader.js';
|
||||
import type { Config } from '@qwen-code/qwen-code-core';
|
||||
import type { PromptArgument } from '@modelcontextprotocol/sdk/types.js';
|
||||
import { describe, it, expect } from 'vitest';
|
||||
|
||||
describe('McpPromptLoader', () => {
|
||||
const mockConfig = {} as Config;
|
||||
|
||||
describe('parseArgs', () => {
|
||||
it('should handle multi-word positional arguments', () => {
|
||||
const loader = new McpPromptLoader(mockConfig);
|
||||
const promptArgs: PromptArgument[] = [
|
||||
{ name: 'arg1', required: true },
|
||||
{ name: 'arg2', required: true },
|
||||
];
|
||||
const userArgs = 'hello world';
|
||||
const result = loader.parseArgs(userArgs, promptArgs);
|
||||
expect(result).toEqual({ arg1: 'hello', arg2: 'world' });
|
||||
});
|
||||
|
||||
it('should handle quoted multi-word positional arguments', () => {
|
||||
const loader = new McpPromptLoader(mockConfig);
|
||||
const promptArgs: PromptArgument[] = [
|
||||
{ name: 'arg1', required: true },
|
||||
{ name: 'arg2', required: true },
|
||||
];
|
||||
const userArgs = '"hello world" foo';
|
||||
const result = loader.parseArgs(userArgs, promptArgs);
|
||||
expect(result).toEqual({ arg1: 'hello world', arg2: 'foo' });
|
||||
});
|
||||
|
||||
it('should handle a single positional argument with multiple words', () => {
|
||||
const loader = new McpPromptLoader(mockConfig);
|
||||
const promptArgs: PromptArgument[] = [{ name: 'arg1', required: true }];
|
||||
const userArgs = 'hello world';
|
||||
const result = loader.parseArgs(userArgs, promptArgs);
|
||||
expect(result).toEqual({ arg1: 'hello world' });
|
||||
});
|
||||
|
||||
it('should handle escaped quotes in positional arguments', () => {
|
||||
const loader = new McpPromptLoader(mockConfig);
|
||||
const promptArgs: PromptArgument[] = [{ name: 'arg1', required: true }];
|
||||
const userArgs = '"hello \\"world\\""';
|
||||
const result = loader.parseArgs(userArgs, promptArgs);
|
||||
expect(result).toEqual({ arg1: 'hello "world"' });
|
||||
});
|
||||
|
||||
it('should handle escaped backslashes in positional arguments', () => {
|
||||
const loader = new McpPromptLoader(mockConfig);
|
||||
const promptArgs: PromptArgument[] = [{ name: 'arg1', required: true }];
|
||||
const userArgs = '"hello\\\\world"';
|
||||
const result = loader.parseArgs(userArgs, promptArgs);
|
||||
expect(result).toEqual({ arg1: 'hello\\world' });
|
||||
});
|
||||
|
||||
it('should handle named args followed by positional args', () => {
|
||||
const loader = new McpPromptLoader(mockConfig);
|
||||
const promptArgs: PromptArgument[] = [
|
||||
{ name: 'named', required: true },
|
||||
{ name: 'pos', required: true },
|
||||
];
|
||||
const userArgs = '--named="value" positional';
|
||||
const result = loader.parseArgs(userArgs, promptArgs);
|
||||
expect(result).toEqual({ named: 'value', pos: 'positional' });
|
||||
});
|
||||
|
||||
it('should handle positional args followed by named args', () => {
|
||||
const loader = new McpPromptLoader(mockConfig);
|
||||
const promptArgs: PromptArgument[] = [
|
||||
{ name: 'pos', required: true },
|
||||
{ name: 'named', required: true },
|
||||
];
|
||||
const userArgs = 'positional --named="value"';
|
||||
const result = loader.parseArgs(userArgs, promptArgs);
|
||||
expect(result).toEqual({ pos: 'positional', named: 'value' });
|
||||
});
|
||||
|
||||
it('should handle positional args interspersed with named args', () => {
|
||||
const loader = new McpPromptLoader(mockConfig);
|
||||
const promptArgs: PromptArgument[] = [
|
||||
{ name: 'pos1', required: true },
|
||||
{ name: 'named', required: true },
|
||||
{ name: 'pos2', required: true },
|
||||
];
|
||||
const userArgs = 'p1 --named="value" p2';
|
||||
const result = loader.parseArgs(userArgs, promptArgs);
|
||||
expect(result).toEqual({ pos1: 'p1', named: 'value', pos2: 'p2' });
|
||||
});
|
||||
|
||||
it('should treat an escaped quote at the start as a literal', () => {
|
||||
const loader = new McpPromptLoader(mockConfig);
|
||||
const promptArgs: PromptArgument[] = [
|
||||
{ name: 'arg1', required: true },
|
||||
{ name: 'arg2', required: true },
|
||||
];
|
||||
const userArgs = '\\"hello world';
|
||||
const result = loader.parseArgs(userArgs, promptArgs);
|
||||
expect(result).toEqual({ arg1: '"hello', arg2: 'world' });
|
||||
});
|
||||
|
||||
it('should handle a complex mix of args', () => {
|
||||
const loader = new McpPromptLoader(mockConfig);
|
||||
const promptArgs: PromptArgument[] = [
|
||||
{ name: 'pos1', required: true },
|
||||
{ name: 'named1', required: true },
|
||||
{ name: 'pos2', required: true },
|
||||
{ name: 'named2', required: true },
|
||||
{ name: 'pos3', required: true },
|
||||
];
|
||||
const userArgs =
|
||||
'p1 --named1="value 1" "p2 has spaces" --named2=value2 "p3 \\"with quotes\\""';
|
||||
const result = loader.parseArgs(userArgs, promptArgs);
|
||||
expect(result).toEqual({
|
||||
pos1: 'p1',
|
||||
named1: 'value 1',
|
||||
pos2: 'p2 has spaces',
|
||||
named2: 'value2',
|
||||
pos3: 'p3 "with quotes"',
|
||||
});
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -4,19 +4,19 @@
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { Config } from '@qwen-code/qwen-code-core';
|
||||
import {
|
||||
Config,
|
||||
getErrorMessage,
|
||||
getMCPServerPrompts,
|
||||
} from '@qwen-code/qwen-code-core';
|
||||
import {
|
||||
import type {
|
||||
CommandContext,
|
||||
CommandKind,
|
||||
SlashCommand,
|
||||
SlashCommandActionReturn,
|
||||
} from '../ui/commands/types.js';
|
||||
import { ICommandLoader } from './types.js';
|
||||
import { PromptArgument } from '@modelcontextprotocol/sdk/types.js';
|
||||
import { CommandKind } from '../ui/commands/types.js';
|
||||
import type { ICommandLoader } from './types.js';
|
||||
import type { PromptArgument } from '@modelcontextprotocol/sdk/types.js';
|
||||
|
||||
/**
|
||||
* Discovers and loads executable slash commands from prompts exposed by
|
||||
@@ -169,7 +169,16 @@ export class McpPromptLoader implements ICommandLoader {
|
||||
return Promise.resolve(promptCommands);
|
||||
}
|
||||
|
||||
private parseArgs(
|
||||
/**
|
||||
* Parses the `userArgs` string representing the prompt arguments (all the text
|
||||
* after the command) into a record matching the shape of the `promptArgs`.
|
||||
*
|
||||
* @param userArgs
|
||||
* @param promptArgs
|
||||
* @returns A record of the parsed arguments
|
||||
* @visibleForTesting
|
||||
*/
|
||||
parseArgs(
|
||||
userArgs: string,
|
||||
promptArgs: PromptArgument[] | undefined,
|
||||
): Record<string, unknown> | Error {
|
||||
@@ -177,28 +186,36 @@ export class McpPromptLoader implements ICommandLoader {
|
||||
const promptInputs: Record<string, unknown> = {};
|
||||
|
||||
// arg parsing: --key="value" or --key=value
|
||||
const namedArgRegex = /--([^=]+)=(?:"((?:\\.|[^"\\])*)"|([^ ]*))/g;
|
||||
const namedArgRegex = /--([^=]+)=(?:"((?:\\.|[^"\\])*)"|([^ ]+))/g;
|
||||
let match;
|
||||
const remainingArgs: string[] = [];
|
||||
let lastIndex = 0;
|
||||
const positionalParts: string[] = [];
|
||||
|
||||
while ((match = namedArgRegex.exec(userArgs)) !== null) {
|
||||
const key = match[1];
|
||||
const value = match[2] ?? match[3]; // Quoted or unquoted value
|
||||
// Extract the quoted or unquoted argument and remove escape chars.
|
||||
const value = (match[2] ?? match[3]).replace(/\\(.)/g, '$1');
|
||||
argValues[key] = value;
|
||||
// Capture text between matches as potential positional args
|
||||
if (match.index > lastIndex) {
|
||||
remainingArgs.push(userArgs.substring(lastIndex, match.index).trim());
|
||||
positionalParts.push(userArgs.substring(lastIndex, match.index));
|
||||
}
|
||||
lastIndex = namedArgRegex.lastIndex;
|
||||
}
|
||||
|
||||
// Capture any remaining text after the last named arg
|
||||
if (lastIndex < userArgs.length) {
|
||||
remainingArgs.push(userArgs.substring(lastIndex).trim());
|
||||
positionalParts.push(userArgs.substring(lastIndex));
|
||||
}
|
||||
|
||||
const positionalArgs = remainingArgs.join(' ').split(/ +/);
|
||||
const positionalArgsString = positionalParts.join('').trim();
|
||||
// extracts either quoted strings or non-quoted sequences of non-space characters.
|
||||
const positionalArgRegex = /(?:"((?:\\.|[^"\\])*)"|([^ ]+))/g;
|
||||
const positionalArgs: string[] = [];
|
||||
while ((match = positionalArgRegex.exec(positionalArgsString)) !== null) {
|
||||
// Extract the quoted or unquoted argument and remove escape chars.
|
||||
positionalArgs.push((match[1] ?? match[2]).replace(/\\(.)/g, '$1'));
|
||||
}
|
||||
|
||||
if (!promptArgs) {
|
||||
return promptInputs;
|
||||
@@ -213,19 +230,27 @@ export class McpPromptLoader implements ICommandLoader {
|
||||
(arg) => arg.required && !promptInputs[arg.name],
|
||||
);
|
||||
|
||||
const missingArgs: string[] = [];
|
||||
for (let i = 0; i < unfilledArgs.length; i++) {
|
||||
if (positionalArgs.length > i && positionalArgs[i]) {
|
||||
promptInputs[unfilledArgs[i].name] = positionalArgs[i];
|
||||
} else {
|
||||
missingArgs.push(unfilledArgs[i].name);
|
||||
if (unfilledArgs.length === 1) {
|
||||
// If we have only one unfilled arg, we don't require quotes we just
|
||||
// join all the given arguments together as if they were quoted.
|
||||
promptInputs[unfilledArgs[0].name] = positionalArgs.join(' ');
|
||||
} else {
|
||||
const missingArgs: string[] = [];
|
||||
for (let i = 0; i < unfilledArgs.length; i++) {
|
||||
if (positionalArgs.length > i) {
|
||||
promptInputs[unfilledArgs[i].name] = positionalArgs[i];
|
||||
} else {
|
||||
missingArgs.push(unfilledArgs[i].name);
|
||||
}
|
||||
}
|
||||
if (missingArgs.length > 0) {
|
||||
const missingArgNames = missingArgs
|
||||
.map((name) => `--${name}`)
|
||||
.join(', ');
|
||||
return new Error(`Missing required argument(s): ${missingArgNames}`);
|
||||
}
|
||||
}
|
||||
|
||||
if (missingArgs.length > 0) {
|
||||
const missingArgNames = missingArgs.map((name) => `--${name}`).join(', ');
|
||||
return new Error(`Missing required argument(s): ${missingArgNames}`);
|
||||
}
|
||||
return promptInputs;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -13,7 +13,7 @@ describe('Argument Processors', () => {
|
||||
const processor = new DefaultArgumentProcessor();
|
||||
|
||||
it('should append the full command if args are provided', async () => {
|
||||
const prompt = 'Parse the command.';
|
||||
const prompt = [{ text: 'Parse the command.' }];
|
||||
const context = createMockCommandContext({
|
||||
invocation: {
|
||||
raw: '/mycommand arg1 "arg two"',
|
||||
@@ -22,11 +22,13 @@ describe('Argument Processors', () => {
|
||||
},
|
||||
});
|
||||
const result = await processor.process(prompt, context);
|
||||
expect(result).toBe('Parse the command.\n\n/mycommand arg1 "arg two"');
|
||||
expect(result).toEqual([
|
||||
{ text: '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 prompt = [{ text: 'Parse the command.' }];
|
||||
const context = createMockCommandContext({
|
||||
invocation: {
|
||||
raw: '/mycommand',
|
||||
@@ -35,7 +37,7 @@ describe('Argument Processors', () => {
|
||||
},
|
||||
});
|
||||
const result = await processor.process(prompt, context);
|
||||
expect(result).toBe('Parse the command.');
|
||||
expect(result).toEqual([{ text: 'Parse the command.' }]);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
@@ -4,8 +4,9 @@
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { IPromptProcessor } from './types.js';
|
||||
import { CommandContext } from '../../ui/commands/types.js';
|
||||
import { appendToLastTextPart } from '@qwen-code/qwen-code-core';
|
||||
import type { IPromptProcessor, PromptPipelineContent } from './types.js';
|
||||
import type { CommandContext } from '../../ui/commands/types.js';
|
||||
|
||||
/**
|
||||
* Appends the user's full command invocation to the prompt if arguments are
|
||||
@@ -14,9 +15,12 @@ import { CommandContext } from '../../ui/commands/types.js';
|
||||
* 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> {
|
||||
if (context.invocation!.args) {
|
||||
return `${prompt}\n\n${context.invocation!.raw}`;
|
||||
async process(
|
||||
prompt: PromptPipelineContent,
|
||||
context: CommandContext,
|
||||
): Promise<PromptPipelineContent> {
|
||||
if (context.invocation?.args) {
|
||||
return appendToLastTextPart(prompt, context.invocation.raw);
|
||||
}
|
||||
return prompt;
|
||||
}
|
||||
|
||||
@@ -0,0 +1,221 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect, beforeEach, vi } from 'vitest';
|
||||
import { createMockCommandContext } from '../../test-utils/mockCommandContext.js';
|
||||
import { type CommandContext } from '../../ui/commands/types.js';
|
||||
import { AtFileProcessor } from './atFileProcessor.js';
|
||||
import { MessageType } from '../../ui/types.js';
|
||||
import type { Config } from '@qwen-code/qwen-code-core';
|
||||
import type { PartUnion } from '@google/genai';
|
||||
|
||||
// Mock the core dependency
|
||||
const mockReadPathFromWorkspace = vi.hoisted(() => vi.fn());
|
||||
vi.mock('@qwen-code/qwen-code-core', async (importOriginal) => {
|
||||
const original = await importOriginal<object>();
|
||||
return {
|
||||
...original,
|
||||
readPathFromWorkspace: mockReadPathFromWorkspace,
|
||||
};
|
||||
});
|
||||
|
||||
describe('AtFileProcessor', () => {
|
||||
let context: CommandContext;
|
||||
let mockConfig: Config;
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
|
||||
mockConfig = {
|
||||
// The processor only passes the config through, so we don't need a full mock.
|
||||
} as unknown as Config;
|
||||
|
||||
context = createMockCommandContext({
|
||||
services: {
|
||||
config: mockConfig,
|
||||
},
|
||||
});
|
||||
|
||||
// Default mock success behavior: return content wrapped in a text part.
|
||||
mockReadPathFromWorkspace.mockImplementation(
|
||||
async (path: string): Promise<PartUnion[]> => [
|
||||
{ text: `content of ${path}` },
|
||||
],
|
||||
);
|
||||
});
|
||||
|
||||
it('should not change the prompt if no @{ trigger is present', async () => {
|
||||
const processor = new AtFileProcessor();
|
||||
const prompt: PartUnion[] = [{ text: 'This is a simple prompt.' }];
|
||||
const result = await processor.process(prompt, context);
|
||||
expect(result).toEqual(prompt);
|
||||
expect(mockReadPathFromWorkspace).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should not change the prompt if config service is missing', async () => {
|
||||
const processor = new AtFileProcessor();
|
||||
const prompt: PartUnion[] = [{ text: 'Analyze @{file.txt}' }];
|
||||
const contextWithoutConfig = createMockCommandContext({
|
||||
services: {
|
||||
config: null,
|
||||
},
|
||||
});
|
||||
const result = await processor.process(prompt, contextWithoutConfig);
|
||||
expect(result).toEqual(prompt);
|
||||
expect(mockReadPathFromWorkspace).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
describe('Parsing Logic', () => {
|
||||
it('should replace a single valid @{path/to/file.txt} placeholder', async () => {
|
||||
const processor = new AtFileProcessor();
|
||||
const prompt: PartUnion[] = [
|
||||
{ text: 'Analyze this file: @{path/to/file.txt}' },
|
||||
];
|
||||
const result = await processor.process(prompt, context);
|
||||
expect(mockReadPathFromWorkspace).toHaveBeenCalledWith(
|
||||
'path/to/file.txt',
|
||||
mockConfig,
|
||||
);
|
||||
expect(result).toEqual([
|
||||
{ text: 'Analyze this file: ' },
|
||||
{ text: 'content of path/to/file.txt' },
|
||||
]);
|
||||
});
|
||||
|
||||
it('should replace multiple different @{...} placeholders', async () => {
|
||||
const processor = new AtFileProcessor();
|
||||
const prompt: PartUnion[] = [
|
||||
{ text: 'Compare @{file1.js} with @{file2.js}' },
|
||||
];
|
||||
const result = await processor.process(prompt, context);
|
||||
expect(mockReadPathFromWorkspace).toHaveBeenCalledTimes(2);
|
||||
expect(mockReadPathFromWorkspace).toHaveBeenCalledWith(
|
||||
'file1.js',
|
||||
mockConfig,
|
||||
);
|
||||
expect(mockReadPathFromWorkspace).toHaveBeenCalledWith(
|
||||
'file2.js',
|
||||
mockConfig,
|
||||
);
|
||||
expect(result).toEqual([
|
||||
{ text: 'Compare ' },
|
||||
{ text: 'content of file1.js' },
|
||||
{ text: ' with ' },
|
||||
{ text: 'content of file2.js' },
|
||||
]);
|
||||
});
|
||||
|
||||
it('should handle placeholders at the beginning, middle, and end', async () => {
|
||||
const processor = new AtFileProcessor();
|
||||
const prompt: PartUnion[] = [
|
||||
{ text: '@{start.txt} in the @{middle.txt} and @{end.txt}' },
|
||||
];
|
||||
const result = await processor.process(prompt, context);
|
||||
expect(result).toEqual([
|
||||
{ text: 'content of start.txt' },
|
||||
{ text: ' in the ' },
|
||||
{ text: 'content of middle.txt' },
|
||||
{ text: ' and ' },
|
||||
{ text: 'content of end.txt' },
|
||||
]);
|
||||
});
|
||||
|
||||
it('should correctly parse paths that contain balanced braces', async () => {
|
||||
const processor = new AtFileProcessor();
|
||||
const prompt: PartUnion[] = [
|
||||
{ text: 'Analyze @{path/with/{braces}/file.txt}' },
|
||||
];
|
||||
const result = await processor.process(prompt, context);
|
||||
expect(mockReadPathFromWorkspace).toHaveBeenCalledWith(
|
||||
'path/with/{braces}/file.txt',
|
||||
mockConfig,
|
||||
);
|
||||
expect(result).toEqual([
|
||||
{ text: 'Analyze ' },
|
||||
{ text: 'content of path/with/{braces}/file.txt' },
|
||||
]);
|
||||
});
|
||||
|
||||
it('should throw an error if the prompt contains an unclosed trigger', async () => {
|
||||
const processor = new AtFileProcessor();
|
||||
const prompt: PartUnion[] = [{ text: 'Hello @{world' }];
|
||||
// The new parser throws an error for unclosed injections.
|
||||
await expect(processor.process(prompt, context)).rejects.toThrow(
|
||||
/Unclosed injection/,
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Integration and Error Handling', () => {
|
||||
it('should leave the placeholder unmodified if readPathFromWorkspace throws', async () => {
|
||||
const processor = new AtFileProcessor();
|
||||
const prompt: PartUnion[] = [
|
||||
{ text: 'Analyze @{not-found.txt} and @{good-file.txt}' },
|
||||
];
|
||||
mockReadPathFromWorkspace.mockImplementation(async (path: string) => {
|
||||
if (path === 'not-found.txt') {
|
||||
throw new Error('File not found');
|
||||
}
|
||||
return [{ text: `content of ${path}` }];
|
||||
});
|
||||
|
||||
const result = await processor.process(prompt, context);
|
||||
expect(result).toEqual([
|
||||
{ text: 'Analyze ' },
|
||||
{ text: '@{not-found.txt}' }, // Placeholder is preserved as a text part
|
||||
{ text: ' and ' },
|
||||
{ text: 'content of good-file.txt' },
|
||||
]);
|
||||
});
|
||||
});
|
||||
|
||||
describe('UI Feedback', () => {
|
||||
it('should call ui.addItem with an ERROR on failure', async () => {
|
||||
const processor = new AtFileProcessor();
|
||||
const prompt: PartUnion[] = [{ text: 'Analyze @{bad-file.txt}' }];
|
||||
mockReadPathFromWorkspace.mockRejectedValue(new Error('Access denied'));
|
||||
|
||||
await processor.process(prompt, context);
|
||||
|
||||
expect(context.ui.addItem).toHaveBeenCalledTimes(1);
|
||||
expect(context.ui.addItem).toHaveBeenCalledWith(
|
||||
{
|
||||
type: MessageType.ERROR,
|
||||
text: "Failed to inject content for '@{bad-file.txt}': Access denied",
|
||||
},
|
||||
expect.any(Number),
|
||||
);
|
||||
});
|
||||
|
||||
it('should call ui.addItem with a WARNING if the file was ignored', async () => {
|
||||
const processor = new AtFileProcessor();
|
||||
const prompt: PartUnion[] = [{ text: 'Analyze @{ignored.txt}' }];
|
||||
// Simulate an ignored file by returning an empty array.
|
||||
mockReadPathFromWorkspace.mockResolvedValue([]);
|
||||
|
||||
const result = await processor.process(prompt, context);
|
||||
|
||||
// The placeholder should be removed, resulting in only the prefix.
|
||||
expect(result).toEqual([{ text: 'Analyze ' }]);
|
||||
|
||||
expect(context.ui.addItem).toHaveBeenCalledTimes(1);
|
||||
expect(context.ui.addItem).toHaveBeenCalledWith(
|
||||
{
|
||||
type: MessageType.INFO,
|
||||
text: "File '@{ignored.txt}' was ignored by .gitignore or .geminiignore and was not included in the prompt.",
|
||||
},
|
||||
expect.any(Number),
|
||||
);
|
||||
});
|
||||
|
||||
it('should NOT call ui.addItem on success', async () => {
|
||||
const processor = new AtFileProcessor();
|
||||
const prompt: PartUnion[] = [{ text: 'Analyze @{good-file.txt}' }];
|
||||
await processor.process(prompt, context);
|
||||
expect(context.ui.addItem).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,96 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import {
|
||||
flatMapTextParts,
|
||||
readPathFromWorkspace,
|
||||
} from '@qwen-code/qwen-code-core';
|
||||
import type { CommandContext } from '../../ui/commands/types.js';
|
||||
import { MessageType } from '../../ui/types.js';
|
||||
import {
|
||||
AT_FILE_INJECTION_TRIGGER,
|
||||
type IPromptProcessor,
|
||||
type PromptPipelineContent,
|
||||
} from './types.js';
|
||||
import { extractInjections } from './injectionParser.js';
|
||||
|
||||
export class AtFileProcessor implements IPromptProcessor {
|
||||
constructor(private readonly commandName?: string) {}
|
||||
|
||||
async process(
|
||||
input: PromptPipelineContent,
|
||||
context: CommandContext,
|
||||
): Promise<PromptPipelineContent> {
|
||||
const config = context.services.config;
|
||||
if (!config) {
|
||||
return input;
|
||||
}
|
||||
|
||||
return flatMapTextParts(input, async (text) => {
|
||||
if (!text.includes(AT_FILE_INJECTION_TRIGGER)) {
|
||||
return [{ text }];
|
||||
}
|
||||
|
||||
const injections = extractInjections(
|
||||
text,
|
||||
AT_FILE_INJECTION_TRIGGER,
|
||||
this.commandName,
|
||||
);
|
||||
if (injections.length === 0) {
|
||||
return [{ text }];
|
||||
}
|
||||
|
||||
const output: PromptPipelineContent = [];
|
||||
let lastIndex = 0;
|
||||
|
||||
for (const injection of injections) {
|
||||
const prefix = text.substring(lastIndex, injection.startIndex);
|
||||
if (prefix) {
|
||||
output.push({ text: prefix });
|
||||
}
|
||||
|
||||
const pathStr = injection.content;
|
||||
try {
|
||||
const fileContentParts = await readPathFromWorkspace(pathStr, config);
|
||||
if (fileContentParts.length === 0) {
|
||||
const uiMessage = `File '@{${pathStr}}' was ignored by .gitignore or .geminiignore and was not included in the prompt.`;
|
||||
context.ui.addItem(
|
||||
{ type: MessageType.INFO, text: uiMessage },
|
||||
Date.now(),
|
||||
);
|
||||
}
|
||||
output.push(...fileContentParts);
|
||||
} catch (error) {
|
||||
const message =
|
||||
error instanceof Error ? error.message : String(error);
|
||||
const uiMessage = `Failed to inject content for '@{${pathStr}}': ${message}`;
|
||||
|
||||
console.error(
|
||||
`[AtFileProcessor] ${uiMessage}. Leaving placeholder in prompt.`,
|
||||
);
|
||||
context.ui.addItem(
|
||||
{ type: MessageType.ERROR, text: uiMessage },
|
||||
Date.now(),
|
||||
);
|
||||
|
||||
const placeholder = text.substring(
|
||||
injection.startIndex,
|
||||
injection.endIndex,
|
||||
);
|
||||
output.push({ text: placeholder });
|
||||
}
|
||||
lastIndex = injection.endIndex;
|
||||
}
|
||||
|
||||
const suffix = text.substring(lastIndex);
|
||||
if (suffix) {
|
||||
output.push({ text: suffix });
|
||||
}
|
||||
|
||||
return output;
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,223 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect } from 'vitest';
|
||||
import { extractInjections } from './injectionParser.js';
|
||||
|
||||
describe('extractInjections', () => {
|
||||
const SHELL_TRIGGER = '!{';
|
||||
const AT_FILE_TRIGGER = '@{';
|
||||
|
||||
describe('Basic Functionality', () => {
|
||||
it('should return an empty array if no trigger is present', () => {
|
||||
const prompt = 'This is a simple prompt without injections.';
|
||||
const result = extractInjections(prompt, SHELL_TRIGGER);
|
||||
expect(result).toEqual([]);
|
||||
});
|
||||
|
||||
it('should extract a single, simple injection', () => {
|
||||
const prompt = 'Run this command: !{ls -la}';
|
||||
const result = extractInjections(prompt, SHELL_TRIGGER);
|
||||
expect(result).toEqual([
|
||||
{
|
||||
content: 'ls -la',
|
||||
startIndex: 18,
|
||||
endIndex: 27,
|
||||
},
|
||||
]);
|
||||
});
|
||||
|
||||
it('should extract multiple injections', () => {
|
||||
const prompt = 'First: !{cmd1}, Second: !{cmd2}';
|
||||
const result = extractInjections(prompt, SHELL_TRIGGER);
|
||||
expect(result).toHaveLength(2);
|
||||
expect(result[0]).toEqual({
|
||||
content: 'cmd1',
|
||||
startIndex: 7,
|
||||
endIndex: 14,
|
||||
});
|
||||
expect(result[1]).toEqual({
|
||||
content: 'cmd2',
|
||||
startIndex: 24,
|
||||
endIndex: 31,
|
||||
});
|
||||
});
|
||||
|
||||
it('should handle different triggers (e.g., @{)', () => {
|
||||
const prompt = 'Read this file: @{path/to/file.txt}';
|
||||
const result = extractInjections(prompt, AT_FILE_TRIGGER);
|
||||
expect(result).toEqual([
|
||||
{
|
||||
content: 'path/to/file.txt',
|
||||
startIndex: 16,
|
||||
endIndex: 35,
|
||||
},
|
||||
]);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Positioning and Edge Cases', () => {
|
||||
it('should handle injections at the start and end of the prompt', () => {
|
||||
const prompt = '!{start} middle text !{end}';
|
||||
const result = extractInjections(prompt, SHELL_TRIGGER);
|
||||
expect(result).toHaveLength(2);
|
||||
expect(result[0]).toEqual({
|
||||
content: 'start',
|
||||
startIndex: 0,
|
||||
endIndex: 8,
|
||||
});
|
||||
expect(result[1]).toEqual({
|
||||
content: 'end',
|
||||
startIndex: 21,
|
||||
endIndex: 27,
|
||||
});
|
||||
});
|
||||
|
||||
it('should handle adjacent injections', () => {
|
||||
const prompt = '!{A}!{B}';
|
||||
const result = extractInjections(prompt, SHELL_TRIGGER);
|
||||
expect(result).toHaveLength(2);
|
||||
expect(result[0]).toEqual({ content: 'A', startIndex: 0, endIndex: 4 });
|
||||
expect(result[1]).toEqual({ content: 'B', startIndex: 4, endIndex: 8 });
|
||||
});
|
||||
|
||||
it('should handle empty injections', () => {
|
||||
const prompt = 'Empty: !{}';
|
||||
const result = extractInjections(prompt, SHELL_TRIGGER);
|
||||
expect(result).toEqual([
|
||||
{
|
||||
content: '',
|
||||
startIndex: 7,
|
||||
endIndex: 10,
|
||||
},
|
||||
]);
|
||||
});
|
||||
|
||||
it('should trim whitespace within the content', () => {
|
||||
const prompt = '!{ \n command with space \t }';
|
||||
const result = extractInjections(prompt, SHELL_TRIGGER);
|
||||
expect(result).toEqual([
|
||||
{
|
||||
content: 'command with space',
|
||||
startIndex: 0,
|
||||
endIndex: 29,
|
||||
},
|
||||
]);
|
||||
});
|
||||
|
||||
it('should ignore similar patterns that are not the exact trigger', () => {
|
||||
const prompt = 'Not a trigger: !(cmd) or {cmd} or ! {cmd}';
|
||||
const result = extractInjections(prompt, SHELL_TRIGGER);
|
||||
expect(result).toEqual([]);
|
||||
});
|
||||
|
||||
it('should ignore extra closing braces before the trigger', () => {
|
||||
const prompt = 'Ignore this } then !{run}';
|
||||
const result = extractInjections(prompt, SHELL_TRIGGER);
|
||||
expect(result).toEqual([
|
||||
{
|
||||
content: 'run',
|
||||
startIndex: 19,
|
||||
endIndex: 25,
|
||||
},
|
||||
]);
|
||||
});
|
||||
|
||||
it('should stop parsing at the first balanced closing brace (non-greedy)', () => {
|
||||
// This tests that the parser doesn't greedily consume extra closing braces
|
||||
const prompt = 'Run !{ls -l}} extra braces';
|
||||
const result = extractInjections(prompt, SHELL_TRIGGER);
|
||||
expect(result).toEqual([
|
||||
{
|
||||
content: 'ls -l',
|
||||
startIndex: 4,
|
||||
endIndex: 12,
|
||||
},
|
||||
]);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Nested Braces (Balanced)', () => {
|
||||
it('should correctly parse content with simple nested braces (e.g., JSON)', () => {
|
||||
const prompt = `Send JSON: !{curl -d '{"key": "value"}'}`;
|
||||
const result = extractInjections(prompt, SHELL_TRIGGER);
|
||||
expect(result).toHaveLength(1);
|
||||
expect(result[0].content).toBe(`curl -d '{"key": "value"}'`);
|
||||
});
|
||||
|
||||
it('should correctly parse content with shell constructs (e.g., awk)', () => {
|
||||
const prompt = `Process text: !{awk '{print $1}' file.txt}`;
|
||||
const result = extractInjections(prompt, SHELL_TRIGGER);
|
||||
expect(result).toHaveLength(1);
|
||||
expect(result[0].content).toBe(`awk '{print $1}' file.txt`);
|
||||
});
|
||||
|
||||
it('should correctly parse multiple levels of nesting', () => {
|
||||
const prompt = `!{level1 {level2 {level3}} suffix}`;
|
||||
const result = extractInjections(prompt, SHELL_TRIGGER);
|
||||
expect(result).toHaveLength(1);
|
||||
expect(result[0].content).toBe(`level1 {level2 {level3}} suffix`);
|
||||
expect(result[0].endIndex).toBe(prompt.length);
|
||||
});
|
||||
|
||||
it('should correctly parse paths containing balanced braces', () => {
|
||||
const prompt = 'Analyze @{path/with/{braces}/file.txt}';
|
||||
const result = extractInjections(prompt, AT_FILE_TRIGGER);
|
||||
expect(result).toHaveLength(1);
|
||||
expect(result[0].content).toBe('path/with/{braces}/file.txt');
|
||||
});
|
||||
|
||||
it('should correctly handle an injection containing the trigger itself', () => {
|
||||
// This works because the parser counts braces, it doesn't look for the trigger again until the current one is closed.
|
||||
const prompt = '!{echo "The trigger is !{ confusing }"}';
|
||||
const expectedContent = 'echo "The trigger is !{ confusing }"';
|
||||
const result = extractInjections(prompt, SHELL_TRIGGER);
|
||||
expect(result).toHaveLength(1);
|
||||
expect(result[0].content).toBe(expectedContent);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Error Handling (Unbalanced/Unclosed)', () => {
|
||||
it('should throw an error for a simple unclosed injection', () => {
|
||||
const prompt = 'This prompt has !{an unclosed trigger';
|
||||
expect(() => extractInjections(prompt, SHELL_TRIGGER)).toThrow(
|
||||
/Invalid syntax: Unclosed injection starting at index 16 \('!{'\)/,
|
||||
);
|
||||
});
|
||||
|
||||
it('should throw an error if the prompt ends inside a nested block', () => {
|
||||
const prompt = 'This fails: !{outer {inner';
|
||||
expect(() => extractInjections(prompt, SHELL_TRIGGER)).toThrow(
|
||||
/Invalid syntax: Unclosed injection starting at index 12 \('!{'\)/,
|
||||
);
|
||||
});
|
||||
|
||||
it('should include the context name in the error message if provided', () => {
|
||||
const prompt = 'Failing !{command';
|
||||
const contextName = 'test-command';
|
||||
expect(() =>
|
||||
extractInjections(prompt, SHELL_TRIGGER, contextName),
|
||||
).toThrow(
|
||||
/Invalid syntax in command 'test-command': Unclosed injection starting at index 8/,
|
||||
);
|
||||
});
|
||||
|
||||
it('should throw if content contains unbalanced braces (e.g., missing closing)', () => {
|
||||
// This is functionally the same as an unclosed injection from the parser's perspective.
|
||||
const prompt = 'Analyze @{path/with/braces{example.txt}';
|
||||
expect(() => extractInjections(prompt, AT_FILE_TRIGGER)).toThrow(
|
||||
/Invalid syntax: Unclosed injection starting at index 8 \('@{'\)/,
|
||||
);
|
||||
});
|
||||
|
||||
it('should clearly state that unbalanced braces in content are not supported in the error', () => {
|
||||
const prompt = 'Analyze @{path/with/braces{example.txt}';
|
||||
expect(() => extractInjections(prompt, AT_FILE_TRIGGER)).toThrow(
|
||||
/Paths or commands with unbalanced braces are not supported directly/,
|
||||
);
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,89 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
/**
|
||||
* Represents a single detected injection site in a prompt string.
|
||||
*/
|
||||
export interface Injection {
|
||||
/** The content extracted from within the braces (e.g., the command or path), trimmed. */
|
||||
content: string;
|
||||
/** The starting index of the injection (inclusive, points to the start of the trigger). */
|
||||
startIndex: number;
|
||||
/** The ending index of the injection (exclusive, points after the closing '}'). */
|
||||
endIndex: number;
|
||||
}
|
||||
|
||||
/**
|
||||
* Iteratively parses a prompt string to extract injections (e.g., !{...} or @{...}),
|
||||
* correctly handling nested braces within the content.
|
||||
*
|
||||
* This parser relies on simple brace counting and does not support escaping.
|
||||
*
|
||||
* @param prompt The prompt string to parse.
|
||||
* @param trigger The opening trigger sequence (e.g., '!{', '@{').
|
||||
* @param contextName Optional context name (e.g., command name) for error messages.
|
||||
* @returns An array of extracted Injection objects.
|
||||
* @throws Error if an unclosed injection is found.
|
||||
*/
|
||||
export function extractInjections(
|
||||
prompt: string,
|
||||
trigger: string,
|
||||
contextName?: string,
|
||||
): Injection[] {
|
||||
const injections: Injection[] = [];
|
||||
let index = 0;
|
||||
|
||||
while (index < prompt.length) {
|
||||
const startIndex = prompt.indexOf(trigger, index);
|
||||
|
||||
if (startIndex === -1) {
|
||||
break;
|
||||
}
|
||||
|
||||
let currentIndex = startIndex + trigger.length;
|
||||
let braceCount = 1;
|
||||
let foundEnd = false;
|
||||
|
||||
while (currentIndex < prompt.length) {
|
||||
const char = prompt[currentIndex];
|
||||
|
||||
if (char === '{') {
|
||||
braceCount++;
|
||||
} else if (char === '}') {
|
||||
braceCount--;
|
||||
if (braceCount === 0) {
|
||||
const injectionContent = prompt.substring(
|
||||
startIndex + trigger.length,
|
||||
currentIndex,
|
||||
);
|
||||
const endIndex = currentIndex + 1;
|
||||
|
||||
injections.push({
|
||||
content: injectionContent.trim(),
|
||||
startIndex,
|
||||
endIndex,
|
||||
});
|
||||
|
||||
index = endIndex;
|
||||
foundEnd = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
currentIndex++;
|
||||
}
|
||||
|
||||
// Check if the inner loop finished without finding the closing brace.
|
||||
if (!foundEnd) {
|
||||
const contextInfo = contextName ? ` in command '${contextName}'` : '';
|
||||
// Enforce strict parsing (Comment 1) and clarify limitations (Comment 2).
|
||||
throw new Error(
|
||||
`Invalid syntax${contextInfo}: Unclosed injection starting at index ${startIndex} ('${trigger}'). Ensure braces are balanced. Paths or commands with unbalanced braces are not supported directly.`,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
return injections;
|
||||
}
|
||||
@@ -7,14 +7,16 @@
|
||||
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 { ApprovalMode, Config } from '@qwen-code/qwen-code-core';
|
||||
import os from 'os';
|
||||
import type { CommandContext } from '../../ui/commands/types.js';
|
||||
import type { Config } from '@qwen-code/qwen-code-core';
|
||||
import { ApprovalMode } from '@qwen-code/qwen-code-core';
|
||||
import os from 'node:os';
|
||||
import { quote } from 'shell-quote';
|
||||
import { createPartFromText } from '@google/genai';
|
||||
import type { PromptPipelineContent } from './types.js';
|
||||
|
||||
// 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.
|
||||
// mirroring the logic in the actual `escapeShellArg` implementation.
|
||||
function getExpectedEscapedArgForPlatform(arg: string): string {
|
||||
if (os.platform() === 'win32') {
|
||||
const comSpec = (process.env['ComSpec'] || 'cmd.exe').toLowerCase();
|
||||
@@ -31,6 +33,11 @@ function getExpectedEscapedArgForPlatform(arg: string): string {
|
||||
}
|
||||
}
|
||||
|
||||
// Helper to create PromptPipelineContent
|
||||
function createPromptPipelineContent(text: string): PromptPipelineContent {
|
||||
return [createPartFromText(text)];
|
||||
}
|
||||
|
||||
const mockCheckCommandPermissions = vi.hoisted(() => vi.fn());
|
||||
const mockShellExecute = vi.hoisted(() => vi.fn());
|
||||
|
||||
@@ -92,7 +99,7 @@ describe('ShellProcessor', () => {
|
||||
|
||||
it('should throw an error if config is missing', async () => {
|
||||
const processor = new ShellProcessor('test-command');
|
||||
const prompt = '!{ls}';
|
||||
const prompt: PromptPipelineContent = createPromptPipelineContent('!{ls}');
|
||||
const contextWithoutConfig = createMockCommandContext({
|
||||
services: {
|
||||
config: null,
|
||||
@@ -106,15 +113,19 @@ describe('ShellProcessor', () => {
|
||||
|
||||
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 prompt: PromptPipelineContent = createPromptPipelineContent(
|
||||
'This is a simple prompt with no injections.',
|
||||
);
|
||||
const result = await processor.process(prompt, context);
|
||||
expect(result).toBe(prompt);
|
||||
expect(result).toEqual(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}';
|
||||
const prompt: PromptPipelineContent = createPromptPipelineContent(
|
||||
'The current status is: !{git status}',
|
||||
);
|
||||
mockCheckCommandPermissions.mockReturnValue({
|
||||
allAllowed: true,
|
||||
disallowedCommands: [],
|
||||
@@ -137,12 +148,14 @@ describe('ShellProcessor', () => {
|
||||
expect.any(Object),
|
||||
false,
|
||||
);
|
||||
expect(result).toBe('The current status is: On branch main');
|
||||
expect(result).toEqual([{ text: '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}';
|
||||
const prompt: PromptPipelineContent = createPromptPipelineContent(
|
||||
'!{git status} in !{pwd}',
|
||||
);
|
||||
mockCheckCommandPermissions.mockReturnValue({
|
||||
allAllowed: true,
|
||||
disallowedCommands: [],
|
||||
@@ -163,12 +176,14 @@ describe('ShellProcessor', () => {
|
||||
|
||||
expect(mockCheckCommandPermissions).toHaveBeenCalledTimes(2);
|
||||
expect(mockShellExecute).toHaveBeenCalledTimes(2);
|
||||
expect(result).toBe('On branch main in /usr/home');
|
||||
expect(result).toEqual([{ text: 'On branch main in /usr/home' }]);
|
||||
});
|
||||
|
||||
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 /}';
|
||||
const prompt: PromptPipelineContent = createPromptPipelineContent(
|
||||
'Do something dangerous: !{rm -rf /}',
|
||||
);
|
||||
mockCheckCommandPermissions.mockReturnValue({
|
||||
allAllowed: false,
|
||||
disallowedCommands: ['rm -rf /'],
|
||||
@@ -181,7 +196,9 @@ 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 /}';
|
||||
const prompt: PromptPipelineContent = createPromptPipelineContent(
|
||||
'Do something dangerous: !{rm -rf /}',
|
||||
);
|
||||
mockCheckCommandPermissions.mockReturnValue({
|
||||
allAllowed: false,
|
||||
disallowedCommands: ['rm -rf /'],
|
||||
@@ -202,12 +219,14 @@ describe('ShellProcessor', () => {
|
||||
expect.any(Object),
|
||||
false,
|
||||
);
|
||||
expect(result).toBe('Do something dangerous: deleted');
|
||||
expect(result).toEqual([{ text: '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}';
|
||||
const prompt: PromptPipelineContent = createPromptPipelineContent(
|
||||
'Do something forbidden: !{reboot}',
|
||||
);
|
||||
mockCheckCommandPermissions.mockReturnValue({
|
||||
allAllowed: false,
|
||||
disallowedCommands: ['reboot'],
|
||||
@@ -227,7 +246,9 @@ describe('ShellProcessor', () => {
|
||||
|
||||
it('should throw ConfirmationRequiredError with the correct command', async () => {
|
||||
const processor = new ShellProcessor('test-command');
|
||||
const prompt = 'Do something dangerous: !{rm -rf /}';
|
||||
const prompt: PromptPipelineContent = createPromptPipelineContent(
|
||||
'Do something dangerous: !{rm -rf /}',
|
||||
);
|
||||
mockCheckCommandPermissions.mockReturnValue({
|
||||
allAllowed: false,
|
||||
disallowedCommands: ['rm -rf /'],
|
||||
@@ -249,7 +270,9 @@ describe('ShellProcessor', () => {
|
||||
|
||||
it('should throw ConfirmationRequiredError with multiple commands if multiple are disallowed', async () => {
|
||||
const processor = new ShellProcessor('test-command');
|
||||
const prompt = '!{cmd1} and !{cmd2}';
|
||||
const prompt: PromptPipelineContent = createPromptPipelineContent(
|
||||
'!{cmd1} and !{cmd2}',
|
||||
);
|
||||
mockCheckCommandPermissions.mockImplementation((cmd) => {
|
||||
if (cmd === 'cmd1') {
|
||||
return { allAllowed: false, disallowedCommands: ['cmd1'] };
|
||||
@@ -274,7 +297,9 @@ describe('ShellProcessor', () => {
|
||||
|
||||
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 /}';
|
||||
const prompt: PromptPipelineContent = createPromptPipelineContent(
|
||||
'First: !{echo "hello"}, Second: !{rm -rf /}',
|
||||
);
|
||||
|
||||
mockCheckCommandPermissions.mockImplementation((cmd) => {
|
||||
if (cmd.includes('rm')) {
|
||||
@@ -293,7 +318,9 @@ describe('ShellProcessor', () => {
|
||||
|
||||
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 /}';
|
||||
const prompt: PromptPipelineContent = createPromptPipelineContent(
|
||||
'Allowed: !{ls -l}, Disallowed: !{rm -rf /}',
|
||||
);
|
||||
|
||||
mockCheckCommandPermissions.mockImplementation((cmd) => ({
|
||||
allAllowed: !cmd.includes('rm'),
|
||||
@@ -313,7 +340,9 @@ describe('ShellProcessor', () => {
|
||||
|
||||
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}';
|
||||
const prompt: PromptPipelineContent = createPromptPipelineContent(
|
||||
'Run !{cmd1} and !{cmd2}',
|
||||
);
|
||||
|
||||
// Add commands to the session allowlist
|
||||
context.session.sessionShellAllowlist = new Set(['cmd1', 'cmd2']);
|
||||
@@ -345,12 +374,14 @@ describe('ShellProcessor', () => {
|
||||
context.session.sessionShellAllowlist,
|
||||
);
|
||||
expect(mockShellExecute).toHaveBeenCalledTimes(2);
|
||||
expect(result).toBe('Run output1 and output2');
|
||||
expect(result).toEqual([{ text: 'Run output1 and output2' }]);
|
||||
});
|
||||
|
||||
it('should trim whitespace from the command inside the injection before interpolation', async () => {
|
||||
const processor = new ShellProcessor('test-command');
|
||||
const prompt = 'Files: !{ ls {{args}} -l }';
|
||||
const prompt: PromptPipelineContent = createPromptPipelineContent(
|
||||
'Files: !{ ls {{args}} -l }',
|
||||
);
|
||||
|
||||
const rawArgs = context.invocation!.args;
|
||||
|
||||
@@ -384,7 +415,8 @@ describe('ShellProcessor', () => {
|
||||
|
||||
it('should handle an empty command inside the injection gracefully (skips execution)', async () => {
|
||||
const processor = new ShellProcessor('test-command');
|
||||
const prompt = 'This is weird: !{}';
|
||||
const prompt: PromptPipelineContent =
|
||||
createPromptPipelineContent('This is weird: !{}');
|
||||
|
||||
const result = await processor.process(prompt, context);
|
||||
|
||||
@@ -392,77 +424,14 @@ describe('ShellProcessor', () => {
|
||||
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/,
|
||||
);
|
||||
});
|
||||
expect(result).toEqual([{ text: 'This is weird: ' }]);
|
||||
});
|
||||
|
||||
describe('Error Reporting', () => {
|
||||
it('should append exit code and command name on failure', async () => {
|
||||
const processor = new ShellProcessor('test-command');
|
||||
const prompt = '!{cmd}';
|
||||
const prompt: PromptPipelineContent =
|
||||
createPromptPipelineContent('!{cmd}');
|
||||
mockShellExecute.mockReturnValue({
|
||||
result: Promise.resolve({
|
||||
...SUCCESS_RESULT,
|
||||
@@ -474,14 +443,17 @@ describe('ShellProcessor', () => {
|
||||
|
||||
const result = await processor.process(prompt, context);
|
||||
|
||||
expect(result).toBe(
|
||||
"some error output\n[Shell command 'cmd' exited with code 1]",
|
||||
);
|
||||
expect(result).toEqual([
|
||||
{
|
||||
text: "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}';
|
||||
const prompt: PromptPipelineContent =
|
||||
createPromptPipelineContent('!{cmd}');
|
||||
mockShellExecute.mockReturnValue({
|
||||
result: Promise.resolve({
|
||||
...SUCCESS_RESULT,
|
||||
@@ -494,14 +466,17 @@ describe('ShellProcessor', () => {
|
||||
|
||||
const result = await processor.process(prompt, context);
|
||||
|
||||
expect(result).toBe(
|
||||
"output\n[Shell command 'cmd' terminated by signal SIGTERM]",
|
||||
);
|
||||
expect(result).toEqual([
|
||||
{
|
||||
text: "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 prompt: PromptPipelineContent =
|
||||
createPromptPipelineContent('!{bad-command}');
|
||||
const spawnError = new Error('spawn EACCES');
|
||||
mockShellExecute.mockReturnValue({
|
||||
result: Promise.resolve({
|
||||
@@ -521,7 +496,9 @@ describe('ShellProcessor', () => {
|
||||
|
||||
it('should report abort status with command name if aborted', async () => {
|
||||
const processor = new ShellProcessor('test-command');
|
||||
const prompt = '!{long-running-command}';
|
||||
const prompt: PromptPipelineContent = createPromptPipelineContent(
|
||||
'!{long-running-command}',
|
||||
);
|
||||
const spawnError = new Error('Aborted');
|
||||
mockShellExecute.mockReturnValue({
|
||||
result: Promise.resolve({
|
||||
@@ -535,9 +512,11 @@ describe('ShellProcessor', () => {
|
||||
});
|
||||
|
||||
const result = await processor.process(prompt, context);
|
||||
expect(result).toBe(
|
||||
"partial output\n[Shell command 'long-running-command' aborted]",
|
||||
);
|
||||
expect(result).toEqual([
|
||||
{
|
||||
text: "partial output\n[Shell command 'long-running-command' aborted]",
|
||||
},
|
||||
]);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -551,29 +530,35 @@ describe('ShellProcessor', () => {
|
||||
|
||||
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 prompt: PromptPipelineContent = createPromptPipelineContent(
|
||||
'The user said: {{args}}',
|
||||
);
|
||||
|
||||
const result = await processor.process(prompt, context);
|
||||
|
||||
expect(result).toBe(`The user said: ${rawArgs}`);
|
||||
expect(result).toEqual([{ text: `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"}';
|
||||
const prompt: PromptPipelineContent = createPromptPipelineContent(
|
||||
'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`);
|
||||
expect(result).toEqual([{ text: `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}';
|
||||
const prompt: PromptPipelineContent = createPromptPipelineContent(
|
||||
'Command: !{grep {{args}} file.txt}',
|
||||
);
|
||||
mockShellExecute.mockReturnValue({
|
||||
result: Promise.resolve({ ...SUCCESS_RESULT, output: 'match found' }),
|
||||
});
|
||||
@@ -591,12 +576,14 @@ describe('ShellProcessor', () => {
|
||||
false,
|
||||
);
|
||||
|
||||
expect(result).toBe('Command: match found');
|
||||
expect(result).toEqual([{ text: '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}}}';
|
||||
const prompt: PromptPipelineContent = createPromptPipelineContent(
|
||||
'User "({{args}})" requested search: !{search {{args}}}',
|
||||
);
|
||||
mockShellExecute.mockReturnValue({
|
||||
result: Promise.resolve({ ...SUCCESS_RESULT, output: 'results' }),
|
||||
});
|
||||
@@ -613,12 +600,15 @@ describe('ShellProcessor', () => {
|
||||
false,
|
||||
);
|
||||
|
||||
expect(result).toBe(`User "(${rawArgs})" requested search: results`);
|
||||
expect(result).toEqual([
|
||||
{ text: `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 prompt: PromptPipelineContent =
|
||||
createPromptPipelineContent('!{rm {{args}}}');
|
||||
|
||||
const expectedEscapedArgs = getExpectedEscapedArgForPlatform(rawArgs);
|
||||
const expectedResolvedCommand = `rm ${expectedEscapedArgs}`;
|
||||
@@ -641,7 +631,8 @@ describe('ShellProcessor', () => {
|
||||
|
||||
it('should report the resolved command if a hard denial occurs', async () => {
|
||||
const processor = new ShellProcessor('test-command');
|
||||
const prompt = '!{rm {{args}}}';
|
||||
const prompt: PromptPipelineContent =
|
||||
createPromptPipelineContent('!{rm {{args}}}');
|
||||
const expectedEscapedArgs = getExpectedEscapedArgForPlatform(rawArgs);
|
||||
const expectedResolvedCommand = `rm ${expectedEscapedArgs}`;
|
||||
mockCheckCommandPermissions.mockReturnValue({
|
||||
@@ -661,7 +652,9 @@ describe('ShellProcessor', () => {
|
||||
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 prompt: PromptPipelineContent = createPromptPipelineContent(
|
||||
'Commit message: !{git commit -m {{args}}}',
|
||||
);
|
||||
|
||||
const expectedEscapedArgs =
|
||||
getExpectedEscapedArgForPlatform(multilineArgs);
|
||||
@@ -690,7 +683,8 @@ describe('ShellProcessor', () => {
|
||||
])('should safely escape args containing $name', async ({ input }) => {
|
||||
const processor = new ShellProcessor('test-command');
|
||||
context.invocation!.args = input;
|
||||
const prompt = '!{echo {{args}}}';
|
||||
const prompt: PromptPipelineContent =
|
||||
createPromptPipelineContent('!{echo {{args}}}');
|
||||
|
||||
const expectedEscapedArgs = getExpectedEscapedArgForPlatform(input);
|
||||
const expectedCommand = `echo ${expectedEscapedArgs}`;
|
||||
|
||||
@@ -10,14 +10,16 @@ import {
|
||||
escapeShellArg,
|
||||
getShellConfiguration,
|
||||
ShellExecutionService,
|
||||
flatMapTextParts,
|
||||
} from '@qwen-code/qwen-code-core';
|
||||
|
||||
import { CommandContext } from '../../ui/commands/types.js';
|
||||
import type { CommandContext } from '../../ui/commands/types.js';
|
||||
import type { IPromptProcessor, PromptPipelineContent } from './types.js';
|
||||
import {
|
||||
IPromptProcessor,
|
||||
SHELL_INJECTION_TRIGGER,
|
||||
SHORTHAND_ARGS_PLACEHOLDER,
|
||||
} from './types.js';
|
||||
import { extractInjections, type Injection } from './injectionParser.js';
|
||||
|
||||
export class ConfirmationRequiredError extends Error {
|
||||
constructor(
|
||||
@@ -30,15 +32,10 @@ export class ConfirmationRequiredError extends Error {
|
||||
}
|
||||
|
||||
/**
|
||||
* Represents a single detected shell injection site in the prompt.
|
||||
* Represents a single detected shell injection site in the prompt,
|
||||
* after resolution of arguments. Extends the base Injection interface.
|
||||
*/
|
||||
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;
|
||||
interface ResolvedShellInjection extends Injection {
|
||||
/** The command after {{args}} has been escaped and substituted. */
|
||||
resolvedCommand?: string;
|
||||
}
|
||||
@@ -56,11 +53,25 @@ interface ShellInjection {
|
||||
export class ShellProcessor implements IPromptProcessor {
|
||||
constructor(private readonly commandName: string) {}
|
||||
|
||||
async process(prompt: string, context: CommandContext): Promise<string> {
|
||||
async process(
|
||||
prompt: PromptPipelineContent,
|
||||
context: CommandContext,
|
||||
): Promise<PromptPipelineContent> {
|
||||
return flatMapTextParts(prompt, (text) =>
|
||||
this.processString(text, context),
|
||||
);
|
||||
}
|
||||
|
||||
private async processString(
|
||||
prompt: string,
|
||||
context: CommandContext,
|
||||
): Promise<PromptPipelineContent> {
|
||||
const userArgsRaw = context.invocation?.args || '';
|
||||
|
||||
if (!prompt.includes(SHELL_INJECTION_TRIGGER)) {
|
||||
return prompt.replaceAll(SHORTHAND_ARGS_PLACEHOLDER, userArgsRaw);
|
||||
return [
|
||||
{ text: prompt.replaceAll(SHORTHAND_ARGS_PLACEHOLDER, userArgsRaw) },
|
||||
];
|
||||
}
|
||||
|
||||
const config = context.services.config;
|
||||
@@ -71,26 +82,37 @@ export class ShellProcessor implements IPromptProcessor {
|
||||
}
|
||||
const { sessionShellAllowlist } = context.session;
|
||||
|
||||
const injections = this.extractInjections(prompt);
|
||||
const injections = extractInjections(
|
||||
prompt,
|
||||
SHELL_INJECTION_TRIGGER,
|
||||
this.commandName,
|
||||
);
|
||||
|
||||
// If extractInjections found no closed blocks (and didn't throw), treat as raw.
|
||||
if (injections.length === 0) {
|
||||
return prompt.replaceAll(SHORTHAND_ARGS_PLACEHOLDER, userArgsRaw);
|
||||
return [
|
||||
{ text: 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 resolvedInjections: ResolvedShellInjection[] = injections.map(
|
||||
(injection) => {
|
||||
const command = injection.content;
|
||||
|
||||
if (command === '') {
|
||||
return { ...injection, resolvedCommand: undefined };
|
||||
}
|
||||
|
||||
const resolvedCommand = command.replaceAll(
|
||||
SHORTHAND_ARGS_PLACEHOLDER,
|
||||
userArgsEscaped,
|
||||
);
|
||||
return { ...injection, resolvedCommand };
|
||||
},
|
||||
);
|
||||
|
||||
const commandsToConfirm = new Set<string>();
|
||||
for (const injection of resolvedInjections) {
|
||||
@@ -180,69 +202,6 @@ export class ShellProcessor implements IPromptProcessor {
|
||||
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;
|
||||
return [{ text: processedPrompt }];
|
||||
}
|
||||
}
|
||||
|
||||
@@ -4,7 +4,13 @@
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { CommandContext } from '../../ui/commands/types.js';
|
||||
import type { CommandContext } from '../../ui/commands/types.js';
|
||||
import type { PartUnion } from '@google/genai';
|
||||
|
||||
/**
|
||||
* Defines the input/output type for prompt processors.
|
||||
*/
|
||||
export type PromptPipelineContent = PartUnion[];
|
||||
|
||||
/**
|
||||
* Defines the interface for a prompt processor, a module that can transform
|
||||
@@ -13,12 +19,8 @@ import { CommandContext } from '../../ui/commands/types.js';
|
||||
*/
|
||||
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.
|
||||
* Processes a prompt input (which may contain text and multi-modal parts),
|
||||
* applying a specific transformation as part of a pipeline.
|
||||
*
|
||||
* @param prompt The current state of the prompt string. This may have been
|
||||
* modified by previous processors in the pipeline.
|
||||
@@ -28,7 +30,10 @@ export interface IPromptProcessor {
|
||||
* @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>;
|
||||
process(
|
||||
prompt: PromptPipelineContent,
|
||||
context: CommandContext,
|
||||
): Promise<PromptPipelineContent>;
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -42,3 +47,8 @@ export const SHORTHAND_ARGS_PLACEHOLDER = '{{args}}';
|
||||
* The trigger string for shell command injection in custom commands.
|
||||
*/
|
||||
export const SHELL_INJECTION_TRIGGER = '!{';
|
||||
|
||||
/**
|
||||
* The trigger string for at file injection in custom commands.
|
||||
*/
|
||||
export const AT_FILE_INJECTION_TRIGGER = '@{';
|
||||
|
||||
@@ -4,7 +4,7 @@
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { SlashCommand } from '../ui/commands/types.js';
|
||||
import type { SlashCommand } from '../ui/commands/types.js';
|
||||
|
||||
/**
|
||||
* Defines the contract for any class that can load and provide slash commands.
|
||||
|
||||
Reference in New Issue
Block a user