abbyswag/qwen-code
1
0
Fork 0
mirror of https://github.com/QwenLM/qwen-code.git synced 2025-12-19 09:33:53 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(commands): Enable @file processing in TOML commands (#6716)

Browse Source
This commit is contained in:
Abhi
2025-08-27 23:22:21 -04:00
committed by GitHub
parent 529c2649b8
commit bfdddcbd99
27 changed files with 1836 additions and 331 deletions
Download Patch File Download Diff File Expand all files Collapse all files

238
packages/cli/src/services/FileCommandLoader.test.ts
View File

@@ -14,6 +14,7 @@ import { createMockCommandContext } from '../test-utils/mockCommandContext.js';
import {
SHELL_INJECTION_TRIGGER,
SHORTHAND_ARGS_PLACEHOLDER,
type PromptPipelineContent,
} from './prompt-processors/types.js';
import {
ConfirmationRequiredError,
@@ -21,8 +22,15 @@ import {
} from './prompt-processors/shellProcessor.js';
import { DefaultArgumentProcessor } from './prompt-processors/argumentProcessor.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,
@@ -68,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(() => {
@@ -110,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');
}
@@ -203,7 +224,7 @@ describe('FileCommandLoader', () => {
const mockConfig = {
getProjectRoot: vi.fn(() => '/path/to/project'),
getExtensions: vi.fn(() => []),
} as Config;
} as unknown as Config;
const loader = new FileCommandLoader(mockConfig);
const commands = await loader.loadCommands(signal);
expect(commands).toHaveLength(1);
@@ -246,7 +267,7 @@ describe('FileCommandLoader', () => {
const mockConfig = {
getProjectRoot: vi.fn(() => process.cwd()),
getExtensions: vi.fn(() => []),
} as Config;
} as unknown as Config;
const loader = new FileCommandLoader(mockConfig);
const commands = await loader.loadCommands(signal);
@@ -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,7 +298,7 @@ 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');
}
@@ -446,6 +467,54 @@ 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', () => {
@@ -487,7 +556,7 @@ describe('FileCommandLoader', () => {
path: extensionDir,
},
]),
} as Config;
} as unknown as Config;
const loader = new FileCommandLoader(mockConfig);
const commands = await loader.loadCommands(signal);
@@ -538,7 +607,7 @@ describe('FileCommandLoader', () => {
path: extensionDir,
},
]),
} as Config;
} as unknown as Config;
const loader = new FileCommandLoader(mockConfig);
const commands = await loader.loadCommands(signal);
@@ -559,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');
@@ -576,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');
@@ -594,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' }]);
}
});
@@ -645,7 +714,7 @@ describe('FileCommandLoader', () => {
path: extensionDir2,
},
]),
} as Config;
} as unknown as Config;
const loader = new FileCommandLoader(mockConfig);
const commands = await loader.loadCommands(signal);
@@ -681,7 +750,7 @@ describe('FileCommandLoader', () => {
path: extensionDir,
},
]),
} as Config;
} as unknown as Config;
const loader = new FileCommandLoader(mockConfig);
const commands = await loader.loadCommands(signal);
expect(commands).toHaveLength(0);
@@ -713,7 +782,7 @@ describe('FileCommandLoader', () => {
getExtensions: vi.fn(() => [
{ name: 'a', version: '1.0.0', isActive: true, path: extensionDir },
]),
} as Config;
} as unknown as Config;
const loader = new FileCommandLoader(mockConfig);
const commands = await loader.loadCommands(signal);
@@ -737,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');
}
@@ -771,7 +842,9 @@ 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' },
]);
}
});
});
@@ -805,7 +878,7 @@ 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 }]);
}
});
});
@@ -859,7 +932,7 @@ describe('FileCommandLoader', () => {
'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);
@@ -875,7 +948,7 @@ 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' }]);
}
});
@@ -938,23 +1011,36 @@ describe('FileCommandLoader', () => {
),
).rejects.toThrow('Something else went wrong');
});
it('assembles the processor pipeline in the correct order (Shell -> Default)', async () => {
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(
@@ -972,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' },
]);
}
});
});
});

39
packages/cli/src/services/FileCommandLoader.ts
View File

@@ -19,15 +19,20 @@ import type {
} from '../ui/commands/types.js';
import { CommandKind } from '../ui/commands/types.js';
import { DefaultArgumentProcessor } from './prompt-processors/argumentProcessor.js';
import type { IPromptProcessor } from './prompt-processors/types.js';
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;
@@ -224,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());
}
@@ -253,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

10
packages/cli/src/services/prompt-processors/argumentProcessor.test.ts
View File

@@ -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.' }]);
});
});
});

12
packages/cli/src/services/prompt-processors/argumentProcessor.ts
View File

