feat: Add Shell Command Execution to Custom Commands (#4917)

3 files changed, 411 insertions, 0 deletions

diff --git a/packages/cli/src/services/prompt-processors/shellProcessor.test.ts b/packages/cli/src/services/prompt-processors/shellProcessor.test.ts
new file mode 100644
index 00000000..a2883923
--- /dev/null
+++ b/packages/cli/src/services/prompt-processors/shellProcessor.test.ts
@@ -0,0 +1,300 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { vi, describe, it, expect, beforeEach } from 'vitest';
+import { ConfirmationRequiredError, ShellProcessor } from './shellProcessor.js';
+import { createMockCommandContext } from '../../test-utils/mockCommandContext.js';
+import { CommandContext } from '../../ui/commands/types.js';
+import { Config } from '@google/gemini-cli-core';
+
+const mockCheckCommandPermissions = vi.hoisted(() => vi.fn());
+const mockShellExecute = vi.hoisted(() => vi.fn());
+
+vi.mock('@google/gemini-cli-core', async (importOriginal) => {
+  const original = await importOriginal<object>();
+  return {
+    ...original,
+    checkCommandPermissions: mockCheckCommandPermissions,
+    ShellExecutionService: {
+      execute: mockShellExecute,
+    },
+  };
+});
+
+describe('ShellProcessor', () => {
+  let context: CommandContext;
+  let mockConfig: Partial<Config>;
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+
+    mockConfig = {
+      getTargetDir: vi.fn().mockReturnValue('/test/dir'),
+    };
+
+    context = createMockCommandContext({
+      services: {
+        config: mockConfig as Config,
+      },
+      session: {
+        sessionShellAllowlist: new Set(),
+      },
+    });
+
+    mockShellExecute.mockReturnValue({
+      result: Promise.resolve({
+        output: 'default shell output',
+      }),
+    });
+    mockCheckCommandPermissions.mockReturnValue({
+      allAllowed: true,
+      disallowedCommands: [],
+    });
+  });
+
+  it('should not change the prompt if no shell injections are present', async () => {
+    const processor = new ShellProcessor('test-command');
+    const prompt = 'This is a simple prompt with no injections.';
+    const result = await processor.process(prompt, context);
+    expect(result).toBe(prompt);
+    expect(mockShellExecute).not.toHaveBeenCalled();
+  });
+
+  it('should process a single valid shell injection if allowed', async () => {
+    const processor = new ShellProcessor('test-command');
+    const prompt = 'The current status is: !{git status}';
+    mockCheckCommandPermissions.mockReturnValue({
+      allAllowed: true,
+      disallowedCommands: [],
+    });
+    mockShellExecute.mockReturnValue({
+      result: Promise.resolve({ output: 'On branch main' }),
+    });
+
+    const result = await processor.process(prompt, context);
+
+    expect(mockCheckCommandPermissions).toHaveBeenCalledWith(
+      'git status',
+      expect.any(Object),
+      context.session.sessionShellAllowlist,
+    );
+    expect(mockShellExecute).toHaveBeenCalledWith(
+      'git status',
+      expect.any(String),
+      expect.any(Function),
+      expect.any(Object),
+    );
+    expect(result).toBe('The current status is: On branch main');
+  });
+
+  it('should process multiple valid shell injections if all are allowed', async () => {
+    const processor = new ShellProcessor('test-command');
+    const prompt = '!{git status} in !{pwd}';
+    mockCheckCommandPermissions.mockReturnValue({
+      allAllowed: true,
+      disallowedCommands: [],
+    });
+
+    mockShellExecute
+      .mockReturnValueOnce({
+        result: Promise.resolve({ output: 'On branch main' }),
+      })
+      .mockReturnValueOnce({
+        result: Promise.resolve({ output: '/usr/home' }),
+      });
+
+    const result = await processor.process(prompt, context);
+
+    expect(mockCheckCommandPermissions).toHaveBeenCalledTimes(2);
+    expect(mockShellExecute).toHaveBeenCalledTimes(2);
+    expect(result).toBe('On branch main in /usr/home');
+  });
+
+  it('should throw ConfirmationRequiredError if a command is not allowed', async () => {
+    const processor = new ShellProcessor('test-command');
+    const prompt = 'Do something dangerous: !{rm -rf /}';
+    mockCheckCommandPermissions.mockReturnValue({
+      allAllowed: false,
+      disallowedCommands: ['rm -rf /'],
+    });
+
+    await expect(processor.process(prompt, context)).rejects.toThrow(
+      ConfirmationRequiredError,
+    );
+  });
+
+  it('should throw ConfirmationRequiredError with the correct command', async () => {
+    const processor = new ShellProcessor('test-command');
+    const prompt = 'Do something dangerous: !{rm -rf /}';
+    mockCheckCommandPermissions.mockReturnValue({
+      allAllowed: false,
+      disallowedCommands: ['rm -rf /'],
+    });
+
+    try {
+      await processor.process(prompt, context);
+      // Fail if it doesn't throw
+      expect(true).toBe(false);
+    } catch (e) {
+      expect(e).toBeInstanceOf(ConfirmationRequiredError);
+      if (e instanceof ConfirmationRequiredError) {
+        expect(e.commandsToConfirm).toEqual(['rm -rf /']);
+      }
+    }
+
+    expect(mockShellExecute).not.toHaveBeenCalled();
+  });
+
+  it('should throw ConfirmationRequiredError with multiple commands if multiple are disallowed', async () => {
+    const processor = new ShellProcessor('test-command');
+    const prompt = '!{cmd1} and !{cmd2}';
+    mockCheckCommandPermissions.mockImplementation((cmd) => {
+      if (cmd === 'cmd1') {
+        return { allAllowed: false, disallowedCommands: ['cmd1'] };
+      }
+      if (cmd === 'cmd2') {
+        return { allAllowed: false, disallowedCommands: ['cmd2'] };
+      }
+      return { allAllowed: true, disallowedCommands: [] };
+    });
+
+    try {
+      await processor.process(prompt, context);
+      // Fail if it doesn't throw
+      expect(true).toBe(false);
+    } catch (e) {
+      expect(e).toBeInstanceOf(ConfirmationRequiredError);
+      if (e instanceof ConfirmationRequiredError) {
+        expect(e.commandsToConfirm).toEqual(['cmd1', 'cmd2']);
+      }
+    }
+  });
+
+  it('should not execute any commands if at least one requires confirmation', async () => {
+    const processor = new ShellProcessor('test-command');
+    const prompt = 'First: !{echo "hello"}, Second: !{rm -rf /}';
+
+    mockCheckCommandPermissions.mockImplementation((cmd) => {
+      if (cmd.includes('rm')) {
+        return { allAllowed: false, disallowedCommands: [cmd] };
+      }
+      return { allAllowed: true, disallowedCommands: [] };
+    });
+
+    await expect(processor.process(prompt, context)).rejects.toThrow(
+      ConfirmationRequiredError,
+    );
+
+    // Ensure no commands were executed because the pipeline was halted.
+    expect(mockShellExecute).not.toHaveBeenCalled();
+  });
+
+  it('should only request confirmation for disallowed commands in a mixed prompt', async () => {
+    const processor = new ShellProcessor('test-command');
+    const prompt = 'Allowed: !{ls -l}, Disallowed: !{rm -rf /}';
+
+    mockCheckCommandPermissions.mockImplementation((cmd) => ({
+      allAllowed: !cmd.includes('rm'),
+      disallowedCommands: cmd.includes('rm') ? [cmd] : [],
+    }));
+
+    try {
+      await processor.process(prompt, context);
+      expect.fail('Should have thrown ConfirmationRequiredError');
+    } catch (e) {
+      expect(e).toBeInstanceOf(ConfirmationRequiredError);
+      if (e instanceof ConfirmationRequiredError) {
+        expect(e.commandsToConfirm).toEqual(['rm -rf /']);
+      }
+    }
+  });
+
+  it('should execute all commands if they are on the session allowlist', async () => {
+    const processor = new ShellProcessor('test-command');
+    const prompt = 'Run !{cmd1} and !{cmd2}';
+
+    // Add commands to the session allowlist
+    context.session.sessionShellAllowlist = new Set(['cmd1', 'cmd2']);
+
+    // checkCommandPermissions should now pass for these
+    mockCheckCommandPermissions.mockReturnValue({
+      allAllowed: true,
+      disallowedCommands: [],
+    });
+
+    mockShellExecute
+      .mockReturnValueOnce({ result: Promise.resolve({ output: 'output1' }) })
+      .mockReturnValueOnce({ result: Promise.resolve({ output: 'output2' }) });
+
+    const result = await processor.process(prompt, context);
+
+    expect(mockCheckCommandPermissions).toHaveBeenCalledWith(
+      'cmd1',
+      expect.any(Object),
+      context.session.sessionShellAllowlist,
+    );
+    expect(mockCheckCommandPermissions).toHaveBeenCalledWith(
+      'cmd2',
+      expect.any(Object),
+      context.session.sessionShellAllowlist,
+    );
+    expect(mockShellExecute).toHaveBeenCalledTimes(2);
+    expect(result).toBe('Run output1 and output2');
+  });
+
+  it('should trim whitespace from the command inside the injection', async () => {
+    const processor = new ShellProcessor('test-command');
+    const prompt = 'Files: !{  ls -l  }';
+    mockCheckCommandPermissions.mockReturnValue({
+      allAllowed: true,
+      disallowedCommands: [],
+    });
+    mockShellExecute.mockReturnValue({
+      result: Promise.resolve({ output: 'total 0' }),
+    });
+
+    await processor.process(prompt, context);
+
+    expect(mockCheckCommandPermissions).toHaveBeenCalledWith(
+      'ls -l', // Verifies that the command was trimmed
+      expect.any(Object),
+      context.session.sessionShellAllowlist,
+    );
+    expect(mockShellExecute).toHaveBeenCalledWith(
+      'ls -l',
+      expect.any(String),
+      expect.any(Function),
+      expect.any(Object),
+    );
+  });
+
+  it('should handle an empty command inside the injection gracefully', async () => {
+    const processor = new ShellProcessor('test-command');
+    const prompt = 'This is weird: !{}';
+    mockCheckCommandPermissions.mockReturnValue({
+      allAllowed: true,
+      disallowedCommands: [],
+    });
+    mockShellExecute.mockReturnValue({
+      result: Promise.resolve({ output: 'empty output' }),
+    });
+
+    const result = await processor.process(prompt, context);
+
+    expect(mockCheckCommandPermissions).toHaveBeenCalledWith(
+      '',
+      expect.any(Object),
+      context.session.sessionShellAllowlist,
+    );
+    expect(mockShellExecute).toHaveBeenCalledWith(
+      '',
+      expect.any(String),
+      expect.any(Function),
+      expect.any(Object),
+    );
+    expect(result).toBe('This is weird: empty output');
+  });
+});
diff --git a/packages/cli/src/services/prompt-processors/shellProcessor.ts b/packages/cli/src/services/prompt-processors/shellProcessor.ts
new file mode 100644
index 00000000..bf811d66
--- /dev/null
+++ b/packages/cli/src/services/prompt-processors/shellProcessor.ts
@@ -0,0 +1,106 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import {
+  checkCommandPermissions,
+  ShellExecutionService,
+} from '@google/gemini-cli-core';
+
+import { CommandContext } from '../../ui/commands/types.js';
+import { IPromptProcessor } from './types.js';
+
+export class ConfirmationRequiredError extends Error {
+  constructor(
+    message: string,
+    public commandsToConfirm: string[],
+  ) {
+    super(message);
+    this.name = 'ConfirmationRequiredError';
+  }
+}
+
+/**
+ * Finds all instances of shell command injections (`!{...}`) in a prompt,
+ * executes them, and replaces the injection site with the command's output.
+ *
+ * This processor ensures that only allowlisted commands are executed. If a
+ * disallowed command is found, it halts execution and reports an error.
+ */
+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 matches = [...prompt.matchAll(ShellProcessor.SHELL_INJECTION_REGEX)];
+    if (matches.length === 0) {
+      return prompt; // No shell commands, nothing to do.
+    }
+
+    // Discover all commands and check permissions.
+    for (const match of matches) {
+      const command = match[1].trim();
+      const { allAllowed, disallowedCommands, blockReason, isHardDenial } =
+        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.'}`,
+          );
+        }
+
+        // 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.
+    if (commandsToConfirm.size > 0) {
+      throw new ConfirmationRequiredError(
+        'Shell command confirmation required',
+        Array.from(commandsToConfirm),
+      );
+    }
+
+    // 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.
+      );
+
+      const executionResult = await result;
+      processedPrompt = processedPrompt.replace(
+        fullMatch,
+        executionResult.output,
+      );
+    }
+
+    return processedPrompt;
+  }
+}
diff --git a/packages/cli/src/services/prompt-processors/types.ts b/packages/cli/src/services/prompt-processors/types.ts
index 2ca61062..2653d2b7 100644
--- a/packages/cli/src/services/prompt-processors/types.ts
+++ b/packages/cli/src/services/prompt-processors/types.ts
@@ -35,3 +35,8 @@ export interface IPromptProcessor {
  * The placeholder string for shorthand argument injection in custom commands.
  */
 export const SHORTHAND_ARGS_PLACEHOLDER = '{{args}}';
+
+/**
+ * The trigger string for shell command injection in custom commands.
+ */
+export const SHELL_INJECTION_TRIGGER = '!{';