From a5ea113a8e485f42cc1136b6c57e337cbf369c57 Mon Sep 17 00:00:00 2001 From: Neha Prasad Date: Mon, 28 Jul 2025 23:27:33 +0530 Subject: [PATCH] fix: Clear previous thoughts when starting new prompts (#4966) --- .../cli/src/ui/hooks/useGeminiStream.test.tsx | 195 ++++++++++++++++++ packages/cli/src/ui/hooks/useGeminiStream.ts | 7 +- 2 files changed, 200 insertions(+), 2 deletions(-) diff --git a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx index 02fae607..085e3e96 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx +++ b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx @@ -319,6 +319,7 @@ describe('useGeminiStream', () => { }, setQuotaErrorOccurred: vi.fn(), getQuotaErrorOccurred: vi.fn(() => false), + getModel: vi.fn(() => 'gemini-2.5-pro'), getContentGeneratorConfig: vi .fn() .mockReturnValue(contentGeneratorConfig), @@ -1473,4 +1474,198 @@ describe('useGeminiStream', () => { } }); }); + + describe('Thought Reset', () => { + it('should reset thought to null when starting a new prompt', async () => { + // First, simulate a response with a thought + mockSendMessageStream.mockReturnValue( + (async function* () { + yield { + type: ServerGeminiEventType.Thought, + value: { + subject: 'Previous thought', + description: 'Old description', + }, + }; + yield { + type: ServerGeminiEventType.Content, + value: 'Some response content', + }; + yield { type: ServerGeminiEventType.Finished, value: 'STOP' }; + })(), + ); + + const { result } = renderHook(() => + useGeminiStream( + new MockedGeminiClientClass(mockConfig), + [], + mockAddItem, + mockSetShowHelp, + mockConfig, + mockOnDebugMessage, + mockHandleSlashCommand, + false, + () => 'vscode' as EditorType, + () => {}, + () => Promise.resolve(), + false, + () => {}, + ), + ); + + // Submit first query to set a thought + await act(async () => { + await result.current.submitQuery('First query'); + }); + + // Wait for the first response to complete + await waitFor(() => { + expect(mockAddItem).toHaveBeenCalledWith( + expect.objectContaining({ + type: 'gemini', + text: 'Some response content', + }), + expect.any(Number), + ); + }); + + // Now simulate a new response without a thought + mockSendMessageStream.mockReturnValue( + (async function* () { + yield { + type: ServerGeminiEventType.Content, + value: 'New response content', + }; + yield { type: ServerGeminiEventType.Finished, value: 'STOP' }; + })(), + ); + + // Submit second query - thought should be reset + await act(async () => { + await result.current.submitQuery('Second query'); + }); + + // The thought should be reset to null when starting the new prompt + // We can verify this by checking that the LoadingIndicator would not show the previous thought + // The actual thought state is internal to the hook, but we can verify the behavior + // by ensuring the second response doesn't show the previous thought + await waitFor(() => { + expect(mockAddItem).toHaveBeenCalledWith( + expect.objectContaining({ + type: 'gemini', + text: 'New response content', + }), + expect.any(Number), + ); + }); + }); + + it('should reset thought to null when user cancels', async () => { + // Mock a stream that yields a thought then gets cancelled + mockSendMessageStream.mockReturnValue( + (async function* () { + yield { + type: ServerGeminiEventType.Thought, + value: { subject: 'Some thought', description: 'Description' }, + }; + yield { type: ServerGeminiEventType.UserCancelled }; + })(), + ); + + const { result } = renderHook(() => + useGeminiStream( + new MockedGeminiClientClass(mockConfig), + [], + mockAddItem, + mockSetShowHelp, + mockConfig, + mockOnDebugMessage, + mockHandleSlashCommand, + false, + () => 'vscode' as EditorType, + () => {}, + () => Promise.resolve(), + false, + () => {}, + ), + ); + + // Submit query + await act(async () => { + await result.current.submitQuery('Test query'); + }); + + // Verify cancellation message was added + await waitFor(() => { + expect(mockAddItem).toHaveBeenCalledWith( + expect.objectContaining({ + type: 'info', + text: 'User cancelled the request.', + }), + expect.any(Number), + ); + }); + + // Verify state is reset to idle + expect(result.current.streamingState).toBe(StreamingState.Idle); + }); + + it('should reset thought to null when there is an error', async () => { + // Mock a stream that yields a thought then encounters an error + mockSendMessageStream.mockReturnValue( + (async function* () { + yield { + type: ServerGeminiEventType.Thought, + value: { subject: 'Some thought', description: 'Description' }, + }; + yield { + type: ServerGeminiEventType.Error, + value: { error: { message: 'Test error' } }, + }; + })(), + ); + + const { result } = renderHook(() => + useGeminiStream( + new MockedGeminiClientClass(mockConfig), + [], + mockAddItem, + mockSetShowHelp, + mockConfig, + mockOnDebugMessage, + mockHandleSlashCommand, + false, + () => 'vscode' as EditorType, + () => {}, + () => Promise.resolve(), + false, + () => {}, + ), + ); + + // Submit query + await act(async () => { + await result.current.submitQuery('Test query'); + }); + + // Verify error message was added + await waitFor(() => { + expect(mockAddItem).toHaveBeenCalledWith( + expect.objectContaining({ + type: 'error', + }), + expect.any(Number), + ); + }); + + // Verify parseAndFormatApiError was called + expect(mockParseAndFormatApiError).toHaveBeenCalledWith( + { message: 'Test error' }, + expect.any(String), + undefined, + 'gemini-2.5-pro', + 'gemini-2.5-flash', + ); + }); + }); }); diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index 456c0fb7..c934a139 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -414,8 +414,9 @@ export const useGeminiStream = ( userMessageTimestamp, ); setIsResponding(false); + setThought(null); // Reset thought when user cancels }, - [addItem, pendingHistoryItemRef, setPendingHistoryItem], + [addItem, pendingHistoryItemRef, setPendingHistoryItem, setThought], ); const handleErrorEvent = useCallback( @@ -437,8 +438,9 @@ export const useGeminiStream = ( }, userMessageTimestamp, ); + setThought(null); // Reset thought when there's an error }, - [addItem, pendingHistoryItemRef, setPendingHistoryItem, config], + [addItem, pendingHistoryItemRef, setPendingHistoryItem, config, setThought], ); const handleFinishedEvent = useCallback( @@ -637,6 +639,7 @@ export const useGeminiStream = ( if (!options?.isContinuation) { startNewPrompt(); + setThought(null); // Reset thought when starting a new prompt } setIsResponding(true);