From 997f461cad7ad175dce892be5f74bc7d8a8554b6 Mon Sep 17 00:00:00 2001 From: Coco Sheng Date: Fri, 1 May 2026 14:58:55 -0400 Subject: [PATCH] fix(cli): prevent Escape from clearing input buffer (#17083) (#26339) --- packages/cli/src/ui/AppContainer.tsx | 41 ++-- .../cli/src/ui/hooks/useAgentStream.test.tsx | 2 +- packages/cli/src/ui/hooks/useAgentStream.ts | 38 +++- .../cli/src/ui/hooks/useGeminiStream.test.tsx | 2 +- packages/cli/src/ui/hooks/useGeminiStream.ts | 204 ++++++++++-------- 5 files changed, 172 insertions(+), 115 deletions(-) diff --git a/packages/cli/src/ui/AppContainer.tsx b/packages/cli/src/ui/AppContainer.tsx index a09f477045..d8b1e1d277 100644 --- a/packages/cli/src/ui/AppContainer.tsx +++ b/packages/cli/src/ui/AppContainer.tsx @@ -1127,18 +1127,21 @@ Logging in with Google... Restarting Gemini CLI to continue. } }, [config, historyManager]); - const cancelHandlerRef = useRef<(shouldRestorePrompt?: boolean) => void>( - () => {}, - ); + const cancelHandlerRef = useRef< + (shouldRestorePrompt?: boolean, clearBuffer?: boolean) => void + >(() => {}); - const onCancelSubmit = useCallback((shouldRestorePrompt?: boolean) => { - if (shouldRestorePrompt) { - setPendingRestorePrompt(true); - } else { - setPendingRestorePrompt(false); - cancelHandlerRef.current(false); - } - }, []); + const onCancelSubmit = useCallback( + (shouldRestorePrompt?: boolean, clearBuffer: boolean = false) => { + if (shouldRestorePrompt) { + setPendingRestorePrompt(true); + } else { + setPendingRestorePrompt(false); + cancelHandlerRef.current(false, clearBuffer); + } + }, + [], + ); useEffect(() => { if (pendingRestorePrompt) { @@ -1321,18 +1324,18 @@ Logging in with Google... Restarting Gemini CLI to continue. }); cancelHandlerRef.current = useCallback( - (shouldRestorePrompt: boolean = true) => { - if (isToolAwaitingConfirmation(pendingHistoryItems)) { + (shouldRestorePrompt: boolean = true, clearBuffer: boolean = false) => { + if (!clearBuffer && isToolAwaitingConfirmation(pendingHistoryItems)) { return; // Don't clear - user may be composing a follow-up message } - if (isToolExecuting(pendingHistoryItems)) { - buffer.setText(''); // Clear for Ctrl+C cancellation - return; - } - // If cancelling (shouldRestorePrompt=false), never modify the buffer - // User is in control - preserve whatever text they typed, pasted, or restored + // If cancelling (shouldRestorePrompt=false): if (!shouldRestorePrompt) { + // Clear the buffer if explicitly requested (e.g., Ctrl+C) + if (clearBuffer) { + buffer.setText(''); + } + // Otherwise (e.g., Escape), user is in control - preserve whatever text they typed return; } diff --git a/packages/cli/src/ui/hooks/useAgentStream.test.tsx b/packages/cli/src/ui/hooks/useAgentStream.test.tsx index 53bb512504..1136a3592e 100644 --- a/packages/cli/src/ui/hooks/useAgentStream.test.tsx +++ b/packages/cli/src/ui/hooks/useAgentStream.test.tsx @@ -202,6 +202,6 @@ describe('useAgentStream', () => { }); expect(mockLegacyAgentProtocol.abort).toHaveBeenCalled(); - expect(mockOnCancelSubmit).toHaveBeenCalledWith(false); + expect(mockOnCancelSubmit).toHaveBeenCalledWith(false, true); }); }); diff --git a/packages/cli/src/ui/hooks/useAgentStream.ts b/packages/cli/src/ui/hooks/useAgentStream.ts index 926ba7cc7c..aea7b76ba5 100644 --- a/packages/cli/src/ui/hooks/useAgentStream.ts +++ b/packages/cli/src/ui/hooks/useAgentStream.ts @@ -36,11 +36,15 @@ import type { UseHistoryManagerReturn } from './useHistoryManager.js'; import { useSessionStats } from '../contexts/SessionContext.js'; import { useStateAndRef } from './useStateAndRef.js'; import { type MinimalTrackedToolCall } from './useTurnActivityMonitor.js'; +import { useKeypress } from './useKeypress.js'; export interface UseAgentStreamOptions { agent?: AgentProtocol; addItem: UseHistoryManagerReturn['addItem']; - onCancelSubmit: (shouldRestorePrompt?: boolean) => void; + onCancelSubmit: ( + shouldRestorePrompt?: boolean, + clearBuffer?: boolean, + ) => void; isShellFocused?: boolean; logger?: Logger | null; } @@ -120,13 +124,16 @@ export const useAgentStream = ({ } }, [addItem, pendingHistoryItemRef, setPendingHistoryItem]); - const cancelOngoingRequest = useCallback(async () => { - if (agent) { - await agent.abort(); - setStreamingState(StreamingState.Idle); - onCancelSubmit(false); - } - }, [agent, onCancelSubmit]); + const cancelOngoingRequest = useCallback( + async (clearBuffer: boolean = true) => { + if (agent) { + await agent.abort(); + setStreamingState(StreamingState.Idle); + onCancelSubmit(false, clearBuffer); + } + }, + [agent, onCancelSubmit], + ); // TODO: Support native handleApprovalModeChange for Plan Mode const handleApprovalModeChange = useCallback( @@ -322,6 +329,21 @@ export const useAgentStream = ({ return () => unsubscribe?.(); }, [agent, handleEvent]); + useKeypress( + (key) => { + if (key.name === 'escape' && !isShellFocused) { + void cancelOngoingRequest(false); + return true; + } + return false; + }, + { + isActive: + streamingState === StreamingState.Responding || + streamingState === StreamingState.WaitingForConfirmation, + }, + ); + const submitQuery = useCallback( async ( query: Part[] | string, diff --git a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx index d6c68ec880..53e7475218 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx +++ b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx @@ -1637,7 +1637,7 @@ describe('useGeminiStream', () => { simulateEscapeKeyPress(); - expect(cancelSubmitSpy).toHaveBeenCalledWith(false); + expect(cancelSubmitSpy).toHaveBeenCalledWith(false, false); }); it('should call setShellInputFocused(false) when escape is pressed', async () => { diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index eee0241a58..14f90ca4d0 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -227,7 +227,10 @@ export const useGeminiStream = ( performMemoryRefresh: () => Promise, modelSwitchedFromQuotaError: boolean, setModelSwitchedFromQuotaError: React.Dispatch>, - onCancelSubmit: (shouldRestorePrompt?: boolean) => void, + onCancelSubmit: ( + shouldRestorePrompt?: boolean, + clearBuffer?: boolean, + ) => void, setShellInputFocused: (value: boolean) => void, terminalWidth: number, terminalHeight: number, @@ -803,100 +806,129 @@ export const useGeminiStream = ( [addItem, config, isLowErrorVerbosity], ); - const cancelOngoingRequest = useCallback(() => { - if ( - streamingState !== StreamingState.Responding && - streamingState !== StreamingState.WaitingForConfirmation - ) { - return; - } - if (turnCancelledRef.current) { - return; - } - turnCancelledRef.current = true; - setRetryStatus(null); - - // A full cancellation means no tools have produced a final result yet. - // This determines if we show a generic "Request cancelled" message. - const isFullCancellation = !toolCalls.some( - (tc) => tc.status === 'success' || tc.status === 'error', - ); - - // Ensure we have an abort controller, creating one if it doesn't exist. - if (!abortControllerRef.current) { - abortControllerRef.current = new AbortController(); - } - - // The order is important here. - // 1. Fire the signal to interrupt any active async operations. - abortControllerRef.current.abort(); - // 2. Call the imperative cancel to clear the queue of pending tools. - cancelAllToolCalls(abortControllerRef.current.signal); - - if (pendingHistoryItemRef.current) { - const isShellCommand = - pendingHistoryItemRef.current.type === 'tool_group' && - pendingHistoryItemRef.current.tools.some( - (t) => t.name === SHELL_COMMAND_NAME, - ); - - // If it is a shell command, we update the status to Canceled and clear the output - // to avoid artifacts, then add it to history immediately. - if (isShellCommand) { - // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion - const toolGroup = pendingHistoryItemRef.current as HistoryItemToolGroup; - const updatedTools = toolGroup.tools.map((tool) => { - if (tool.name === SHELL_COMMAND_NAME) { - return { - ...tool, - status: CoreToolCallStatus.Cancelled, - resultDisplay: tool.resultDisplay, - }; - } - return tool; - }); - addItem({ ...toolGroup, tools: updatedTools } as HistoryItemWithoutId); - } else { - addItem(pendingHistoryItemRef.current); + const cancelOngoingRequest = useCallback( + (clearBuffer: boolean = true) => { + // If we are already cancelled, do nothing + if (turnCancelledRef.current) { + if (clearBuffer) { + onCancelSubmit(false, true); + } + return; } - } - setPendingHistoryItem(null); - // If it was a full cancellation, add the info message now. - // Otherwise, we let handleCompletedTools figure out the next step, - // which might involve sending partial results back to the model. - if (isFullCancellation) { - // If shell is active, we delay this message to ensure correct ordering - // (Shell item first, then Info message). - if (!activeShellPtyId) { - addItem({ - type: MessageType.INFO, - text: 'Request cancelled.', - }); - setIsResponding(false); + const hasActiveTools = toolCalls.some( + (tc) => + tc.status === CoreToolCallStatus.Executing || + tc.status === CoreToolCallStatus.Scheduled || + tc.status === CoreToolCallStatus.Validating, + ); + + // If we are not responding, not waiting for confirmation, and have no active tools, + // there is nothing to abort. + if ( + streamingState === StreamingState.Idle && + !isRespondingRef.current && + !hasActiveTools + ) { + // Even if we are "idle", if we are called with clearBuffer=true (Ctrl+C), + // we still want to clear the buffer. + if (clearBuffer) { + onCancelSubmit(false, true); + } + return; } - } - onCancelSubmit(false); - setShellInputFocused(false); - }, [ - streamingState, - addItem, - setPendingHistoryItem, - onCancelSubmit, - pendingHistoryItemRef, - setShellInputFocused, - cancelAllToolCalls, - toolCalls, - activeShellPtyId, - setIsResponding, - ]); + turnCancelledRef.current = true; + setRetryStatus(null); + + // A full cancellation means no tools have produced a final result yet. + // This determines if we show a generic "Request cancelled" message. + const isFullCancellation = !toolCalls.some( + (tc) => tc.status === 'success' || tc.status === 'error', + ); + + // Ensure we have an abort controller, creating one if it doesn't exist. + if (!abortControllerRef.current) { + abortControllerRef.current = new AbortController(); + } + + // The order is important here. + // 1. Fire the signal to interrupt any active async operations. + abortControllerRef.current.abort(); + // 2. Call the imperative cancel to clear the queue of pending tools. + cancelAllToolCalls(abortControllerRef.current.signal); + + if (pendingHistoryItemRef.current) { + // If it is a shell command, we update the status to Canceled and clear the output + // to avoid artifacts, then add it to history immediately. + if ( + pendingHistoryItemRef.current.type === 'tool_group' && + pendingHistoryItemRef.current.tools.some( + (t) => t.name === SHELL_COMMAND_NAME, + ) + ) { + const toolGroup = pendingHistoryItemRef.current; + const updatedTools = toolGroup.tools.map((tool) => { + if (tool.name === SHELL_COMMAND_NAME) { + return { + ...tool, + status: CoreToolCallStatus.Cancelled, + resultDisplay: tool.resultDisplay, + }; + } + return tool; + }); + const newToolGroup: HistoryItemToolGroup = { + ...toolGroup, + tools: updatedTools, + }; + addItem(newToolGroup); + } else { + addItem(pendingHistoryItemRef.current); + } + } + setPendingHistoryItem(null); + + // If it was a full cancellation, add the info message now. + // Otherwise, we let handleCompletedTools figure out the next step, + // which might involve sending partial results back to the model. + if (isFullCancellation) { + // If shell is active, we delay this message to ensure correct ordering + // (Shell item first, then Info message). + if (!activeShellPtyId) { + addItem({ + type: MessageType.INFO, + text: 'Request cancelled.', + }); + setIsResponding(false); + } + } + + onCancelSubmit(false, clearBuffer); + setShellInputFocused(false); + }, + [ + streamingState, + addItem, + setPendingHistoryItem, + onCancelSubmit, + pendingHistoryItemRef, + isRespondingRef, + setShellInputFocused, + cancelAllToolCalls, + toolCalls, + activeShellPtyId, + setIsResponding, + ], + ); useKeypress( (key) => { if (key.name === 'escape' && !isShellFocused) { - cancelOngoingRequest(); + cancelOngoingRequest(false); + return true; } + return false; }, { isActive: