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

packages/cli/src/services/prompt-processors/shellProcessor.test.ts

@@ -4,11 +4,32 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { vi, describe, it, expect, beforeEach } from 'vitest';
+import { describe, it, expect, beforeEach, vi } 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';
+import os from 'os';
+import { quote } from 'shell-quote';
+
+// Helper function to determine the expected escaped string based on the current OS,
+// mirroring the logic in the actual `escapeShellArg` implementation. This makes
+// our tests robust and platform-agnostic.
+function getExpectedEscapedArgForPlatform(arg: string): string {
+  if (os.platform() === 'win32') {
+    const comSpec = (process.env.ComSpec || 'cmd.exe').toLowerCase();
+    const isPowerShell =
+      comSpec.endsWith('powershell.exe') || comSpec.endsWith('pwsh.exe');
+
+    if (isPowerShell) {
+      return `'${arg.replace(/'/g, "''")}'`;
+    } else {
+      return `"${arg.replace(/"/g, '""')}"`;
+    }
+  } else {
+    return quote([arg]);
+  }
+}
 
 const mockCheckCommandPermissions = vi.hoisted(() => vi.fn());
 const mockShellExecute = vi.hoisted(() => vi.fn());
@@ -24,6 +45,15 @@ vi.mock('@google/gemini-cli-core', async (importOriginal) => {
   };
 });
 
