diff --git a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx index 1f2ef5f90c..a1251f4143 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx +++ b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx @@ -3510,6 +3510,116 @@ describe('useGeminiStream', () => { expect(result.current.loopDetectionConfirmationRequest).not.toBeNull(); }); }); + + describe('Race Condition Prevention', () => { + it('should reject concurrent submitQuery when already responding', async () => { + // Stream that stays open (simulates "still responding") + mockSendMessageStream.mockReturnValue( + (async function* () { + yield { + type: ServerGeminiEventType.Content, + value: 'First response', + }; + // Keep the stream open + await new Promise(() => {}); + })(), + ); + + const { result } = renderTestHook(); + + // Start first query without awaiting (fire-and-forget, like existing tests) + await act(async () => { + // eslint-disable-next-line @typescript-eslint/no-floating-promises + result.current.submitQuery('first query'); + }); + + // Wait for the stream to start responding + await waitFor(() => { + expect(result.current.streamingState).toBe(StreamingState.Responding); + }); + + // Try a second query while first is still responding + await act(async () => { + // eslint-disable-next-line @typescript-eslint/no-floating-promises + result.current.submitQuery('second query'); + }); + + // Should have only called sendMessageStream once (second was rejected) + expect(mockSendMessageStream).toHaveBeenCalledTimes(1); + }); + + it('should allow continuation queries via loop detection retry', async () => { + const mockLoopDetectionService = { + disableForSession: vi.fn(), + }; + const mockClient = { + ...new MockedGeminiClientClass(mockConfig), + getLoopDetectionService: () => mockLoopDetectionService, + }; + mockConfig.getGeminiClient = vi.fn().mockReturnValue(mockClient); + + // First call triggers loop detection + mockSendMessageStream.mockReturnValueOnce( + (async function* () { + yield { + type: ServerGeminiEventType.LoopDetected, + }; + })(), + ); + + // Retry call succeeds + mockSendMessageStream.mockReturnValueOnce( + (async function* () { + yield { + type: ServerGeminiEventType.Content, + value: 'Retry success', + }; + yield { + type: ServerGeminiEventType.Finished, + value: { reason: 'STOP' }, + }; + })(), + ); + + const { result } = renderTestHook(); + + await act(async () => { + await result.current.submitQuery('test query'); + }); + + await waitFor(() => { + expect( + result.current.loopDetectionConfirmationRequest, + ).not.toBeNull(); + }); + + // User selects "disable" which triggers a continuation query + await act(async () => { + result.current.loopDetectionConfirmationRequest?.onComplete({ + userSelection: 'disable', + }); + }); + + // Verify disableForSession was called + expect( + mockLoopDetectionService.disableForSession, + ).toHaveBeenCalledTimes(1); + + // Continuation query should have gone through (2 total calls) + await waitFor(() => { + expect(mockSendMessageStream).toHaveBeenCalledTimes(2); + expect(mockSendMessageStream).toHaveBeenNthCalledWith( + 2, + 'test query', + expect.any(AbortSignal), + expect.any(String), + undefined, + false, + 'test query', + ); + }); + }); + }); }); describe('Agent Execution Events', () => { diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index d254902a94..d2e485db1f 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -216,7 +216,15 @@ export const useGeminiStream = ( const previousApprovalModeRef = useRef( config.getApprovalMode(), ); - const [isResponding, setIsResponding] = useState(false); + const [isResponding, setIsRespondingState] = useState(false); + const isRespondingRef = useRef(false); + const setIsResponding = useCallback( + (value: boolean) => { + setIsRespondingState(value); + isRespondingRef.current = value; + }, + [setIsRespondingState], + ); const [thought, thoughtRef, setThought] = useStateAndRef(null); const [pendingHistoryItem, pendingHistoryItemRef, setPendingHistoryItem] = @@ -320,11 +328,14 @@ export const useGeminiStream = ( return (executingShellTool as TrackedExecutingToolCall | undefined)?.pid; }, [toolCalls]); - const onExec = useCallback(async (done: Promise) => { - setIsResponding(true); - await done; - setIsResponding(false); - }, []); + const onExec = useCallback( + async (done: Promise) => { + setIsResponding(true); + await done; + setIsResponding(false); + }, + [setIsResponding], + ); const { handleShellCommand, @@ -538,7 +549,7 @@ export const useGeminiStream = ( setIsResponding(false); } prevActiveShellPtyIdRef.current = activeShellPtyId; - }, [activeShellPtyId, addItem]); + }, [activeShellPtyId, addItem, setIsResponding]); useEffect(() => { if ( @@ -700,6 +711,7 @@ export const useGeminiStream = ( cancelAllToolCalls, toolCalls, activeShellPtyId, + setIsResponding, ]); useKeypress( @@ -952,7 +964,13 @@ export const useGeminiStream = ( setIsResponding(false); setThought(null); // Reset thought when user cancels }, - [addItem, pendingHistoryItemRef, setPendingHistoryItem, setThought], + [ + addItem, + pendingHistoryItemRef, + setPendingHistoryItem, + setThought, + setIsResponding, + ], ); const handleErrorEvent = useCallback( @@ -1358,14 +1376,15 @@ export const useGeminiStream = ( async ({ metadata: spanMetadata }) => { spanMetadata.input = query; - const queryId = `${Date.now()}-${Math.random()}`; - activeQueryIdRef.current = queryId; if ( - (streamingState === StreamingState.Responding || + (isRespondingRef.current || + streamingState === StreamingState.Responding || streamingState === StreamingState.WaitingForConfirmation) && !options?.isContinuation ) return; + const queryId = `${Date.now()}-${Math.random()}`; + activeQueryIdRef.current = queryId; const userMessageTimestamp = Date.now(); @@ -1452,7 +1471,7 @@ export const useGeminiStream = ( loopDetectedRef.current = false; // Show the confirmation dialog to choose whether to disable loop detection setLoopDetectionConfirmationRequest({ - onComplete: (result: { + onComplete: async (result: { userSelection: 'disable' | 'keep'; }) => { setLoopDetectionConfirmationRequest(null); @@ -1468,8 +1487,7 @@ export const useGeminiStream = ( }); if (lastQueryRef.current && lastPromptIdRef.current) { - // eslint-disable-next-line @typescript-eslint/no-floating-promises - submitQuery( + await submitQuery( lastQueryRef.current, { isContinuation: true }, lastPromptIdRef.current, @@ -1537,6 +1555,7 @@ export const useGeminiStream = ( maybeAddSuppressedToolErrorNote, maybeAddLowVerbosityFailureNote, settings.merged.billing?.overageStrategy, + setIsResponding, ], ); @@ -1803,6 +1822,7 @@ export const useGeminiStream = ( isLowErrorVerbosity, maybeAddSuppressedToolErrorNote, maybeAddLowVerbosityFailureNote, + setIsResponding, ], );