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

1 files changed, 34 insertions, 0 deletions

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