From a588d5cd10e71bbb99912fb4d7f43c162a2a68ba Mon Sep 17 00:00:00 2001
From: Taylor Mullen <ntaylormullen@google.com>
Date: Tue, 6 May 2025 22:26:38 -0700
Subject: [PATCH] Prevent UI hang on long tool confirmations.

Problem:
When a tool confirmation dialog appeared for a potentially long-running
operation (e.g., `npm install`), accepting the confirmation would cause
the UI to appear to hang. The confirmation dialog would remain visible,
and no further UI updates would occur until the long-running task
completed. This provided a poor user experience as the application
seemed unresponsive.

Fix:
This change addresses the issue by ensuring the UI is updated to remove
the confirmation dialog *before* the long-running operation begins.
It also marks the tool as executing so a spinner can be shown.

Fixes https://b.corp.google.com/issues/415844994

Signed, sealed, delivered, it's yours!
   - Gemini, your friendly neighborhood code-slinger
---
 .../ui/components/messages/ToolMessage.tsx    |  3 ++-
 packages/cli/src/ui/hooks/useGeminiStream.ts  | 27 +++++++++++++++++++
 packages/cli/src/ui/types.ts                  |  1 +
 3 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/packages/cli/src/ui/components/messages/ToolMessage.tsx b/packages/cli/src/ui/components/messages/ToolMessage.tsx
index 4d5fca37..1d8c478c 100644
--- a/packages/cli/src/ui/components/messages/ToolMessage.tsx
+++ b/packages/cli/src/ui/components/messages/ToolMessage.tsx
@@ -25,7 +25,8 @@ export const ToolMessage: React.FC<IndividualToolCallDisplay> = ({
       <Box minHeight={1}>
         {/* Status Indicator */}
         <Box minWidth={statusIndicatorWidth}>
-          {status === ToolCallStatus.Pending && <Spinner type="dots" />}
+          {(status === ToolCallStatus.Pending ||
+            status === ToolCallStatus.Executing) && <Spinner type="dots" />}
           {status === ToolCallStatus.Success && (
             <Text color={Colors.AccentGreen}>✔</Text>
           )}
diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts
index 611ea3e7..40c688a3 100644
--- a/packages/cli/src/ui/hooks/useGeminiStream.ts
+++ b/packages/cli/src/ui/hooks/useGeminiStream.ts
@@ -396,6 +396,32 @@ export const useGeminiStream = (
           // Call the original server-side handler first
           originalConfirmationDetails.onConfirm(outcome);
 
+          // Ensure UI updates before potentially long-running operations
+          if (currentToolGroupMessageId !== null) {
+            updateItem(
+              currentToolGroupMessageId,
+              (currentItem: HistoryItem) => {
+                if (currentItem?.type !== 'tool_group')
+                  return currentItem as Partial<Omit<HistoryItem, 'id'>>;
+                return {
+                  ...currentItem,
+                  tools: (currentItem.tools || []).map((tool) =>
+                    tool.callId === request.callId
+                      ? {
+                          ...tool,
+                          confirmationDetails: undefined,
+                          status: ToolCallStatus.Executing,
+                        }
+                      : tool,
+                  ),
+                } as Partial<Omit<HistoryItem, 'id'>>;
+              },
+            );
+            refreshStatic();
+          }
+
+          await new Promise((resolve) => setTimeout(resolve, 0)); // Allow UI to re-render
+
           if (outcome === ToolConfirmationOutcome.Cancel) {
             let resultDisplay: ToolResultDisplay | undefined;
             if ('fileDiff' in originalConfirmationDetails) {
@@ -470,6 +496,7 @@ export const useGeminiStream = (
       setShowHelp,
       toolRegistry,
       setInitError,
+      refreshStatic,
     ],
   );
 
diff --git a/packages/cli/src/ui/types.ts b/packages/cli/src/ui/types.ts
index 4e1ab9cc..d6279598 100644
--- a/packages/cli/src/ui/types.ts
+++ b/packages/cli/src/ui/types.ts
@@ -27,6 +27,7 @@ export enum ToolCallStatus {
   Pending = 'Pending',
   Canceled = 'Canceled',
   Confirming = 'Confirming',
+  Executing = 'Executing',
   Success = 'Success',
   Error = 'Error',
 }