bug: fix cancel after a tool has been used (#1270)

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,