feat(cli): Introduce arguments for shell execution in custom commands (#5966)

1 files changed, 192 insertions, 48 deletions

diff --git a/packages/cli/src/services/prompt-processors/shellProcessor.ts b/packages/cli/src/services/prompt-processors/shellProcessor.ts
index bf811d66..9531d57e 100644
--- a/packages/cli/src/services/prompt-processors/shellProcessor.ts
+++ b/packages/cli/src/services/prompt-processors/shellProcessor.ts
@@ -6,11 +6,17 @@
 
 import {
   checkCommandPermissions,
+  escapeShellArg,
+  getShellConfiguration,
   ShellExecutionService,
 } from '@google/gemini-cli-core';
 
 import { CommandContext } from '../../ui/commands/types.js';
-import { IPromptProcessor } from './types.js';
+import {
+  IPromptProcessor,
+  SHELL_INJECTION_TRIGGER,
+  SHORTHAND_ARGS_PLACEHOLDER,
+} from './types.js';
 
 export class ConfirmationRequiredError extends Error {
   constructor(
@@ -23,60 +29,89 @@ export class ConfirmationRequiredError extends Error {
 }
 
 /**
- * Finds all instances of shell command injections (`!{...}`) in a prompt,
- * executes them, and replaces the injection site with the command's output.
+ * Represents a single detected shell injection site in the prompt.
+ */
+interface ShellInjection {
+  /** The shell command extracted from within !{...}, trimmed. */
+  command: string;
+  /** The starting index of the injection (inclusive, points to '!'). */
+  startIndex: number;
+  /** The ending index of the injection (exclusive, points after '}'). */
+  endIndex: number;
+  /** The command after {{args}} has been escaped and substituted. */
+  resolvedCommand?: string;
+}
+
+/**
+ * Handles prompt interpolation, including shell command execution (`!{...}`)
+ * and context-aware argument injection (`{{args}}`).
  *
- * This processor ensures that only allowlisted commands are executed. If a
- * disallowed command is found, it halts execution and reports an error.
+ * This processor ensures that:
+ * 1. `{{args}}` outside `!{...}` are replaced with raw input.
+ * 2. `{{args}}` inside `!{...}` are replaced with shell-escaped input.
+ * 3. Shell commands are executed securely after argument substitution.
+ * 4. Parsing correctly handles nested braces.
  */
 export class ShellProcessor implements IPromptProcessor {
-  /**
-   * A regular expression to find all instances of `!{...}`. The inner
-   * capture group extracts the command itself.
-   */
-  private static readonly SHELL_INJECTION_REGEX = /!\{([^}]*)\}/g;
-
-  /**
-   * @param commandName The name of the custom command being executed, used
-   *   for logging and error messages.
-   */
   constructor(private readonly commandName: string) {}
 
   async process(prompt: string, context: CommandContext): Promise<string> {
-    const { config, sessionShellAllowlist } = {
-      ...context.services,
-      ...context.session,
-    };
-    const commandsToExecute: Array<{ fullMatch: string; command: string }> = [];
-    const commandsToConfirm = new Set<string>();
+    const userArgsRaw = context.invocation?.args || '';
 
-    const matches = [...prompt.matchAll(ShellProcessor.SHELL_INJECTION_REGEX)];
-    if (matches.length === 0) {
-      return prompt; // No shell commands, nothing to do.
+    if (!prompt.includes(SHELL_INJECTION_TRIGGER)) {
+      return prompt.replaceAll(SHORTHAND_ARGS_PLACEHOLDER, userArgsRaw);
     }
 
-    // Discover all commands and check permissions.
-    for (const match of matches) {
-      const command = match[1].trim();
+    const config = context.services.config;
+    if (!config) {
+      throw new Error(
+        `Security configuration not loaded. Cannot verify shell command permissions for '${this.commandName}'. Aborting.`,
+      );
+    }
+    const { sessionShellAllowlist } = context.session;
+
+    const injections = this.extractInjections(prompt);
+    // If extractInjections found no closed blocks (and didn't throw), treat as raw.
+    if (injections.length === 0) {
+      return prompt.replaceAll(SHORTHAND_ARGS_PLACEHOLDER, userArgsRaw);
+    }
+
+    const { shell } = getShellConfiguration();
+    const userArgsEscaped = escapeShellArg(userArgsRaw, shell);
+
+    const resolvedInjections = injections.map((injection) => {
+      if (injection.command === '') {
+        return injection;
+      }
+      // Replace {{args}} inside the command string with the escaped version.
+      const resolvedCommand = injection.command.replaceAll(
+        SHORTHAND_ARGS_PLACEHOLDER,
+        userArgsEscaped,
+      );
+      return { ...injection, resolvedCommand };
+    });
+
+    const commandsToConfirm = new Set<string>();
+    for (const injection of resolvedInjections) {
+      const command = injection.resolvedCommand;
+
+      if (!command) continue;
+
+      // Security check on the final, escaped command string.
       const { allAllowed, disallowedCommands, blockReason, isHardDenial } =
-        checkCommandPermissions(command, config!, sessionShellAllowlist);
+        checkCommandPermissions(command, config, sessionShellAllowlist);
 
       if (!allAllowed) {
-        // If it's a hard denial, this is a non-recoverable security error.
         if (isHardDenial) {
           throw new Error(
-            `${this.commandName} cannot be run. ${blockReason || 'A shell command in this custom command is explicitly blocked in your config settings.'}`,
+            `${this.commandName} cannot be run. Blocked command: "${command}". Reason: ${blockReason || 'Blocked by configuration.'}`,
           );
         }
-
-        // Add each soft denial disallowed command to the set for confirmation.
         disallowedCommands.forEach((uc) => commandsToConfirm.add(uc));
       }
-      commandsToExecute.push({ fullMatch: match[0], command });
     }
 
-    // If any commands require confirmation, throw a special error to halt the
-    // pipeline and trigger the UI flow.
+    // Handle confirmation requirements.
     if (commandsToConfirm.size > 0) {
       throw new ConfirmationRequiredError(
         'Shell command confirmation required',
@@ -84,23 +119,132 @@ export class ShellProcessor implements IPromptProcessor {
       );
     }
 
-    // Execute all commands (only runs if no confirmation was needed).
-    let processedPrompt = prompt;
-    for (const { fullMatch, command } of commandsToExecute) {
-      const { result } = ShellExecutionService.execute(
-        command,
-        config!.getTargetDir(),
-        () => {}, // No streaming needed.
-        new AbortController().signal, // For now, we don't support cancellation from here.
-      );
+    let processedPrompt = '';
+    let lastIndex = 0;
 
-      const executionResult = await result;
-      processedPrompt = processedPrompt.replace(
-        fullMatch,
-        executionResult.output,
+    for (const injection of resolvedInjections) {
+      // Append the text segment BEFORE the injection, substituting {{args}} with RAW input.
+      const segment = prompt.substring(lastIndex, injection.startIndex);
+      processedPrompt += segment.replaceAll(
+        SHORTHAND_ARGS_PLACEHOLDER,
+        userArgsRaw,
       );
+
+      // Execute the resolved command (which already has ESCAPED input).
+      if (injection.resolvedCommand) {
+        const { result } = ShellExecutionService.execute(
+          injection.resolvedCommand,
+          config.getTargetDir(),
+          () => {},
+          new AbortController().signal,
+        );
+
+        const executionResult = await result;
+
+        // Handle Spawn Errors
+        if (executionResult.error && !executionResult.aborted) {
+          throw new Error(
+            `Failed to start shell command in '${this.commandName}': ${executionResult.error.message}. Command: ${injection.resolvedCommand}`,
+          );
+        }
+
+        // Append the output, making stderr explicit for the model.
+        if (executionResult.stdout) {
+          processedPrompt += executionResult.stdout;
+        }
+        if (executionResult.stderr) {
+          if (executionResult.stdout) {
+            processedPrompt += '\n';
+          }
+          processedPrompt += `--- STDERR ---\n${executionResult.stderr}`;
+        }
+
+        // Append a status message if the command did not succeed.
+        if (executionResult.aborted) {
+          processedPrompt += `\n[Shell command '${injection.resolvedCommand}' aborted]`;
+        } else if (
+          executionResult.exitCode !== 0 &&
+          executionResult.exitCode !== null
+        ) {
+          processedPrompt += `\n[Shell command '${injection.resolvedCommand}' exited with code ${executionResult.exitCode}]`;
+        } else if (executionResult.signal !== null) {
+          processedPrompt += `\n[Shell command '${injection.resolvedCommand}' terminated by signal ${executionResult.signal}]`;
+        }
+      }
+
+      lastIndex = injection.endIndex;
     }
 
+    // Append the remaining text AFTER the last injection, substituting {{args}} with RAW input.
+    const finalSegment = prompt.substring(lastIndex);
+    processedPrompt += finalSegment.replaceAll(
+      SHORTHAND_ARGS_PLACEHOLDER,
+      userArgsRaw,
+    );
+
     return processedPrompt;
   }
+
+  /**
+   * Iteratively parses the prompt string to extract shell injections (!{...}),
+   * correctly handling nested braces within the command.
+   *
+   * @param prompt The prompt string to parse.
+   * @returns An array of extracted ShellInjection objects.
+   * @throws Error if an unclosed injection (`!{`) is found.
+   */
+  private extractInjections(prompt: string): ShellInjection[] {
+    const injections: ShellInjection[] = [];
+    let index = 0;
+
+    while (index < prompt.length) {
+      const startIndex = prompt.indexOf(SHELL_INJECTION_TRIGGER, index);
+
+      if (startIndex === -1) {
+        break;
+      }
+
+      let currentIndex = startIndex + SHELL_INJECTION_TRIGGER.length;
+      let braceCount = 1;
+      let foundEnd = false;
+
+      while (currentIndex < prompt.length) {
+        const char = prompt[currentIndex];
+
+        // We count literal braces. This parser does not interpret shell quoting/escaping.
+        if (char === '{') {
+          braceCount++;
+        } else if (char === '}') {
+          braceCount--;
+          if (braceCount === 0) {
+            const commandContent = prompt.substring(
+              startIndex + SHELL_INJECTION_TRIGGER.length,
+              currentIndex,
+            );
+            const endIndex = currentIndex + 1;
+
+            injections.push({
+              command: commandContent.trim(),
+              startIndex,
+              endIndex,
+            });
+
+            index = endIndex;
+            foundEnd = true;
+            break;
+          }
+        }
+        currentIndex++;
+      }
+
+      // Check if the inner loop finished without finding the closing brace.
+      if (!foundEnd) {
+        throw new Error(
+          `Invalid syntax in command '${this.commandName}': Unclosed shell injection starting at index ${startIndex} ('!{'). Ensure braces are balanced.`,
+        );
+      }
+    }
+
+    return injections;
+  }
 }