From bbe95f1eaa8f5351c58e0866ba938415db7891e4 Mon Sep 17 00:00:00 2001 From: Abhi <43648792+abhipatel12@users.noreply.github.com> Date: Wed, 23 Jul 2025 16:11:23 -0400 Subject: feat(commands): Implement argument handling for custom commands via a prompt pipeline (#4702) --- .../cli/src/services/FileCommandLoader.test.ts | 90 +++++++++++++++++++- packages/cli/src/services/FileCommandLoader.ts | 52 ++++++++++-- .../prompt-processors/argumentProcessor.test.ts | 99 ++++++++++++++++++++++ .../prompt-processors/argumentProcessor.ts | 34 ++++++++ .../cli/src/services/prompt-processors/types.ts | 37 ++++++++ 5 files changed, 303 insertions(+), 9 deletions(-) create mode 100644 packages/cli/src/services/prompt-processors/argumentProcessor.test.ts create mode 100644 packages/cli/src/services/prompt-processors/argumentProcessor.ts create mode 100644 packages/cli/src/services/prompt-processors/types.ts (limited to 'packages/cli/src/services') diff --git a/packages/cli/src/services/FileCommandLoader.test.ts b/packages/cli/src/services/FileCommandLoader.test.ts index 518c9230..b86d36ac 100644 --- a/packages/cli/src/services/FileCommandLoader.test.ts +++ b/packages/cli/src/services/FileCommandLoader.test.ts @@ -14,8 +14,6 @@ import mock from 'mock-fs'; import { assert } from 'vitest'; import { createMockCommandContext } from '../test-utils/mockCommandContext.js'; -const mockContext = createMockCommandContext(); - describe('FileCommandLoader', () => { const signal: AbortSignal = new AbortController().signal; @@ -39,7 +37,16 @@ describe('FileCommandLoader', () => { expect(command).toBeDefined(); expect(command.name).toBe('test'); - const result = await command.action?.(mockContext, ''); + const result = await command.action?.( + createMockCommandContext({ + invocation: { + raw: '/test', + name: 'test', + args: '', + }, + }), + '', + ); if (result?.type === 'submit_prompt') { expect(result.content).toBe('This is a test prompt'); } else { @@ -122,7 +129,16 @@ describe('FileCommandLoader', () => { const command = commands[0]; expect(command).toBeDefined(); - const result = await command.action?.(mockContext, ''); + const result = await command.action?.( + createMockCommandContext({ + invocation: { + raw: '/test', + name: 'test', + args: '', + }, + }), + '', + ); if (result?.type === 'submit_prompt') { expect(result.content).toBe('Project prompt'); } else { @@ -232,4 +248,70 @@ describe('FileCommandLoader', () => { // Verify that the ':' in the filename was replaced with an '_' expect(command.name).toBe('legacy_command'); }); + + describe('Shorthand Argument Processor Integration', () => { + it('correctly processes a command with {{args}}', async () => { + const userCommandsDir = getUserCommandsDir(); + mock({ + [userCommandsDir]: { + 'shorthand.toml': + 'prompt = "The user wants to: {{args}}"\ndescription = "Shorthand test"', + }, + }); + + const loader = new FileCommandLoader(null as unknown as Config); + const commands = await loader.loadCommands(signal); + const command = commands.find((c) => c.name === 'shorthand'); + expect(command).toBeDefined(); + + const result = await command!.action?.( + createMockCommandContext({ + invocation: { + raw: '/shorthand do something cool', + name: 'shorthand', + args: 'do something cool', + }, + }), + 'do something cool', + ); + expect(result?.type).toBe('submit_prompt'); + if (result?.type === 'submit_prompt') { + expect(result.content).toBe('The user wants to: do something cool'); + } + }); + }); + + describe('Default Argument Processor Integration', () => { + it('correctly processes a command without {{args}}', async () => { + const userCommandsDir = getUserCommandsDir(); + mock({ + [userCommandsDir]: { + 'model_led.toml': + 'prompt = "This is the instruction."\ndescription = "Default processor test"', + }, + }); + + const loader = new FileCommandLoader(null as unknown as Config); + const commands = await loader.loadCommands(signal); + const command = commands.find((c) => c.name === 'model_led'); + expect(command).toBeDefined(); + + const result = await command!.action?.( + createMockCommandContext({ + invocation: { + raw: '/model_led 1.2.0 added "a feature"', + name: 'model_led', + args: '1.2.0 added "a feature"', + }, + }), + '1.2.0 added "a feature"', + ); + expect(result?.type).toBe('submit_prompt'); + 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); + } + }); + }); }); diff --git a/packages/cli/src/services/FileCommandLoader.ts b/packages/cli/src/services/FileCommandLoader.ts index 1b96cb35..994762c1 100644 --- a/packages/cli/src/services/FileCommandLoader.ts +++ b/packages/cli/src/services/FileCommandLoader.ts @@ -15,7 +15,20 @@ import { getUserCommandsDir, } from '@google/gemini-cli-core'; import { ICommandLoader } from './types.js'; -import { CommandKind, SlashCommand } from '../ui/commands/types.js'; +import { + CommandContext, + CommandKind, + SlashCommand, + SubmitPromptActionReturn, +} from '../ui/commands/types.js'; +import { + DefaultArgumentProcessor, + ShorthandArgumentProcessor, +} from './prompt-processors/argumentProcessor.js'; +import { + IPromptProcessor, + SHORTHAND_ARGS_PLACEHOLDER, +} from './prompt-processors/types.js'; /** * Defines the Zod schema for a command definition file. This serves as the @@ -156,16 +169,45 @@ export class FileCommandLoader implements ICommandLoader { .map((segment) => segment.replaceAll(':', '_')) .join(':'); + const processors: IPromptProcessor[] = []; + + // The presence of '{{args}}' is the switch that determines the behavior. + if (validDef.prompt.includes(SHORTHAND_ARGS_PLACEHOLDER)) { + processors.push(new ShorthandArgumentProcessor()); + } else { + processors.push(new DefaultArgumentProcessor()); + } + return { name: commandName, description: validDef.description || `Custom command from ${path.basename(filePath)}`, kind: CommandKind.FILE, - action: async () => ({ - type: 'submit_prompt', - content: validDef.prompt, - }), + action: async ( + context: CommandContext, + _args: string, + ): Promise => { + if (!context.invocation) { + console.error( + `[FileCommandLoader] Critical error: Command '${commandName}' was executed without invocation context.`, + ); + return { + type: 'submit_prompt', + content: validDef.prompt, // Fallback to unprocessed prompt + }; + } + + let processedPrompt = validDef.prompt; + for (const processor of processors) { + processedPrompt = await processor.process(processedPrompt, context); + } + + return { + type: 'submit_prompt', + content: processedPrompt, + }; + }, }; } } diff --git a/packages/cli/src/services/prompt-processors/argumentProcessor.test.ts b/packages/cli/src/services/prompt-processors/argumentProcessor.test.ts new file mode 100644 index 00000000..6af578a9 --- /dev/null +++ b/packages/cli/src/services/prompt-processors/argumentProcessor.test.ts @@ -0,0 +1,99 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { + DefaultArgumentProcessor, + ShorthandArgumentProcessor, +} from './argumentProcessor.js'; +import { createMockCommandContext } from '../../test-utils/mockCommandContext.js'; + +describe('Argument Processors', () => { + describe('ShorthandArgumentProcessor', () => { + const processor = new ShorthandArgumentProcessor(); + + it('should replace a single {{args}} instance', async () => { + const prompt = 'Refactor the following code: {{args}}'; + const context = createMockCommandContext({ + invocation: { + raw: '/refactor make it faster', + name: 'refactor', + args: 'make it faster', + }, + }); + const result = await processor.process(prompt, context); + expect(result).toBe('Refactor the following code: make it faster'); + }); + + it('should replace multiple {{args}} instances', async () => { + const prompt = 'User said: {{args}}. I repeat: {{args}}!'; + const context = createMockCommandContext({ + invocation: { + raw: '/repeat hello world', + name: 'repeat', + args: 'hello world', + }, + }); + const result = await processor.process(prompt, context); + expect(result).toBe('User said: hello world. I repeat: hello world!'); + }); + + it('should handle an empty args string', async () => { + const prompt = 'The user provided no input: {{args}}.'; + const context = createMockCommandContext({ + invocation: { + raw: '/input', + name: 'input', + args: '', + }, + }); + const result = await processor.process(prompt, context); + expect(result).toBe('The user provided no input: .'); + }); + + it('should not change the prompt if {{args}} is not present', async () => { + const prompt = 'This is a static prompt.'; + const context = createMockCommandContext({ + invocation: { + raw: '/static some arguments', + name: 'static', + args: 'some arguments', + }, + }); + const result = await processor.process(prompt, context); + expect(result).toBe('This is a static prompt.'); + }); + }); + + describe('DefaultArgumentProcessor', () => { + const processor = new DefaultArgumentProcessor(); + + it('should append the full command if args are provided', async () => { + const prompt = 'Parse the command.'; + const context = createMockCommandContext({ + invocation: { + raw: '/mycommand arg1 "arg two"', + name: 'mycommand', + args: 'arg1 "arg two"', + }, + }); + const result = await processor.process(prompt, context); + expect(result).toBe('Parse the command.\n\n/mycommand arg1 "arg two"'); + }); + + it('should NOT append the full command if no args are provided', async () => { + const prompt = 'Parse the command.'; + const context = createMockCommandContext({ + invocation: { + raw: '/mycommand', + name: 'mycommand', + args: '', + }, + }); + const result = await processor.process(prompt, context); + expect(result).toBe('Parse the command.'); + }); + }); +}); diff --git a/packages/cli/src/services/prompt-processors/argumentProcessor.ts b/packages/cli/src/services/prompt-processors/argumentProcessor.ts new file mode 100644 index 00000000..a7efeea9 --- /dev/null +++ b/packages/cli/src/services/prompt-processors/argumentProcessor.ts @@ -0,0 +1,34 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { IPromptProcessor, SHORTHAND_ARGS_PLACEHOLDER } from './types.js'; +import { CommandContext } from '../../ui/commands/types.js'; + +/** + * Replaces all instances of `{{args}}` in a prompt with the user-provided + * argument string. + */ +export class ShorthandArgumentProcessor implements IPromptProcessor { + async process(prompt: string, context: CommandContext): Promise { + return prompt.replaceAll( + SHORTHAND_ARGS_PLACEHOLDER, + context.invocation!.args, + ); + } +} + +/** + * Appends the user's full command invocation to the prompt if arguments are + * provided, allowing the model to perform its own argument parsing. + */ +export class DefaultArgumentProcessor implements IPromptProcessor { + async process(prompt: string, context: CommandContext): Promise { + if (context.invocation!.args) { + return `${prompt}\n\n${context.invocation!.raw}`; + } + return prompt; + } +} diff --git a/packages/cli/src/services/prompt-processors/types.ts b/packages/cli/src/services/prompt-processors/types.ts new file mode 100644 index 00000000..2ca61062 --- /dev/null +++ b/packages/cli/src/services/prompt-processors/types.ts @@ -0,0 +1,37 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { CommandContext } from '../../ui/commands/types.js'; + +/** + * Defines the interface for a prompt processor, a module that can transform + * a prompt string before it is sent to the model. Processors are chained + * together to create a processing pipeline. + */ +export interface IPromptProcessor { + /** + * Processes a prompt string, applying a specific transformation as part of a pipeline. + * + * Each processor in a command's pipeline receives the output of the previous + * processor. This method provides the full command context, allowing for + * complex transformations that may require access to invocation details, + * application services, or UI state. + * + * @param prompt The current state of the prompt string. This may have been + * modified by previous processors in the pipeline. + * @param context The full command context, providing access to invocation + * details (like `context.invocation.raw` and `context.invocation.args`), + * application services, and UI handlers. + * @returns A promise that resolves to the transformed prompt string, which + * will be passed to the next processor or, if it's the last one, sent to the model. + */ + process(prompt: string, context: CommandContext): Promise; +} + +/** + * The placeholder string for shorthand argument injection in custom commands. + */ +export const SHORTHAND_ARGS_PLACEHOLDER = '{{args}}'; -- cgit v1.2.3