diff options
Diffstat (limited to 'packages/cli/src/ui/hooks/useGeminiStream.ts')
| -rw-r--r-- | packages/cli/src/ui/hooks/useGeminiStream.ts | 36 |
1 files changed, 29 insertions, 7 deletions
diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index 234652db..fcfa1c57 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -92,6 +92,7 @@ export const useGeminiStream = ( ) => { const [initError, setInitError] = useState<string | null>(null); const abortControllerRef = useRef<AbortController | null>(null); + const turnCancelledRef = useRef(false); const [isResponding, setIsResponding] = useState<boolean>(false); const [thought, setThought] = useState<ThoughtSummary | null>(null); const [pendingHistoryItemRef, setPendingHistoryItem] = @@ -168,15 +169,25 @@ export const useGeminiStream = ( return StreamingState.Idle; }, [isResponding, toolCalls]); - useEffect(() => { - if (streamingState === StreamingState.Idle) { - abortControllerRef.current = null; - } - }, [streamingState]); - useInput((_input, key) => { - if (streamingState !== StreamingState.Idle && key.escape) { + if (streamingState === StreamingState.Responding && key.escape) { + if (turnCancelledRef.current) { + return; + } + turnCancelledRef.current = true; abortControllerRef.current?.abort(); + if (pendingHistoryItemRef.current) { + addItem(pendingHistoryItemRef.current, Date.now()); + } + addItem( + { + type: MessageType.INFO, + text: 'Request cancelled.', + }, + Date.now(), + ); + setPendingHistoryItem(null); + setIsResponding(false); } }); @@ -189,6 +200,9 @@ export const useGeminiStream = ( queryToSend: PartListUnion | null; shouldProceed: boolean; }> => { + if (turnCancelledRef.current) { + return { queryToSend: null, shouldProceed: false }; + } if (typeof query === 'string' && query.trim().length === 0) { return { queryToSend: null, shouldProceed: false }; } @@ -285,6 +299,10 @@ export const useGeminiStream = ( currentGeminiMessageBuffer: string, userMessageTimestamp: number, ): string => { + if (turnCancelledRef.current) { + // Prevents additional output after a user initiated cancel. + return ''; + } let newGeminiMessageBuffer = currentGeminiMessageBuffer + eventValue; if ( pendingHistoryItemRef.current?.type !== 'gemini' && @@ -335,6 +353,9 @@ export const useGeminiStream = ( const handleUserCancelledEvent = useCallback( (userMessageTimestamp: number) => { + if (turnCancelledRef.current) { + return; + } if (pendingHistoryItemRef.current) { if (pendingHistoryItemRef.current.type === 'tool_group') { const updatedTools = pendingHistoryItemRef.current.tools.map( @@ -469,6 +490,7 @@ export const useGeminiStream = ( abortControllerRef.current = new AbortController(); const abortSignal = abortControllerRef.current.signal; + turnCancelledRef.current = false; const { queryToSend, shouldProceed } = await prepareQueryForGemini( query, |
