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

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

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