allow aborting of shell mode (!) commands, similar to shell tool commands. fix bug that prevented aborts after first abort. more robust killing logic (#616)

3 files changed, 92 insertions, 31 deletions

diff --git a/packages/cli/src/ui/hooks/shellCommandProcessor.ts b/packages/cli/src/ui/hooks/shellCommandProcessor.ts
index 59e337b4..d6c837d1 100644
--- a/packages/cli/src/ui/hooks/shellCommandProcessor.ts
+++ b/packages/cli/src/ui/hooks/shellCommandProcessor.ts
@@ -37,7 +37,7 @@ export const useShellCommandProcessor = (
    * @returns True if the query was handled as a shell command, false otherwise.
    */
   const handleShellCommand = useCallback(
-    (rawQuery: PartListUnion): boolean => {
+    (rawQuery: PartListUnion, signal?: AbortSignal): boolean => {
       if (typeof rawQuery !== 'string') {
         return false;
       }
@@ -120,6 +120,7 @@ export const useShellCommandProcessor = (
           const child = spawn('bash', ['-c', commandToExecute], {
             cwd: targetDir,
             stdio: ['ignore', 'pipe', 'pipe'],
+            detached: true, // Important for process group killing
           });
 
           let exited = false;
@@ -148,18 +149,75 @@ export const useShellCommandProcessor = (
             error = err;
           });
 
-          child.on('exit', (code, signal) => {
+          const abortHandler = () => {
+            if (child.pid && !exited) {
+              onDebugMessage(
+                `Aborting shell command (PID: ${child.pid}) due to signal.`,
+              );
+              try {
+                // attempt to SIGTERM process group (negative PID)
+                // if SIGTERM fails after 200ms, attempt SIGKILL
+                process.kill(-child.pid, 'SIGTERM');
+                setTimeout(() => {
+                  // if SIGTERM fails, attempt SIGKILL
+                  try {
+                    if (child.pid && !exited) {
+                      process.kill(-child.pid, 'SIGKILL');
+                    }
+                  } catch (_e) {
+                    console.error(
+                      `failed to kill shell process ${child.pid}: ${_e}`,
+                    );
+                  }
+                }, 200);
+              } catch (_e) {
+                // if group kill fails, fall back to killing just the main process
+                try {
+                  if (child.pid) {
+                    child.kill('SIGKILL');
+                  }
+                } catch (_e) {
+                  console.error(
+                    `failed to kill shell process ${child.pid}: ${_e}`,
+                  );
+                }
+              }
+            }
+          };
+
+          if (signal) {
+            if (signal.aborted) {
+              abortHandler();
+              // No need to add listener if already aborted
+            } else {
+              signal.addEventListener('abort', abortHandler, { once: true });
+            }
+          }
+
+          child.on('exit', (code, processSignal) => {
             exited = true;
+            if (signal) {
+              signal.removeEventListener('abort', abortHandler);
+            }
             setPendingHistoryItem(null);
             output = output.trim() || '(Command produced no output)';
             if (error) {
               const text = `${error.message.replace(commandToExecute, rawQuery)}\n${output}`;
               addItemToHistory({ type: 'error', text }, userMessageTimestamp);
-            } else if (code !== 0) {
+            } else if (code !== null && code !== 0) {
               const text = `Command exited with code ${code}\n${output}`;
               addItemToHistory({ type: 'error', text }, userMessageTimestamp);
-            } else if (signal) {
-              const text = `Command terminated with signal ${signal}\n${output}`;
+            } else if (signal?.aborted) {
+              addItemToHistory(
+                {
+                  type: 'info',
+                  text: `Command was cancelled.\n${output}`,
+                },
+                userMessageTimestamp,
+              );
+            } else if (processSignal) {
+              const text = `Command terminated with signal ${processSignal}
+${output}`;
               addItemToHistory({ type: 'error', text }, userMessageTimestamp);
             } else {
               addItemToHistory(
diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts
index afaf0ccd..66fb317a 100644
--- a/packages/cli/src/ui/hooks/useGeminiStream.ts
+++ b/packages/cli/src/ui/hooks/useGeminiStream.ts
@@ -199,7 +199,7 @@ export const useGeminiStream = (
           return { queryToSend: null, shouldProceed: false }; // Handled by scheduling the tool
         }
 
-        if (shellModeActive && handleShellCommand(trimmedQuery)) {
+        if (shellModeActive && handleShellCommand(trimmedQuery, signal)) {
           return { queryToSend: null, shouldProceed: false };
         }
 
@@ -492,7 +492,7 @@ export const useGeminiStream = (
       const userMessageTimestamp = Date.now();
       setShowHelp(false);
 
-      abortControllerRef.current ??= new AbortController();
+      abortControllerRef.current = new AbortController();
       const signal = abortControllerRef.current.signal;
 
       const { queryToSend, shouldProceed } = await prepareQueryForGemini(
diff --git a/packages/server/src/tools/shell.ts b/packages/server/src/tools/shell.ts
index 38848c4f..a4634a3e 100644
--- a/packages/server/src/tools/shell.ts
+++ b/packages/server/src/tools/shell.ts
@@ -25,8 +25,6 @@ export interface ShellToolParams {
 }
 import { spawn } from 'child_process';
 
-const OUTPUT_UPDATE_INTERVAL_MS = 1000;
-
 export class ShellTool extends BaseTool<ShellToolParams, ToolResult> {
   static Name: string = 'execute_bash_command';
   private whitelist: Set<string> = new Set();
@@ -126,7 +124,7 @@ export class ShellTool extends BaseTool<ShellToolParams, ToolResult> {
   async execute(
     params: ShellToolParams,
     abortSignal: AbortSignal,
-    updateOutput?: (output: string) => void,
+    onOutputChunk?: (chunk: string) => void,
   ): Promise<ToolResult> {
     const validationError = this.validateToolParams(params);
     if (validationError) {
@@ -157,19 +155,6 @@ export class ShellTool extends BaseTool<ShellToolParams, ToolResult> {
     let exited = false;
     let stdout = '';
     let output = '';
-    let lastUpdateTime = Date.now();
-
-    const appendOutput = (str: string) => {
-      output += str;
-      if (
-        updateOutput &&
-        Date.now() - lastUpdateTime > OUTPUT_UPDATE_INTERVAL_MS
-      ) {
-        updateOutput(output);
-        lastUpdateTime = Date.now();
-      }
-    };
-
     shell.stdout.on('data', (data: Buffer) => {
       // continue to consume post-exit for background processes
       // removing listeners can overflow OS buffer and block subprocesses
@@ -177,7 +162,10 @@ export class ShellTool extends BaseTool<ShellToolParams, ToolResult> {
       if (!exited) {
         const str = data.toString();
         stdout += str;
-        appendOutput(str);
+        output += str;
+        if (onOutputChunk) {
+          onOutputChunk(str);
+        }
       }
     });
 
@@ -186,7 +174,10 @@ export class ShellTool extends BaseTool<ShellToolParams, ToolResult> {
       if (!exited) {
         const str = data.toString();
         stderr += str;
-        appendOutput(str);
+        output += str;
+        if (onOutputChunk) {
+          onOutputChunk(str);
+        }
       }
     });
 
@@ -210,16 +201,28 @@ export class ShellTool extends BaseTool<ShellToolParams, ToolResult> {
     shell.on('exit', exitHandler);
 
     const abortHandler = () => {
-      if (shell.pid) {
+      if (shell.pid && !exited) {
         try {
-          // Kill the entire process group
+          // attempt to SIGTERM process group (negative PID)
+          // if SIGTERM fails after 200ms, attempt SIGKILL
           process.kill(-shell.pid, 'SIGTERM');
+          setTimeout(() => {
+            try {
+              if (shell.pid && !exited) {
+                process.kill(-shell.pid, 'SIGKILL');
+              }
+            } catch (_e) {
+              console.error(`failed to kill shell process ${shell.pid}: ${_e}`);
+            }
+          }, 200);
         } catch (_e) {
-          // Fallback to killing the main process if group kill fails
+          // if group kill fails, fall back to killing just the main process
           try {
-            shell.kill('SIGKILL'); // or 'SIGTERM'
-          } catch (_killError) {
-            // Ignore errors if the process is already dead
+            if (shell.pid) {
+              shell.kill('SIGKILL');
+            }
+          } catch (_e) {
+            console.error(`failed to kill shell process ${shell.pid}: ${_e}`);
           }
         }
       }