From dd3b1cb653e30e9aaeb4a22764e34a38922e716d Mon Sep 17 00:00:00 2001 From: Sandy Tao Date: Tue, 21 Oct 2025 13:27:57 -0700 Subject: [PATCH] feat(cli): continue request after disabling loop detection (#11416) --- .../cli/src/ui/hooks/useGeminiStream.test.tsx | 60 ++++++++++++++- packages/cli/src/ui/hooks/useGeminiStream.ts | 77 ++++++++++--------- 2 files changed, 99 insertions(+), 38 deletions(-) diff --git a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx index f10d0f41c1..02db0f466e 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx +++ b/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 () => { diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index e391f6c5da..a0190a3c4b 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -185,6 +185,8 @@ export const useGeminiStream = ( return undefined; }, [toolCalls]); + const lastQueryRef = useRef(null); + const lastPromptIdRef = useRef(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, @@ -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, ], );