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

3 files changed, 170 insertions, 0 deletions

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