feat(commands): Implement argument handling for custom commands via a prompt pipeline (#4702)

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);
+      }
+    });
+  });
 });

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<SubmitPromptActionReturn> => {
+        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,
+        };
+      },
     };
   }
 }

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

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<string> {
+    return prompt.replaceAll(
+      SHORTHAND_ARGS_PLACEHOLDER,
+      context.invocation!.args,
+    );
+  }
+}
+
+/**
+ * Appends the user's full command invocation to the prompt if arguments are
+ * provided, allowing the model to perform its own argument parsing.
+ */
+export class DefaultArgumentProcessor implements IPromptProcessor {
+  async process(prompt: string, context: CommandContext): Promise<string> {
+    if (context.invocation!.args) {
+      return `${prompt}\n\n${context.invocation!.raw}`;
+    }
+    return prompt;
+  }
+}

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<string>;
+}
+
+/**
+ * The placeholder string for shorthand argument injection in custom commands.
+ */
+export const SHORTHAND_ARGS_PLACEHOLDER = '{{args}}';

packages/cli/src/test-utils/mockCommandContext.ts

@@ -28,6 +28,11 @@ export const createMockCommandContext = (
   overrides: DeepPartial<CommandContext> = {},
 ): CommandContext => {
   const defaultMocks: CommandContext = {
+    invocation: {
+      raw: '',
+      name: '',
+      args: '',
+    },
     services: {
       config: null,
       settings: { merged: {} } as LoadedSettings,

packages/cli/src/ui/commands/types.ts

@@ -14,6 +14,15 @@ import { SessionStatsState } from '../contexts/SessionContext.js';
 
 // Grouped dependencies for clarity and easier mocking
 export interface CommandContext {
+  // Invocation properties for when commands are called.
+  invocation?: {
+    /** The raw, untrimmed input string from the user. */
+    raw: string;
+    /** The primary name of the command that was matched. */
+    name: string;
+    /** The arguments string that follows the command name. */
+    args: string;
+  };
   // Core services and configuration
   services: {
     // TODO(abhipatel12): Ensure that config is never null.
@@ -132,7 +141,7 @@ export interface SlashCommand {
   // The action to run. Optional for parent commands that only group sub-commands.
   action?: (
     context: CommandContext,
-    args: string,
+    args: string, // TODO: Remove args. CommandContext now contains the complete invocation.
   ) =>
     | void
     | SlashCommandActionReturn

packages/cli/src/ui/hooks/slashCommandProcessor.ts

@@ -238,7 +238,18 @@ export const useSlashCommandProcessor = (
         const args = parts.slice(pathIndex).join(' ');
 
         if (commandToExecute.action) {
-          const result = await commandToExecute.action(commandContext, args);
+          const fullCommandContext: CommandContext = {
+            ...commandContext,
+            invocation: {
+              raw: trimmed,
+              name: commandToExecute.name,
+              args,
+            },
+          };
+          const result = await commandToExecute.action(
+            fullCommandContext,
+            args,
+          );
 
           if (result) {
             switch (result.type) {
@@ -288,9 +299,9 @@ export const useSlashCommandProcessor = (
                 await config
                   ?.getGeminiClient()
                   ?.setHistory(result.clientHistory);
-                commandContext.ui.clear();
+                fullCommandContext.ui.clear();
                 result.history.forEach((item, index) => {
-                  commandContext.ui.addItem(item, index);
+                  fullCommandContext.ui.addItem(item, index);
                 });
                 return { type: 'handled' };
               }