feat(cli): continue request after disabling loop detection (#11416)

packages/cli/src/ui/hooks/useGeminiStream.test.tsx

@@ -2677,7 +2677,8 @@ describe('useGeminiStream', () => {
       };
       mockConfig.getGeminiClient = vi.fn().mockReturnValue(mockClient);
 
-      mockSendMessageStream.mockReturnValue(
+      // Mock for the initial request
+      mockSendMessageStream.mockReturnValueOnce(
         (async function* () {
           yield {
             type: ServerGeminiEventType.LoopDetected,
@@ -2685,6 +2686,20 @@ describe('useGeminiStream', () => {
         })(),
       );
 
+      // Mock for the retry request
+      mockSendMessageStream.mockReturnValueOnce(
+        (async function* () {
+          yield {
+            type: ServerGeminiEventType.Content,
+            value: 'Retry successful',
+          };
+          yield {
+            type: ServerGeminiEventType.Finished,
+            value: { reason: 'STOP' },
+          };
+        })(),
+      );
+
       const { result } = renderTestHook();
 
       await act(async () => {
@@ -2715,10 +2730,21 @@ describe('useGeminiStream', () => {
       expect(mockAddItem).toHaveBeenCalledWith(
         {
           type: 'info',
-          text: 'Loop detection has been disabled for this session. Please try your request again.',
+          text: 'Loop detection has been disabled for this session. Retrying request...',
         },
         expect.any(Number),
       );
+
+      // Verify that the request was retried
+      await waitFor(() => {
+        expect(mockSendMessageStream).toHaveBeenCalledTimes(2);
+        expect(mockSendMessageStream).toHaveBeenNthCalledWith(
+          2,
+          'test query',
+          expect.any(AbortSignal),
+          expect.any(String),
+        );
+      });
     });
 
     it('should keep loop detection enabled and show message when user selects "keep"', async () => {
@@ -2771,6 +2797,9 @@ describe('useGeminiStream', () => {
         },
         expect.any(Number),
       );
+
+      // Verify that the request was NOT retried
+      expect(mockSendMessageStream).toHaveBeenCalledTimes(1);
     });
 
     it('should handle multiple loop detection events properly', async () => {
@@ -2821,6 +2850,20 @@ describe('useGeminiStream', () => {
         })(),
       );
 
+      // Mock for the retry request
+      mockSendMessageStream.mockReturnValueOnce(
+        (async function* () {
+          yield {
+            type: ServerGeminiEventType.Content,
+            value: 'Retry successful',
+          };
+          yield {
+            type: ServerGeminiEventType.Finished,
+            value: { reason: 'STOP' },
+          };
+        })(),
+      );
+
       // Second loop detection
       await act(async () => {
         await result.current.submitQuery('second query');
@@ -2843,10 +2886,21 @@ describe('useGeminiStream', () => {
       expect(mockAddItem).toHaveBeenCalledWith(
         {
           type: 'info',
-          text: 'Loop detection has been disabled for this session. Please try your request again.',
+          text: 'Loop detection has been disabled for this session. Retrying request...',
         },
         expect.any(Number),
       );
+
+      // Verify that the request was retried
+      await waitFor(() => {
+        expect(mockSendMessageStream).toHaveBeenCalledTimes(3); // 1st query, 2nd query, retry of 2nd query
+        expect(mockSendMessageStream).toHaveBeenNthCalledWith(
+          3,
+          'second query',
+          expect.any(AbortSignal),
+          expect.any(String),
+        );
+      });
     });
 
     it('should process LoopDetected event after moving pending history to history', async () => {

packages/cli/src/ui/hooks/useGeminiStream.ts

@@ -185,6 +185,8 @@ export const useGeminiStream = (
     return undefined;
   }, [toolCalls]);
 
+  const lastQueryRef = useRef<PartListUnion | null>(null);
+  const lastPromptIdRef = useRef<string | null>(null);
   const loopDetectedRef = useRef(false);
   const [
     loopDetectionConfirmationRequest,
@@ -668,39 +670,6 @@ export const useGeminiStream = (
     [addItem, onCancelSubmit, config],
   );
 
-  const handleLoopDetectionConfirmation = useCallback(
-    (result: { userSelection: 'disable' | 'keep' }) => {
-      setLoopDetectionConfirmationRequest(null);
-
-      if (result.userSelection === 'disable') {
-        config.getGeminiClient().getLoopDetectionService().disableForSession();
-        addItem(
-          {
-            type: 'info',
-            text: `Loop detection has been disabled for this session. Please try your request again.`,
-          },
-          Date.now(),
-        );
-      } else {
-        addItem(
-          {
-            type: 'info',
-            text: `A potential loop was detected. This can happen due to repetitive tool calls or other model behavior. The request has been halted.`,
-          },
-          Date.now(),
-        );
-      }
-    },
-    [config, addItem],
-  );
-
-  const handleLoopDetectedEvent = useCallback(() => {
-    // Show the confirmation dialog to choose whether to disable loop detection
-    setLoopDetectionConfirmationRequest({
-      onComplete: handleLoopDetectionConfirmation,
-    });
-  }, [handleLoopDetectionConfirmation]);
-
   const processGeminiStreamEvents = useCallback(
     async (
       stream: AsyncIterable<GeminiEvent>,
@@ -850,6 +819,10 @@ export const useGeminiStream = (
         setIsResponding(true);
         setInitError(null);
 
+        // Store query and prompt_id for potential retry on loop detection
+        lastQueryRef.current = queryToSend;
+        lastPromptIdRef.current = prompt_id;
+
         try {
           const stream = geminiClient.sendMessageStream(
             queryToSend,
@@ -872,7 +845,42 @@ export const useGeminiStream = (
           }
           if (loopDetectedRef.current) {
             loopDetectedRef.current = false;
-            handleLoopDetectedEvent();
+            // Show the confirmation dialog to choose whether to disable loop detection
+            setLoopDetectionConfirmationRequest({
+              onComplete: (result: { userSelection: 'disable' | 'keep' }) => {
+                setLoopDetectionConfirmationRequest(null);
+
+                if (result.userSelection === 'disable') {
+                  config
+                    .getGeminiClient()
+                    .getLoopDetectionService()
+                    .disableForSession();
+                  addItem(
+                    {
+                      type: 'info',
+                      text: `Loop detection has been disabled for this session. Retrying request...`,
+                    },
+                    Date.now(),
+                  );
+
+                  if (lastQueryRef.current && lastPromptIdRef.current) {
+                    submitQuery(
+                      lastQueryRef.current,
+                      { isContinuation: true },
+                      lastPromptIdRef.current,
+                    );
+                  }
+                } else {
+                  addItem(
+                    {
+                      type: 'info',
+                      text: `A potential loop was detected. This can happen due to repetitive tool calls or other model behavior. The request has been halted.`,
+                    },
+                    Date.now(),
+                  );
+                }
+              },
+            });
           }
         } catch (error: unknown) {
           if (error instanceof UnauthorizedError) {
@@ -911,7 +919,6 @@ export const useGeminiStream = (
       config,
       startNewPrompt,
       getPromptCount,
-      handleLoopDetectedEvent,
     ],
   );