From 8440b971f5cca47cced8458892f4ec0932745d70 Mon Sep 17 00:00:00 2001 From: Taylor Mullen Date: Sun, 25 May 2025 16:01:10 -0700 Subject: [PATCH] Fix(cli): Prevent premature input box reactivation during tool confirmation - Introduced a 'validating' state for tool calls to prevent the input box from reappearing while waiting for a tool's `shouldConfirmExecute` method to complete. - When a tool call is initiated, it's now immediately set to a 'validating' status. This ensures the UI remains in a busy/responding state. - `useGeminiStream` now considers the 'validating' state as part of `StreamingState.Responding`. - `useToolScheduler` has been updated to: - Set the initial status of new tool calls to 'validating'. - Asynchronously perform the `shouldConfirmExecute` check. - Transition to 'awaiting_approval' or 'scheduled' based on the check's outcome. - This resolves an issue where a slow `shouldConfirmExecute` could lead to the input prompt becoming active again before the tool call lifecycle was fully determined. While 'validating' is currently treated similarly to 'executing' in the UI, this new state provides a foundation for more customized user experiences during this phase in the future. Fixes https://github.com/google-gemini/gemini-cli/issues/527 --- packages/cli/src/ui/hooks/useGeminiStream.ts | 5 +- packages/cli/src/ui/hooks/useToolScheduler.ts | 90 ++++++++++++++----- 2 files changed, 73 insertions(+), 22 deletions(-) diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index 9dcb005b..76d29189 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -109,7 +109,10 @@ export const useGeminiStream = ( if ( isResponding || toolCalls.some( - (t) => t.status === 'executing' || t.status === 'scheduled', + (t) => + t.status === 'executing' || + t.status === 'scheduled' || + t.status === 'validating', ) ) { return StreamingState.Responding; diff --git a/packages/cli/src/ui/hooks/useToolScheduler.ts b/packages/cli/src/ui/hooks/useToolScheduler.ts index a5770d36..1bb44133 100644 --- a/packages/cli/src/ui/hooks/useToolScheduler.ts +++ b/packages/cli/src/ui/hooks/useToolScheduler.ts @@ -20,6 +20,12 @@ import { ToolCallStatus, } from '../types.js'; +type ValidatingToolCall = { + status: 'validating'; + request: ToolCallRequestInfo; + tool: Tool; +}; + type ScheduledToolCall = { status: 'scheduled'; request: ToolCallRequestInfo; @@ -62,6 +68,7 @@ type WaitingToolCall = { export type Status = ToolCall['status']; export type ToolCall = + | ValidatingToolCall | ScheduledToolCall | ErroredToolCall | SuccessfulToolCall @@ -99,9 +106,12 @@ export function useToolScheduler( 'Cannot schedule tool calls while other tool calls are running', ); } - const requests = Array.isArray(request) ? request : [request]; - const newCalls: ToolCall[] = await Promise.all( - requests.map(async (r): Promise => { + const requestsToProcess = Array.isArray(request) ? request : [request]; + + // Step 1: Create initial calls with 'validating' status (or 'error' if tool not found) + // and add them to the state immediately to make the UI busy. + const initialNewCalls: ToolCall[] = requestsToProcess.map( + (r): ToolCall => { const tool = toolRegistry.getTool(r.name); if (!tool) { return { @@ -113,16 +123,27 @@ export function useToolScheduler( ), }; } + // Set to 'validating' immediately. This will make streamingState 'Responding'. + return { status: 'validating', request: r, tool }; + }, + ); + setToolCalls((prevCalls) => prevCalls.concat(initialNewCalls)); + // Step 2: Asynchronously check for confirmation and update status for each new call. + initialNewCalls.forEach(async (initialCall) => { + // If the call was already marked as an error (tool not found), skip further processing. + if (initialCall.status !== 'validating') return; + + const { request: r, tool } = initialCall; + try { const userApproval = await tool.shouldConfirmExecute(r.args); if (userApproval) { - return { - status: 'awaiting_approval', - request: r, - tool, - confirmationDetails: { + // Confirmation is needed. Update status to 'awaiting_approval'. + setToolCalls( + setStatus(r.callId, 'awaiting_approval', { ...userApproval, onConfirm: async (outcome) => { + // This onConfirm is triggered by user interaction later. await userApproval.onConfirm(outcome); setToolCalls( outcome === ToolConfirmationOutcome.Cancel @@ -131,21 +152,30 @@ export function useToolScheduler( 'cancelled', 'User did not allow tool call', ) - : setStatus(r.callId, 'scheduled'), + : // If confirmed, it goes to 'scheduled' to be picked up by the execution effect. + setStatus(r.callId, 'scheduled'), ); }, - }, - }; + }), + ); + } else { + // No confirmation needed, move to 'scheduled' for execution. + setToolCalls(setStatus(r.callId, 'scheduled')); } - - return { - status: 'scheduled', - request: r, - tool, - }; - }), - ); - setToolCalls((t) => t.concat(newCalls)); + } catch (e) { + // Handle errors from tool.shouldConfirmExecute() itself. + setToolCalls( + setStatus( + r.callId, + 'error', + toolErrorResponse( + r, + e instanceof Error ? e : new Error(String(e)), + ), + ), + ); + } + }); }, [isRunning, setToolCalls, toolRegistry], ); @@ -273,7 +303,7 @@ function setStatus( ): (t: ToolCall[]) => ToolCall[]; function setStatus( targetCallId: string, - status: 'executing' | 'scheduled', + status: 'executing' | 'scheduled' | 'validating', ): (t: ToolCall[]) => ToolCall[]; function setStatus( targetCallId: string, @@ -338,6 +368,13 @@ function setStatus( }; return next; } + case 'validating': { + const next: ValidatingToolCall = { + ...(t as ValidatingToolCall), // Added type assertion for safety + status: 'validating', + }; + return next; + } case 'executing': { const next: ExecutingToolCall = { ...t, @@ -373,6 +410,8 @@ const toolErrorResponse = ( function mapStatus(status: Status): ToolCallStatus { switch (status) { + case 'validating': + return ToolCallStatus.Executing; case 'awaiting_approval': return ToolCallStatus.Confirming; case 'executing': @@ -445,6 +484,15 @@ export function mapToDisplay( status: mapStatus(t.status), confirmationDetails: undefined, }; + case 'validating': // Add this case + return { + callId: t.request.callId, + name: t.tool.displayName, + description: t.tool.getDescription(t.request.args), + resultDisplay: undefined, + status: mapStatus(t.status), + confirmationDetails: undefined, + }; case 'scheduled': return { callId: t.request.callId,