@@ -4,7 +4,8 @@
* SPDX-License-Identifier: Apache-2.0
*/
import type { IPromptProcessor } from './types.js';
import { appendToLastTextPart } from '@google/gemini-cli-core';
import type { IPromptProcessor, PromptPipelineContent } from './types.js';
import type { CommandContext } from '../../ui/commands/types.js';
/**
@@ -14,9 +15,12 @@ import type { 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;
}

221
packages/cli/src/services/prompt-processors/atFileProcessor.test.ts Normal file
View File

@@ -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 '@google/gemini-cli-core';
import type { PartUnion } from '@google/genai';
// Mock the core dependency
const mockReadPathFromWorkspace = vi.hoisted(() => vi.fn());
vi.mock('@google/gemini-cli-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();
});
});
});

96
packages/cli/src/services/prompt-processors/atFileProcessor.ts Normal file
View File

@@ -0,0 +1,96 @@
/**
* @license
* Copyright 2025 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import {
flatMapTextParts,
readPathFromWorkspace,
} from '@google/gemini-cli-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;
});
}
}

223
packages/cli/src/services/prompt-processors/injectionParser.test.ts Normal file
View File

@@ -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/,
);
});
});
});

89
packages/cli/src/services/prompt-processors/injectionParser.ts Normal file
View File

@@ -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;
}

215
packages/cli/src/services/prompt-processors/shellProcessor.test.ts
View File

@@ -12,10 +12,11 @@ import type { Config } from '@google/gemini-cli-core';
import { ApprovalMode } from '@google/gemini-cli-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();
@@ -32,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());
@@ -93,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,
@@ -107,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: [],
@@ -138,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: [],
@@ -164,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 /'],
@@ -182,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 /'],
@@ -203,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'],
@@ -228,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 /'],
@@ -250,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'] };
@@ -275,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')) {
@@ -294,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'),
@@ -314,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']);
@@ -346,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;
@@ -385,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);
@@ -393,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,
@@ -475,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,
@@ -495,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({
@@ -522,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({
@@ -536,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]",
},
]);
});
});
@@ -552,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' }),
});
@@ -592,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' }),
});
@@ -614,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}`;
@@ -642,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({
@@ -662,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);
@@ -691,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}`;

135
packages/cli/src/services/prompt-processors/shellProcessor.ts
View File

@@ -10,14 +10,16 @@ import {
escapeShellArg,
getShellConfiguration,
ShellExecutionService,
flatMapTextParts,
} from '@google/gemini-cli-core';
import type { CommandContext } from '../../ui/commands/types.js';
import type { IPromptProcessor } from './types.js';
import type { IPromptProcessor, PromptPipelineContent } from './types.js';
import {
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 }];
}
}

24
packages/cli/src/services/prompt-processors/types.ts
View File

@@ -5,6 +5,12 @@
*/
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 type { 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
packages/cli/src/ui/commands/types.ts
View File

@@ -5,7 +5,7 @@
*/
import { type ReactNode } from 'react';
import type { Content } from '@google/genai';
import type { Content, PartListUnion } from '@google/genai';
import type { HistoryItemWithoutId, HistoryItem } from '../types.js';
import type { Config, GitService, Logger } from '@google/gemini-cli-core';
import type { LoadedSettings } from '../../config/settings.js';
@@ -122,7 +122,7 @@ export interface LoadHistoryActionReturn {
*/
export interface SubmitPromptActionReturn {
type: 'submit_prompt';
content: string;
content: PartListUnion;
}
/**

8
packages/cli/src/ui/hooks/slashCommandProcessor.test.ts
View File

@@ -491,7 +491,7 @@ describe('useSlashCommandProcessor', () => {
description: 'A command from a file',
action: async () => ({
type: 'submit_prompt',
content: 'The actual prompt from the TOML file.',
content: [{ text: 'The actual prompt from the TOML file.' }],
}),
},
CommandKind.FILE,
@@ -507,7 +507,7 @@ describe('useSlashCommandProcessor', () => {
expect(actionResult).toEqual({
type: 'submit_prompt',
content: 'The actual prompt from the TOML file.',
content: [{ text: 'The actual prompt from the TOML file.' }],
});
expect(mockAddItem).toHaveBeenCalledWith(
@@ -523,7 +523,7 @@ describe('useSlashCommandProcessor', () => {
description: 'A command from mcp',
action: async () => ({
type: 'submit_prompt',
content: 'The actual prompt from the mcp command.',
content: [{ text: 'The actual prompt from the mcp command.' }],
}),
},
CommandKind.MCP_PROMPT,
@@ -539,7 +539,7 @@ describe('useSlashCommandProcessor', () => {
expect(actionResult).toEqual({
type: 'submit_prompt',
content: 'The actual prompt from the mcp command.',
content: [{ text: 'The actual prompt from the mcp command.' }],
});
expect(mockAddItem).toHaveBeenCalledWith(

3
packages/cli/src/ui/types.ts
View File

@@ -9,6 +9,7 @@ import type {
ToolCallConfirmationDetails,
ToolResultDisplay,
} from '@google/gemini-cli-core';
import type { PartListUnion } from '@google/genai';
// Only defining the state enum needed by the UI
export enum StreamingState {
@@ -239,7 +240,7 @@ export interface ConsoleMessageItem {
*/
export interface SubmitPromptResult {
type: 'submit_prompt';
content: string;
content: PartListUnion;
}
/**