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

packages/cli/src/services/FileCommandLoader.test.ts

@@ -22,7 +22,8 @@ import {
   ConfirmationRequiredError,
   ShellProcessor,
 } from './prompt-processors/shellProcessor.js';
-import { ShorthandArgumentProcessor } from './prompt-processors/argumentProcessor.js';
+import { DefaultArgumentProcessor } from './prompt-processors/argumentProcessor.js';
+import { CommandContext } from '../ui/commands/types.js';
 
 const mockShellProcess = vi.hoisted(() => vi.fn());
 vi.mock('./prompt-processors/shellProcessor.js', () => ({
@@ -46,9 +47,6 @@ vi.mock('./prompt-processors/argumentProcessor.js', async (importOriginal) => {
       typeof import('./prompt-processors/argumentProcessor.js')
     >();
   return {
-    ShorthandArgumentProcessor: vi
-      .fn()
-      .mockImplementation(() => new original.ShorthandArgumentProcessor()),
     DefaultArgumentProcessor: vi
       .fn()
       .mockImplementation(() => new original.DefaultArgumentProcessor()),
@@ -71,7 +69,16 @@ describe('FileCommandLoader', () => {
 
   beforeEach(() => {
     vi.clearAllMocks();
-    mockShellProcess.mockImplementation((prompt) => Promise.resolve(prompt));
+    mockShellProcess.mockImplementation(
+      (prompt: string, context: CommandContext) => {
+        const userArgsRaw = context?.invocation?.args || '';
+        const processedPrompt = prompt.replaceAll(
+          SHORTHAND_ARGS_PLACEHOLDER,
+          userArgsRaw,
+        );
+        return Promise.resolve(processedPrompt);
+      },
+    );
   });
 
   afterEach(() => {
@@ -198,7 +205,7 @@ describe('FileCommandLoader', () => {
     const mockConfig = {
       getProjectRoot: vi.fn(() => '/path/to/project'),
       getExtensions: vi.fn(() => []),
-    } as Config;
+    } as unknown as Config;
     const loader = new FileCommandLoader(mockConfig);
     const commands = await loader.loadCommands(signal);
     expect(commands).toHaveLength(1);
@@ -239,7 +246,7 @@ describe('FileCommandLoader', () => {
     const mockConfig = {
       getProjectRoot: vi.fn(() => process.cwd()),
       getExtensions: vi.fn(() => []),
-    } as Config;
+    } as unknown as Config;
     const loader = new FileCommandLoader(mockConfig);
     const commands = await loader.loadCommands(signal);
 
@@ -379,6 +386,68 @@ describe('FileCommandLoader', () => {
     expect(command.name).toBe('legacy_command');
   });
 
+  describe('Processor Instantiation Logic', () => {
+    it('instantiates only DefaultArgumentProcessor if no {{args}} or !{} are present', async () => {
+      const userCommandsDir = getUserCommandsDir();
+      mock({
+        [userCommandsDir]: {
+          'simple.toml': `prompt = "Just a regular prompt"`,
+        },
+      });
+
+      const loader = new FileCommandLoader(null as unknown as Config);
+      await loader.loadCommands(signal);
+
+      expect(ShellProcessor).not.toHaveBeenCalled();
+      expect(DefaultArgumentProcessor).toHaveBeenCalledTimes(1);
+    });
+
+    it('instantiates only ShellProcessor if {{args}} is present (but not !{})', async () => {
+      const userCommandsDir = getUserCommandsDir();
+      mock({
+        [userCommandsDir]: {
+          'args.toml': `prompt = "Prompt with {{args}}"`,
+        },
+      });
+
+      const loader = new FileCommandLoader(null as unknown as Config);
+      await loader.loadCommands(signal);
+
+      expect(ShellProcessor).toHaveBeenCalledTimes(1);
+      expect(DefaultArgumentProcessor).not.toHaveBeenCalled();
+    });
+
+    it('instantiates ShellProcessor and DefaultArgumentProcessor if !{} is present (but not {{args}})', async () => {
+      const userCommandsDir = getUserCommandsDir();
+      mock({
+        [userCommandsDir]: {
+          'shell.toml': `prompt = "Prompt with !{cmd}"`,
+        },
+      });
+
+      const loader = new FileCommandLoader(null as unknown as Config);
+      await loader.loadCommands(signal);
+
+      expect(ShellProcessor).toHaveBeenCalledTimes(1);
+      expect(DefaultArgumentProcessor).toHaveBeenCalledTimes(1);
+    });
+
+    it('instantiates only ShellProcessor if both {{args}} and !{} are present', async () => {
+      const userCommandsDir = getUserCommandsDir();
+      mock({
+        [userCommandsDir]: {
+          'both.toml': `prompt = "Prompt with {{args}} and !{cmd}"`,
+        },
+      });
+
+      const loader = new FileCommandLoader(null as unknown as Config);
+      await loader.loadCommands(signal);
+
+      expect(ShellProcessor).toHaveBeenCalledTimes(1);
+      expect(DefaultArgumentProcessor).not.toHaveBeenCalled();
+    });
+  });
+
   describe('Extension Command Loading', () => {
     it('loads commands from active extensions', async () => {
       const userCommandsDir = getUserCommandsDir();
@@ -416,7 +485,7 @@ describe('FileCommandLoader', () => {
             path: extensionDir,
           },
         ]),
-      } as Config;
+      } as unknown as Config;
       const loader = new FileCommandLoader(mockConfig);
       const commands = await loader.loadCommands(signal);
 
@@ -465,7 +534,7 @@ describe('FileCommandLoader', () => {
             path: extensionDir,
           },
         ]),
-      } as Config;
+      } as unknown as Config;
       const loader = new FileCommandLoader(mockConfig);
       const commands = await loader.loadCommands(signal);
 
@@ -572,7 +641,7 @@ describe('FileCommandLoader', () => {
             path: extensionDir2,
           },
         ]),
-      } as Config;
+      } as unknown as Config;
       const loader = new FileCommandLoader(mockConfig);
       const commands = await loader.loadCommands(signal);
 
@@ -608,7 +677,7 @@ describe('FileCommandLoader', () => {
             path: extensionDir,
           },
         ]),
-      } as Config;
+      } as unknown as Config;
       const loader = new FileCommandLoader(mockConfig);
       const commands = await loader.loadCommands(signal);
       expect(commands).toHaveLength(0);
@@ -640,7 +709,7 @@ describe('FileCommandLoader', () => {
         getExtensions: vi.fn(() => [
           { name: 'a', version: '1.0.0', isActive: true, path: extensionDir },
         ]),
-      } as Config;
+      } as unknown as Config;
       const loader = new FileCommandLoader(mockConfig);
       const commands = await loader.loadCommands(signal);
 
@@ -671,7 +740,7 @@ describe('FileCommandLoader', () => {
     });
   });
 
-  describe('Shorthand Argument Processor Integration', () => {
+  describe('Argument Handling Integration (via ShellProcessor)', () => {
     it('correctly processes a command with {{args}}', async () => {
       const userCommandsDir = getUserCommandsDir();
       mock({
@@ -738,6 +807,19 @@ describe('FileCommandLoader', () => {
   });
 
   describe('Shell Processor Integration', () => {
+    it('instantiates ShellProcessor if {{args}} is present (even without shell trigger)', async () => {
+      const userCommandsDir = getUserCommandsDir();
+      mock({
+        [userCommandsDir]: {
+          'args_only.toml': `prompt = "Hello {{args}}"`,
+        },
+      });
+
+      const loader = new FileCommandLoader(null as unknown as Config);
+      await loader.loadCommands(signal);
+
+      expect(ShellProcessor).toHaveBeenCalledWith('args_only');
+    });
     it('instantiates ShellProcessor if the trigger is present', async () => {
       const userCommandsDir = getUserCommandsDir();
       mock({
@@ -752,7 +834,7 @@ describe('FileCommandLoader', () => {
       expect(ShellProcessor).toHaveBeenCalledWith('shell');
     });
 
-    it('does not instantiate ShellProcessor if trigger is missing', async () => {
+    it('does not instantiate ShellProcessor if no triggers ({{args}} or !{}) are present', async () => {
       const userCommandsDir = getUserCommandsDir();
       mock({
         [userCommandsDir]: {
@@ -852,32 +934,30 @@ describe('FileCommandLoader', () => {
         ),
       ).rejects.toThrow('Something else went wrong');
     });
-
-    it('assembles the processor pipeline in the correct order (Shell -> Argument)', async () => {
+    it('assembles the processor pipeline in the correct order (Shell -> Default)', async () => {
       const userCommandsDir = getUserCommandsDir();
       mock({
         [userCommandsDir]: {
+          // This prompt uses !{} but NOT {{args}}, so both processors should be active.
           'pipeline.toml': `
-            prompt = "Shell says: ${SHELL_INJECTION_TRIGGER}echo foo} and user says: ${SHORTHAND_ARGS_PLACEHOLDER}"
-          `,
+              prompt = "Shell says: ${SHELL_INJECTION_TRIGGER}echo foo}."
+            `,
         },
       });
 
-      // Mock the process methods to track call order
-      const argProcessMock = vi
+      const defaultProcessMock = vi
         .fn()
-        .mockImplementation((p) => `${p}-arg-processed`);
+        .mockImplementation((p) => Promise.resolve(`${p}-default-processed`));
 
-      // Redefine the mock for this specific test
       mockShellProcess.mockImplementation((p) =>
         Promise.resolve(`${p}-shell-processed`),
       );
 
-      vi.mocked(ShorthandArgumentProcessor).mockImplementation(
+      vi.mocked(DefaultArgumentProcessor).mockImplementation(
         () =>
           ({
-            process: argProcessMock,
-          }) as unknown as ShorthandArgumentProcessor,
+            process: defaultProcessMock,
+          }) as unknown as DefaultArgumentProcessor,
       );
 
       const loader = new FileCommandLoader(null as unknown as Config);
@@ -885,7 +965,7 @@ describe('FileCommandLoader', () => {
       const command = commands.find((c) => c.name === 'pipeline');
       expect(command).toBeDefined();
 
-      await command!.action!(
+      const result = await command!.action!(
         createMockCommandContext({
           invocation: {
             raw: '/pipeline bar',
@@ -896,20 +976,27 @@ describe('FileCommandLoader', () => {
         'bar',
       );
 
-      // Verify that the shell processor was called before the argument processor
       expect(mockShellProcess.mock.invocationCallOrder[0]).toBeLessThan(
-        argProcessMock.mock.invocationCallOrder[0],
+        defaultProcessMock.mock.invocationCallOrder[0],
       );
 
-      // Also verify the flow of the prompt through the processors
+      // Verify the flow of the prompt through the processors
+      // 1. Shell processor runs first
       expect(mockShellProcess).toHaveBeenCalledWith(
-        expect.any(String),
+        expect.stringContaining(SHELL_INJECTION_TRIGGER),
         expect.any(Object),
       );
-      expect(argProcessMock).toHaveBeenCalledWith(
-        expect.stringContaining('-shell-processed'), // It receives the output of the shell processor
+      // 2. Default processor runs second
+      expect(defaultProcessMock).toHaveBeenCalledWith(
+        expect.stringContaining('-shell-processed'),
         expect.any(Object),
       );
+
+      if (result?.type === 'submit_prompt') {
+        expect(result.content).toContain('-shell-processed-default-processed');
+      } else {
+        assert.fail('Incorrect action type');
+      }
     });
   });
 });