fix: cancel parallel tool calls mid-execution (#489)

2 files changed, 12 insertions, 2 deletions

diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts
index d3ecad95..80f6945c 100644
--- a/packages/cli/src/ui/hooks/useGeminiStream.ts
+++ b/packages/cli/src/ui/hooks/useGeminiStream.ts
@@ -395,7 +395,10 @@ export const useGeminiStream = (
   };
 
   const streamingState: StreamingState =
-    isResponding || toolCalls.some((t) => t.status === 'awaiting_approval')
+    isResponding ||
+    toolCalls.some(
+      (t) => t.status === 'awaiting_approval' || t.status === 'executing',
+    )
       ? StreamingState.Responding
       : StreamingState.Idle;
 
diff --git a/packages/cli/src/ui/hooks/useToolScheduler.ts b/packages/cli/src/ui/hooks/useToolScheduler.ts
index e14241b6..8bcc0ae9 100644
--- a/packages/cli/src/ui/hooks/useToolScheduler.ts
+++ b/packages/cli/src/ui/hooks/useToolScheduler.ts
@@ -185,12 +185,19 @@ export function useToolScheduler(
   useEffect(() => {
     // effect for executing scheduled tool calls
     if (toolCalls.every((t) => t.status === 'scheduled')) {
+      const signal = abortController.signal;
       toolCalls.forEach((c) => {
         const callId = c.request.callId;
         setToolCalls(setStatus(c.request.callId, 'executing'));
         c.tool
-          .execute(c.request.args, abortController.signal)
+          .execute(c.request.args, signal)
           .then((result) => {
+            if (signal.aborted) {
+              setToolCalls(
+                setStatus(callId, 'cancelled', 'Cancelled during execution'),
+              );
+              return;
+            }
             const functionResponse: Part = {
               functionResponse: {
                 name: c.request.name,