+const SUCCESS_RESULT = {
+  stdout: 'default shell output',
+  stderr: '',
+  exitCode: 0,
+  error: null,
+  aborted: false,
+  signal: null,
+};
+
 describe('ShellProcessor', () => {
   let context: CommandContext;
   let mockConfig: Partial<Config>;
@@ -36,6 +66,11 @@ describe('ShellProcessor', () => {
     };
 
     context = createMockCommandContext({
+      invocation: {
+        raw: '/cmd default args',
+        name: 'cmd',
+        args: 'default args',
+      },
       services: {
         config: mockConfig as Config,
       },
@@ -45,16 +80,29 @@ describe('ShellProcessor', () => {
     });
 
     mockShellExecute.mockReturnValue({
-      result: Promise.resolve({
-        output: 'default shell output',
-      }),
+      result: Promise.resolve(SUCCESS_RESULT),
     });
+
     mockCheckCommandPermissions.mockReturnValue({
       allAllowed: true,
       disallowedCommands: [],
     });
   });
 
+  it('should throw an error if config is missing', async () => {
+    const processor = new ShellProcessor('test-command');
+    const prompt = '!{ls}';
+    const contextWithoutConfig = createMockCommandContext({
+      services: {
+        config: null,
+      },
+    });
+
+    await expect(
+      processor.process(prompt, contextWithoutConfig),
+    ).rejects.toThrow(/Security configuration not loaded/);
+  });
+
   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.';
@@ -71,7 +119,7 @@ describe('ShellProcessor', () => {
       disallowedCommands: [],
     });
     mockShellExecute.mockReturnValue({
-      result: Promise.resolve({ output: 'On branch main' }),
+      result: Promise.resolve({ ...SUCCESS_RESULT, stdout: 'On branch main' }),
     });
 
     const result = await processor.process(prompt, context);
@@ -100,10 +148,13 @@ describe('ShellProcessor', () => {
 
     mockShellExecute
       .mockReturnValueOnce({
-        result: Promise.resolve({ output: 'On branch main' }),
+        result: Promise.resolve({
+          ...SUCCESS_RESULT,
+          stdout: 'On branch main',
+        }),
       })
       .mockReturnValueOnce({
-        result: Promise.resolve({ output: '/usr/home' }),
+        result: Promise.resolve({ ...SUCCESS_RESULT, stdout: '/usr/home' }),
       });
 
     const result = await processor.process(prompt, context);
@@ -226,8 +277,12 @@ describe('ShellProcessor', () => {
     });
 
     mockShellExecute
-      .mockReturnValueOnce({ result: Promise.resolve({ output: 'output1' }) })
-      .mockReturnValueOnce({ result: Promise.resolve({ output: 'output2' }) });
+      .mockReturnValueOnce({
+        result: Promise.resolve({ ...SUCCESS_RESULT, stdout: 'output1' }),
+      })
+      .mockReturnValueOnce({
+        result: Promise.resolve({ ...SUCCESS_RESULT, stdout: 'output2' }),
+      });
 
     const result = await processor.process(prompt, context);
 
@@ -245,56 +300,406 @@ describe('ShellProcessor', () => {
     expect(result).toBe('Run output1 and output2');
   });
 
-  it('should trim whitespace from the command inside the injection', async () => {
+  it('should trim whitespace from the command inside the injection before interpolation', async () => {
     const processor = new ShellProcessor('test-command');
-    const prompt = 'Files: !{  ls -l  }';
+    const prompt = 'Files: !{  ls {{args}} -l  }';
+
+    const rawArgs = context.invocation!.args;
+
+    const expectedEscapedArgs = getExpectedEscapedArgForPlatform(rawArgs);
+
+    const expectedCommand = `ls ${expectedEscapedArgs} -l`;
+
     mockCheckCommandPermissions.mockReturnValue({
       allAllowed: true,
       disallowedCommands: [],
     });
     mockShellExecute.mockReturnValue({
-      result: Promise.resolve({ output: 'total 0' }),
+      result: Promise.resolve({ ...SUCCESS_RESULT, stdout: 'total 0' }),
     });
 
     await processor.process(prompt, context);
 
     expect(mockCheckCommandPermissions).toHaveBeenCalledWith(
-      'ls -l', // Verifies that the command was trimmed
+      expectedCommand,
       expect.any(Object),
       context.session.sessionShellAllowlist,
     );
     expect(mockShellExecute).toHaveBeenCalledWith(
-      'ls -l',
+      expectedCommand,
       expect.any(String),
       expect.any(Function),
       expect.any(Object),
     );
   });
 
-  it('should handle an empty command inside the injection gracefully', async () => {
+  it('should handle an empty command inside the injection gracefully (skips execution)', 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');
+    expect(mockCheckCommandPermissions).not.toHaveBeenCalled();
+    expect(mockShellExecute).not.toHaveBeenCalled();
+
+    // It replaces !{} with an empty string.
+    expect(result).toBe('This is weird: ');
+  });
+
+  describe('Robust Parsing (Balanced Braces)', () => {
+    it('should correctly parse commands containing nested braces (e.g., awk)', async () => {
+      const processor = new ShellProcessor('test-command');
+      const command = "awk '{print $1}' file.txt";
+      const prompt = `Output: !{${command}}`;
+      mockShellExecute.mockReturnValue({
+        result: Promise.resolve({ ...SUCCESS_RESULT, stdout: 'result' }),
+      });
+
+      const result = await processor.process(prompt, context);
+
+      expect(mockCheckCommandPermissions).toHaveBeenCalledWith(
+        command,
+        expect.any(Object),
+        context.session.sessionShellAllowlist,
+      );
+      expect(mockShellExecute).toHaveBeenCalledWith(
+        command,
+        expect.any(String),
+        expect.any(Function),
+        expect.any(Object),
+      );
+      expect(result).toBe('Output: result');
+    });
+
+    it('should handle deeply nested braces correctly', async () => {
+      const processor = new ShellProcessor('test-command');
+      const command = "echo '{{a},{b}}'";
+      const prompt = `!{${command}}`;
+      mockShellExecute.mockReturnValue({
+        result: Promise.resolve({ ...SUCCESS_RESULT, stdout: '{{a},{b}}' }),
+      });
+
+      const result = await processor.process(prompt, context);
+      expect(mockShellExecute).toHaveBeenCalledWith(
+        command,
+        expect.any(String),
+        expect.any(Function),
+        expect.any(Object),
+      );
+      expect(result).toBe('{{a},{b}}');
+    });
+
+    it('should throw an error for unclosed shell injections', async () => {
+      const processor = new ShellProcessor('test-command');
+      const prompt = 'This prompt is broken: !{ls -l';
+
+      await expect(processor.process(prompt, context)).rejects.toThrow(
+        /Unclosed shell injection/,
+      );
+    });
+
+    it('should throw an error for unclosed nested braces', async () => {
+      const processor = new ShellProcessor('test-command');
+      const prompt = 'Broken: !{echo {a}';
+
+      await expect(processor.process(prompt, context)).rejects.toThrow(
+        /Unclosed shell injection/,
+      );
+    });
+  });
+
+  describe('Error Reporting', () => {
+    it('should append stderr information if the command produces it', async () => {
+      const processor = new ShellProcessor('test-command');
+      const prompt = '!{cmd}';
+      mockShellExecute.mockReturnValue({
+        result: Promise.resolve({
+          ...SUCCESS_RESULT,
+          stdout: 'some output',
+          stderr: 'some error',
+        }),
+      });
+
+      const result = await processor.process(prompt, context);
+
+      expect(result).toBe('some output\n--- STDERR ---\nsome error');
+    });
+
+    it('should handle stderr-only output correctly', async () => {
+      const processor = new ShellProcessor('test-command');
+      const prompt = '!{cmd}';
+      mockShellExecute.mockReturnValue({
+        result: Promise.resolve({
+          ...SUCCESS_RESULT,
+          stdout: '',
+          stderr: 'error only',
+        }),
+      });
+
+      const result = await processor.process(prompt, context);
+
+      expect(result).toBe('--- STDERR ---\nerror only');
+    });
+
+    it('should append exit code and command name if the command fails', async () => {
+      const processor = new ShellProcessor('test-command');
+      const prompt = 'Run a failing command: !{exit 1}';
+      mockShellExecute.mockReturnValue({
+        result: Promise.resolve({
+          ...SUCCESS_RESULT,
+          stdout: 'some error output',
+          stderr: '',
+          exitCode: 1,
+        }),
+      });
+
+      const result = await processor.process(prompt, context);
+
+      expect(result).toBe(
+        "Run a failing command: some error output\n[Shell command 'exit 1' exited with code 1]",
+      );
+    });
+
+    it('should append signal info and command name if terminated by signal', async () => {
+      const processor = new ShellProcessor('test-command');
+      const prompt = '!{cmd}';
+      mockShellExecute.mockReturnValue({
+        result: Promise.resolve({
+          ...SUCCESS_RESULT,
+          stdout: 'output',
+          stderr: '',
+          exitCode: null,
+          signal: 'SIGTERM',
+        }),
+      });
+
+      const result = await processor.process(prompt, context);
+
+      expect(result).toBe(
+        "output\n[Shell command 'cmd' terminated by signal SIGTERM]",
+      );
+    });
+
+    it('should append stderr and exit code information correctly', async () => {
+      const processor = new ShellProcessor('test-command');
+      const prompt = '!{cmd}';
+      mockShellExecute.mockReturnValue({
+        result: Promise.resolve({
+          ...SUCCESS_RESULT,
+          stdout: 'out',
+          stderr: 'err',
+          exitCode: 127,
+        }),
+      });
+
+      const result = await processor.process(prompt, context);
+
+      expect(result).toBe(
+        "out\n--- STDERR ---\nerr\n[Shell command 'cmd' exited with code 127]",
+      );
+    });
+
+    it('should throw a detailed error if the shell fails to spawn', async () => {
+      const processor = new ShellProcessor('test-command');
+      const prompt = '!{bad-command}';
+      const spawnError = new Error('spawn EACCES');
+      mockShellExecute.mockReturnValue({
+        result: Promise.resolve({
+          ...SUCCESS_RESULT,
+          stdout: '',
+          stderr: '',
+          exitCode: null,
+          error: spawnError,
+          aborted: false,
+        }),
+      });
+
+      await expect(processor.process(prompt, context)).rejects.toThrow(
+        "Failed to start shell command in 'test-command': spawn EACCES. Command: bad-command",
+      );
+    });
+
+    it('should report abort status with command name if aborted', async () => {
+      const processor = new ShellProcessor('test-command');
+      const prompt = '!{long-running-command}';
+      const spawnError = new Error('Aborted');
+      mockShellExecute.mockReturnValue({
+        result: Promise.resolve({
+          ...SUCCESS_RESULT,
+          stdout: 'partial output',
+          stderr: '',
+          exitCode: null,
+          error: spawnError,
+          aborted: true, // Key difference
+        }),
+      });
+
+      const result = await processor.process(prompt, context);
+      expect(result).toBe(
+        "partial output\n[Shell command 'long-running-command' aborted]",
+      );
+    });
+  });
+
+  describe('Context-Aware Argument Interpolation ({{args}})', () => {
+    const rawArgs = 'user input';
+
+    beforeEach(() => {
+      // Update context for these tests to use specific arguments
+      context.invocation!.args = rawArgs;
+    });
+
+    it('should perform raw replacement if no shell injections are present (optimization path)', async () => {
+      const processor = new ShellProcessor('test-command');
+      const prompt = 'The user said: {{args}}';
+
+      const result = await processor.process(prompt, context);
+
+      expect(result).toBe(`The user said: ${rawArgs}`);
+      expect(mockShellExecute).not.toHaveBeenCalled();
+    });
+
+    it('should perform raw replacement outside !{} blocks', async () => {
+      const processor = new ShellProcessor('test-command');
+      const prompt = 'Outside: {{args}}. Inside: !{echo "hello"}';
+      mockShellExecute.mockReturnValue({
+        result: Promise.resolve({ ...SUCCESS_RESULT, stdout: 'hello' }),
+      });
+
+      const result = await processor.process(prompt, context);
+
+      expect(result).toBe(`Outside: ${rawArgs}. Inside: hello`);
+    });
+
+    it('should perform escaped replacement inside !{} blocks', async () => {
+      const processor = new ShellProcessor('test-command');
+      const prompt = 'Command: !{grep {{args}} file.txt}';
+      mockShellExecute.mockReturnValue({
+        result: Promise.resolve({ ...SUCCESS_RESULT, stdout: 'match found' }),
+      });
+
+      const result = await processor.process(prompt, context);
+
+      const expectedEscapedArgs = getExpectedEscapedArgForPlatform(rawArgs);
+      const expectedCommand = `grep ${expectedEscapedArgs} file.txt`;
+
+      expect(mockShellExecute).toHaveBeenCalledWith(
+        expectedCommand,
+        expect.any(String),
+        expect.any(Function),
+        expect.any(Object),
+      );
+
+      expect(result).toBe('Command: match found');
+    });
+
+    it('should handle both raw (outside) and escaped (inside) injection simultaneously', async () => {
+      const processor = new ShellProcessor('test-command');
+      const prompt = 'User "({{args}})" requested search: !{search {{args}}}';
+      mockShellExecute.mockReturnValue({
+        result: Promise.resolve({ ...SUCCESS_RESULT, stdout: 'results' }),
+      });
+
+      const result = await processor.process(prompt, context);
+
+      const expectedEscapedArgs = getExpectedEscapedArgForPlatform(rawArgs);
+      const expectedCommand = `search ${expectedEscapedArgs}`;
+      expect(mockShellExecute).toHaveBeenCalledWith(
+        expectedCommand,
+        expect.any(String),
+        expect.any(Function),
+        expect.any(Object),
+      );
+
+      expect(result).toBe(`User "(${rawArgs})" requested search: results`);
+    });
+
+    it('should perform security checks on the final, resolved (escaped) command', async () => {
+      const processor = new ShellProcessor('test-command');
+      const prompt = '!{rm {{args}}}';
+
+      const expectedEscapedArgs = getExpectedEscapedArgForPlatform(rawArgs);
+      const expectedResolvedCommand = `rm ${expectedEscapedArgs}`;
+      mockCheckCommandPermissions.mockReturnValue({
+        allAllowed: false,
+        disallowedCommands: [expectedResolvedCommand],
+        isHardDenial: false,
+      });
+
+      await expect(processor.process(prompt, context)).rejects.toThrow(
+        ConfirmationRequiredError,
+      );
+
+      expect(mockCheckCommandPermissions).toHaveBeenCalledWith(
+        expectedResolvedCommand,
+        expect.any(Object),
+        context.session.sessionShellAllowlist,
+      );
+    });
+
+    it('should report the resolved command if a hard denial occurs', async () => {
+      const processor = new ShellProcessor('test-command');
+      const prompt = '!{rm {{args}}}';
+      const expectedEscapedArgs = getExpectedEscapedArgForPlatform(rawArgs);
+      const expectedResolvedCommand = `rm ${expectedEscapedArgs}`;
+      mockCheckCommandPermissions.mockReturnValue({
+        allAllowed: false,
+        disallowedCommands: [expectedResolvedCommand],
+        isHardDenial: true,
+        blockReason: 'It is forbidden.',
+      });
+
+      await expect(processor.process(prompt, context)).rejects.toThrow(
+        `Blocked command: "${expectedResolvedCommand}". Reason: It is forbidden.`,
+      );
+    });
+  });
+  describe('Real-World Escaping Scenarios', () => {
+    it('should correctly handle multiline arguments', async () => {
+      const processor = new ShellProcessor('test-command');
+      const multilineArgs = 'first line\nsecond line';
+      context.invocation!.args = multilineArgs;
+      const prompt = 'Commit message: !{git commit -m {{args}}}';
+
+      const expectedEscapedArgs =
+        getExpectedEscapedArgForPlatform(multilineArgs);
+      const expectedCommand = `git commit -m ${expectedEscapedArgs}`;
+
+      await processor.process(prompt, context);
+
+      expect(mockShellExecute).toHaveBeenCalledWith(
+        expectedCommand,
+        expect.any(String),
+        expect.any(Function),
+        expect.any(Object),
+      );
+    });
+
+    it.each([
+      { name: 'spaces', input: 'file with spaces.txt' },
+      { name: 'double quotes', input: 'a "quoted" string' },
+      { name: 'single quotes', input: "it's a string" },
+      { name: 'command substitution (backticks)', input: '`reboot`' },
+      { name: 'command substitution (dollar)', input: '$(reboot)' },
+      { name: 'variable expansion', input: '$HOME' },
+      { name: 'command chaining (semicolon)', input: 'a; reboot' },
+      { name: 'command chaining (ampersand)', input: 'a && reboot' },
+    ])('should safely escape args containing $name', async ({ input }) => {
+      const processor = new ShellProcessor('test-command');
+      context.invocation!.args = input;
+      const prompt = '!{echo {{args}}}';
+
+      const expectedEscapedArgs = getExpectedEscapedArgForPlatform(input);
+      const expectedCommand = `echo ${expectedEscapedArgs}`;
+
+      await processor.process(prompt, context);
+
+      expect(mockShellExecute).toHaveBeenCalledWith(
+        expectedCommand,
+        expect.any(String),
+        expect.any(Function),
+        expect.any(Object),
+      );
+    });
   });
 });