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

packages/cli/src/ui/hooks/slashCommandProcessor.ts

@@ -9,7 +9,12 @@ import { type PartListUnion } from '@google/genai';
 import process from 'node:process';
 import { UseHistoryManagerReturn } from './useHistoryManager.js';
 import { useStateAndRef } from './useStateAndRef.js';
-import { Config, GitService, Logger } from '@google/gemini-cli-core';
+import {
+  Config,
+  GitService,
+  Logger,
+  ToolConfirmationOutcome,
+} from '@google/gemini-cli-core';
 import { useSessionStats } from '../contexts/SessionContext.js';
 import {
   Message,
@@ -44,9 +49,21 @@ export const useSlashCommandProcessor = (
   setQuittingMessages: (message: HistoryItem[]) => void,
   openPrivacyNotice: () => void,
   toggleVimEnabled: () => Promise<boolean>,
+  setIsProcessing: (isProcessing: boolean) => void,
 ) => {
   const session = useSessionStats();
   const [commands, setCommands] = useState<readonly SlashCommand[]>([]);
+  const [shellConfirmationRequest, setShellConfirmationRequest] =
+    useState<null | {
+      commands: string[];
+      onConfirm: (
+        outcome: ToolConfirmationOutcome,
+        approvedCommands?: string[],
+      ) => void;
+    }>(null);
+  const [sessionShellAllowlist, setSessionShellAllowlist] = useState(
+    new Set<string>(),
+  );
   const gitService = useMemo(() => {
     if (!config?.getProjectRoot()) {
       return;
@@ -144,6 +161,7 @@ export const useSlashCommandProcessor = (
       },
       session: {
         stats: session.stats,
+        sessionShellAllowlist,
       },
     }),
     [
@@ -161,6 +179,7 @@ export const useSlashCommandProcessor = (
       setPendingCompressionItem,
       toggleCorgiMode,
       toggleVimEnabled,
+      sessionShellAllowlist,
     ],
   );
 
@@ -189,69 +208,87 @@ export const useSlashCommandProcessor = (
   const handleSlashCommand = useCallback(
     async (
       rawQuery: PartListUnion,
+      oneTimeShellAllowlist?: Set<string>,
     ): Promise<SlashCommandProcessorResult | false> => {
-      if (typeof rawQuery !== 'string') {
-        return false;
-      }
-
-      const trimmed = rawQuery.trim();
-      if (!trimmed.startsWith('/') && !trimmed.startsWith('?')) {
-        return false;
-      }
-
-      const userMessageTimestamp = Date.now();
-      addItem({ type: MessageType.USER, text: trimmed }, userMessageTimestamp);
-
-      const parts = trimmed.substring(1).trim().split(/\s+/);
-      const commandPath = parts.filter((p) => p); // The parts of the command, e.g., ['memory', 'add']
-
-      let currentCommands = commands;
-      let commandToExecute: SlashCommand | undefined;
-      let pathIndex = 0;
-
-      for (const part of commandPath) {
-        // TODO: For better performance and architectural clarity, this two-pass
-        // search could be replaced. A more optimal approach would be to
-        // pre-compute a single lookup map in `CommandService.ts` that resolves
-        // all name and alias conflicts during the initial loading phase. The
-        // processor would then perform a single, fast lookup on that map.
-
-        // First pass: check for an exact match on the primary command name.
-        let foundCommand = currentCommands.find((cmd) => cmd.name === part);
-
-        // Second pass: if no primary name matches, check for an alias.
-        if (!foundCommand) {
-          foundCommand = currentCommands.find((cmd) =>
-            cmd.altNames?.includes(part),
-          );
+      setIsProcessing(true);
+      try {
+        if (typeof rawQuery !== 'string') {
+          return false;
         }
 
-        if (foundCommand) {
-          commandToExecute = foundCommand;
-          pathIndex++;
-          if (foundCommand.subCommands) {
-            currentCommands = foundCommand.subCommands;
+        const trimmed = rawQuery.trim();
+        if (!trimmed.startsWith('/') && !trimmed.startsWith('?')) {
+          return false;
+        }
+
+        const userMessageTimestamp = Date.now();
+        addItem(
+          { type: MessageType.USER, text: trimmed },
+          userMessageTimestamp,
+        );
+
+        const parts = trimmed.substring(1).trim().split(/\s+/);
+        const commandPath = parts.filter((p) => p); // The parts of the command, e.g., ['memory', 'add']
+
+        let currentCommands = commands;
+        let commandToExecute: SlashCommand | undefined;
+        let pathIndex = 0;
+
+        for (const part of commandPath) {
+          // TODO: For better performance and architectural clarity, this two-pass
+          // search could be replaced. A more optimal approach would be to
+          // pre-compute a single lookup map in `CommandService.ts` that resolves
+          // all name and alias conflicts during the initial loading phase. The
+          // processor would then perform a single, fast lookup on that map.
+
+          // First pass: check for an exact match on the primary command name.
+          let foundCommand = currentCommands.find((cmd) => cmd.name === part);
+
+          // Second pass: if no primary name matches, check for an alias.
+          if (!foundCommand) {
+            foundCommand = currentCommands.find((cmd) =>
+              cmd.altNames?.includes(part),
+            );
+          }
+
+          if (foundCommand) {
+            commandToExecute = foundCommand;
+            pathIndex++;
+            if (foundCommand.subCommands) {
+              currentCommands = foundCommand.subCommands;
+            } else {
+              break;
+            }
           } else {
             break;
           }
-        } else {
-          break;
         }
-      }
 
-      if (commandToExecute) {
-        const args = parts.slice(pathIndex).join(' ');
+        if (commandToExecute) {
+          const args = parts.slice(pathIndex).join(' ');
+
+          if (commandToExecute.action) {
+            const fullCommandContext: CommandContext = {
+              ...commandContext,
+              invocation: {
+                raw: trimmed,
+                name: commandToExecute.name,
+                args,
+              },
+            };
+
+            // If a one-time list is provided for a "Proceed" action, temporarily
+            // augment the session allowlist for this single execution.
+            if (oneTimeShellAllowlist && oneTimeShellAllowlist.size > 0) {
+              fullCommandContext.session = {
+                ...fullCommandContext.session,
+                sessionShellAllowlist: new Set([
+                  ...fullCommandContext.session.sessionShellAllowlist,
+                  ...oneTimeShellAllowlist,
+                ]),
+              };
+            }
 
-        if (commandToExecute.action) {
-          const fullCommandContext: CommandContext = {
-            ...commandContext,
-            invocation: {
-              raw: trimmed,
-              name: commandToExecute.name,
-              args,
-            },
-          };
-          try {
             const result = await commandToExecute.action(
               fullCommandContext,
               args,
@@ -323,6 +360,46 @@ export const useSlashCommandProcessor = (
                     type: 'submit_prompt',
                     content: result.content,
                   };
+                case 'confirm_shell_commands': {
+                  const { outcome, approvedCommands } = await new Promise<{
+                    outcome: ToolConfirmationOutcome;
+                    approvedCommands?: string[];
+                  }>((resolve) => {
+                    setShellConfirmationRequest({
+                      commands: result.commandsToConfirm,
+                      onConfirm: (
+                        resolvedOutcome,
+                        resolvedApprovedCommands,
+                      ) => {
+                        setShellConfirmationRequest(null); // Close the dialog
+                        resolve({
+                          outcome: resolvedOutcome,
+                          approvedCommands: resolvedApprovedCommands,
+                        });
+                      },
+                    });
+                  });
+
+                  if (
+                    outcome === ToolConfirmationOutcome.Cancel ||
+                    !approvedCommands ||
+                    approvedCommands.length === 0
+                  ) {
+                    return { type: 'handled' };
+                  }
+
+                  if (outcome === ToolConfirmationOutcome.ProceedAlways) {
+                    setSessionShellAllowlist(
+                      (prev) => new Set([...prev, ...approvedCommands]),
+                    );
+                  }
+
+                  return await handleSlashCommand(
+                    result.originalInvocation.raw,
+                    // Pass the approved commands as a one-time grant for this execution.
+                    new Set(approvedCommands),
+                  );
+                }
                 default: {
                   const unhandled: never = result;
                   throw new Error(
@@ -331,37 +408,39 @@ export const useSlashCommandProcessor = (
                 }
               }
             }
-          } catch (e) {
-            addItem(
-              {
-                type: MessageType.ERROR,
-                text: e instanceof Error ? e.message : String(e),
-              },
-              Date.now(),
-            );
+
+            return { type: 'handled' };
+          } else if (commandToExecute.subCommands) {
+            const helpText = `Command '/${commandToExecute.name}' requires a subcommand. Available:\n${commandToExecute.subCommands
+              .map((sc) => `  - ${sc.name}: ${sc.description || ''}`)
+              .join('\n')}`;
+            addMessage({
+              type: MessageType.INFO,
+              content: helpText,
+              timestamp: new Date(),
+            });
             return { type: 'handled' };
           }
-
-          return { type: 'handled' };
-        } else if (commandToExecute.subCommands) {
-          const helpText = `Command '/${commandToExecute.name}' requires a subcommand. Available:\n${commandToExecute.subCommands
-            .map((sc) => `  - ${sc.name}: ${sc.description || ''}`)
-            .join('\n')}`;
-          addMessage({
-            type: MessageType.INFO,
-            content: helpText,
-            timestamp: new Date(),
-          });
-          return { type: 'handled' };
         }
-      }
 
-      addMessage({
-        type: MessageType.ERROR,
-        content: `Unknown command: ${trimmed}`,
-        timestamp: new Date(),
-      });
-      return { type: 'handled' };
+        addMessage({
+          type: MessageType.ERROR,
+          content: `Unknown command: ${trimmed}`,
+          timestamp: new Date(),
+        });
+        return { type: 'handled' };
+      } catch (e) {
+        addItem(
+          {
+            type: MessageType.ERROR,
+            text: e instanceof Error ? e.message : String(e),
+          },
+          Date.now(),
+        );
+        return { type: 'handled' };
+      } finally {
+        setIsProcessing(false);
+      }
     },
     [
       config,
@@ -375,6 +454,9 @@ export const useSlashCommandProcessor = (
       openPrivacyNotice,
       openEditorDialog,
       setQuittingMessages,
+      setShellConfirmationRequest,
+      setSessionShellAllowlist,
+      setIsProcessing,
     ],
   );
 
@@ -383,5 +465,6 @@ export const useSlashCommandProcessor = (
     slashCommands: commands,
     pendingHistoryItems,
     commandContext,
+    shellConfirmationRequest,
   };
 };