@google/gemini-cli 0.12.0-nightly.20251023.c4c0c0d1 → 0.12.0-preview.0

package/dist/src/ui/hooks/useGeminiStream.test.js

@@ -4,12 +4,13 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 import { describe, it, expect, vi, beforeEach } from 'vitest';
-import { renderHook, act, waitFor } from '@testing-library/react';
+import { act } from 'react';
+import { renderHook } from '../../test-utils/render.js';
 import { useGeminiStream } from './useGeminiStream.js';
 import { useKeypress } from './useKeypress.js';
 import * as atCommandProcessor from './atCommandProcessor.js';
 import { useReactToolScheduler } from './useReactToolScheduler.js';
-import { ApprovalMode, AuthType, GeminiEventType as ServerGeminiEventType, ToolErrorType, ToolConfirmationOutcome, tokenLimit, } from '@google/gemini-cli-core';
+import { ApprovalMode, AuthType, GeminiEventType as ServerGeminiEventType, ToolErrorType, ToolConfirmationOutcome, tokenLimit, debugLogger, } from '@google/gemini-cli-core';
 import { MessageType, StreamingState } from '../types.js';
 // --- MOCKS ---
 const mockSendMessageStream = vi
@@ -175,8 +176,9 @@ describe('useGeminiStream', () => {
         mockUseReactToolScheduler.mockReturnValue([
             [], // Default to empty array for toolCalls
             mockScheduleToolCalls,
-            mockCancelAllToolCalls,
             mockMarkToolsAsSubmitted,
+            vi.fn(), // setToolCallsForDisplay
+            mockCancelAllToolCalls,
         ]);
         // Reset mocks for GeminiClient instance methods (startChat and sendMessageStream)
         // The GeminiClient constructor itself is mocked at the module level.
@@ -197,35 +199,62 @@ describe('useGeminiStream', () => {
         setValue: vi.fn(),
     };
     const renderTestHook = (initialToolCalls = [], geminiClient) => {
-        let currentToolCalls = initialToolCalls;
-        const setToolCalls = (newToolCalls) => {
-            currentToolCalls = newToolCalls;
-        };
-        mockUseReactToolScheduler.mockImplementation(() => [
-            currentToolCalls,
-            mockScheduleToolCalls,
-            mockCancelAllToolCalls,
-            mockMarkToolsAsSubmitted,
-        ]);
         const client = geminiClient || mockConfig.getGeminiClient();
+        const initialProps = {
+            client,
+            history: [],
+            addItem: mockAddItem,
+            config: mockConfig,
+            onDebugMessage: mockOnDebugMessage,
+            handleSlashCommand: mockHandleSlashCommand,
+            shellModeActive: false,
+            loadedSettings: mockLoadedSettings,
+            toolCalls: initialToolCalls,
+        };
         const { result, rerender } = renderHook((props) => {
-            // Update the mock's return value if new toolCalls are passed in props
-            if (props.toolCalls) {
-                setToolCalls(props.toolCalls);
-            }
+            // This mock needs to be stateful. When setToolCallsForDisplay is called,
+            // it should trigger a rerender with the new state.
+            const mockSetToolCallsForDisplay = vi.fn((updater) => {
+                const newToolCalls = typeof updater === 'function' ? updater(props.toolCalls) : updater;
+                rerender({ ...props, toolCalls: newToolCalls });
+            });
+            // Create a stateful mock for cancellation that updates the toolCalls state.
+            const statefulCancelAllToolCalls = vi.fn((...args) => {
+                // Call the original spy so `toHaveBeenCalled` checks still work.
+                mockCancelAllToolCalls(...args);
+                const newToolCalls = props.toolCalls.map((tc) => {
+                    // Only cancel tools that are in a cancellable state.
+                    if (tc.status === 'awaiting_approval' ||
+                        tc.status === 'executing' ||
+                        tc.status === 'scheduled' ||
+                        tc.status === 'validating') {
+                        // A real cancelled tool call has a response object.
+                        // We need to simulate this to avoid type errors downstream.
+                        return {
+                            ...tc,
+                            status: 'cancelled',
+                            response: {
+                                callId: tc.request.callId,
+                                responseParts: [],
+                                resultDisplay: 'Request cancelled.',
+                            },
+                            responseSubmittedToGemini: true, // Mark as "processed"
+                        };
+                    }
+                    return tc;
+                });
+                rerender({ ...props, toolCalls: newToolCalls });
+            });
+            mockUseReactToolScheduler.mockImplementation(() => [
+                props.toolCalls,
+                mockScheduleToolCalls,
+                mockMarkToolsAsSubmitted,
+                mockSetToolCallsForDisplay,
+                statefulCancelAllToolCalls, // Use the stateful mock
+            ]);
             return useGeminiStream(props.client, props.history, props.addItem, props.config, props.loadedSettings, props.onDebugMessage, props.handleSlashCommand, props.shellModeActive, () => 'vscode', () => { }, () => Promise.resolve(), false, () => { }, () => { }, () => { }, () => { }, 80, 24);
         }, {
-            initialProps: {
-                client,
-                history: [],
-                addItem: mockAddItem,
-                config: mockConfig,
-                onDebugMessage: mockOnDebugMessage,
-                handleSlashCommand: mockHandleSlashCommand,
-                shellModeActive: false,
-                loadedSettings: mockLoadedSettings,
-                toolCalls: initialToolCalls,
-            },
+            initialProps,
         });
         return {
             result,
@@ -341,7 +370,7 @@ describe('useGeminiStream', () => {
         let capturedOnComplete = null;
         mockUseReactToolScheduler.mockImplementation((onComplete) => {
             capturedOnComplete = onComplete;
-            return [[], mockScheduleToolCalls, mockMarkToolsAsSubmitted];
+            return [[], mockScheduleToolCalls, mockMarkToolsAsSubmitted, vi.fn()];
         });
         renderHook(() => useGeminiStream(new MockedGeminiClientClass(mockConfig), [], mockAddItem, mockConfig, mockLoadedSettings, mockOnDebugMessage, mockHandleSlashCommand, false, () => 'vscode', () => { }, () => Promise.resolve(), false, () => { }, () => { }, () => { }, () => { }, 80, 24));
         // Trigger the onComplete callback with completed tools
@@ -350,7 +379,7 @@ describe('useGeminiStream', () => {
                 await capturedOnComplete(completedToolCalls);
             }
         });
-        await waitFor(() => {
+        await vi.waitFor(() => {
             expect(mockMarkToolsAsSubmitted).toHaveBeenCalledTimes(1);
             expect(mockSendMessageStream).toHaveBeenCalledTimes(1);
         });
@@ -390,7 +419,7 @@ describe('useGeminiStream', () => {
         let capturedOnComplete = null;
         mockUseReactToolScheduler.mockImplementation((onComplete) => {
             capturedOnComplete = onComplete;
-            return [[], mockScheduleToolCalls, mockMarkToolsAsSubmitted];
+            return [[], mockScheduleToolCalls, mockMarkToolsAsSubmitted, vi.fn()];
         });
         renderHook(() => useGeminiStream(client, [], mockAddItem, mockConfig, mockLoadedSettings, mockOnDebugMessage, mockHandleSlashCommand, false, () => 'vscode', () => { }, () => Promise.resolve(), false, () => { }, () => { }, () => { }, () => { }, 80, 24));
         // Trigger the onComplete callback with cancelled tools
@@ -399,7 +428,7 @@ describe('useGeminiStream', () => {
                 await capturedOnComplete(cancelledToolCalls);
             }
         });
-        await waitFor(() => {
+        await vi.waitFor(() => {
             expect(mockMarkToolsAsSubmitted).toHaveBeenCalledWith(['1']);
             expect(client.addHistory).toHaveBeenCalledWith({
                 role: 'user',
@@ -473,7 +502,7 @@ describe('useGeminiStream', () => {
         let capturedOnComplete = null;
         mockUseReactToolScheduler.mockImplementation((onComplete) => {
             capturedOnComplete = onComplete;
-            return [[], mockScheduleToolCalls, mockMarkToolsAsSubmitted];
+            return [[], mockScheduleToolCalls, mockMarkToolsAsSubmitted, vi.fn()];
         });
         renderHook(() => useGeminiStream(client, [], mockAddItem, mockConfig, mockLoadedSettings, mockOnDebugMessage, mockHandleSlashCommand, false, () => 'vscode', () => { }, () => Promise.resolve(), false, () => { }, () => { }, () => { }, () => { }, 80, 24));
         // Trigger the onComplete callback with multiple cancelled tools
@@ -482,7 +511,7 @@ describe('useGeminiStream', () => {
                 await capturedOnComplete(allCancelledTools);
             }
         });
-        await waitFor(() => {
+        await vi.waitFor(() => {
             // The tools should be marked as submitted locally
             expect(mockMarkToolsAsSubmitted).toHaveBeenCalledWith([
                 'cancel-1',
@@ -552,6 +581,7 @@ describe('useGeminiStream', () => {
                 currentToolCalls,
                 mockScheduleToolCalls,
                 mockMarkToolsAsSubmitted,
+                vi.fn(), // setToolCallsForDisplay
             ];
         });
         const { result, rerender } = renderHook(() => useGeminiStream(new MockedGeminiClientClass(mockConfig), [], mockAddItem, mockConfig, mockLoadedSettings, mockOnDebugMessage, mockHandleSlashCommand, false, () => 'vscode', () => { }, () => Promise.resolve(), false, () => { }, () => { }, () => { }, () => { }, 80, 24));
@@ -565,6 +595,7 @@ describe('useGeminiStream', () => {
                 completedToolCalls,
                 mockScheduleToolCalls,
                 mockMarkToolsAsSubmitted,
+                vi.fn(), // setToolCallsForDisplay
             ];
         });
         act(() => {
@@ -580,7 +611,7 @@ describe('useGeminiStream', () => {
             }
         });
         // 5. Wait for submitQuery to be called
-        await waitFor(() => {
+        await vi.waitFor(() => {
             expect(mockSendMessageStream).toHaveBeenCalledWith(toolCallResponseParts, expect.any(AbortSignal), 'prompt-id-4');
         });
         // 6. After submission, the state should remain Responding until the stream completes.
@@ -618,13 +649,13 @@ describe('useGeminiStream', () => {
                 result.current.submitQuery('test query');
             });
             // Wait for the first part of the response
-            await waitFor(() => {
+            await vi.waitFor(() => {
                 expect(result.current.streamingState).toBe(StreamingState.Responding);
             });
             // Simulate escape key press
             simulateEscapeKeyPress();
             // Verify cancellation message is added
-            await waitFor(() => {
+            await vi.waitFor(() => {
                 expect(mockAddItem).toHaveBeenCalledWith({
                     type: MessageType.INFO,
                     text: 'Request cancelled.',
@@ -690,7 +721,7 @@ describe('useGeminiStream', () => {
             await act(async () => {
                 result.current.submitQuery('long running query');
             });
-            await waitFor(() => {
+            await vi.waitFor(() => {
                 expect(result.current.streamingState).toBe(StreamingState.Responding);
             });
             // Cancel the request
@@ -707,7 +738,7 @@ describe('useGeminiStream', () => {
             // The final state should be idle after cancellation
             expect(result.current.streamingState).toBe(StreamingState.Idle);
         });
-        it('should not cancel if a tool call is in progress (not just responding)', async () => {
+        it('should cancel if a tool call is in progress', async () => {
             const toolCalls = [
                 {
                     request: { callId: 'call1', name: 'tool1', args: {} },
@@ -727,14 +758,64 @@ describe('useGeminiStream', () => {
                     liveOutput: '...',
                 },
             ];
-            const abortSpy = vi.spyOn(AbortController.prototype, 'abort');
             const { result } = renderTestHook(toolCalls);
             // State is `Responding` because a tool is running
             expect(result.current.streamingState).toBe(StreamingState.Responding);
             // Try to cancel
             simulateEscapeKeyPress();
-            // Nothing should happen because the state is not `Responding`
-            expect(abortSpy).not.toHaveBeenCalled();
+            // The cancel function should be called
+            expect(mockCancelAllToolCalls).toHaveBeenCalled();
+        });
+        it('should cancel a request when a tool is awaiting confirmation', async () => {
+            const mockOnConfirm = vi.fn().mockResolvedValue(undefined);
+            const toolCalls = [
+                {
+                    request: {
+                        callId: 'confirm-call',
+                        name: 'some_tool',
+                        args: {},
+                        isClientInitiated: false,
+                        prompt_id: 'prompt-id-1',
+                    },
+                    status: 'awaiting_approval',
+                    responseSubmittedToGemini: false,
+                    tool: {
+                        name: 'some_tool',
+                        description: 'a tool',
+                        build: vi.fn().mockImplementation((_) => ({
+                            getDescription: () => `Mock description`,
+                        })),
+                    },
+                    invocation: {
+                        getDescription: () => `Mock description`,
+                    },
+                    confirmationDetails: {
+                        type: 'edit',
+                        title: 'Confirm Edit',
+                        onConfirm: mockOnConfirm,
+                        fileName: 'file.txt',
+                        filePath: '/test/file.txt',
+                        fileDiff: 'fake diff',
+                        originalContent: 'old',
+                        newContent: 'new',
+                    },
+                },
+            ];
+            const { result } = renderTestHook(toolCalls);
+            // State is `WaitingForConfirmation` because a tool is awaiting approval
+            expect(result.current.streamingState).toBe(StreamingState.WaitingForConfirmation);
+            // Try to cancel
+            simulateEscapeKeyPress();
+            // The imperative cancel function should be called on the scheduler
+            expect(mockCancelAllToolCalls).toHaveBeenCalled();
+            // A cancellation message should be added to history
+            await vi.waitFor(() => {
+                expect(mockAddItem).toHaveBeenCalledWith(expect.objectContaining({
+                    text: 'Request cancelled.',
+                }), expect.any(Number));
+            });
+            // The final state should be idle
+            expect(result.current.streamingState).toBe(StreamingState.Idle);
         });
     });
     describe('Slash Command Handling', () => {
@@ -749,7 +830,7 @@ describe('useGeminiStream', () => {
             await act(async () => {
                 await result.current.submitQuery('/memory add "test fact"');
             });
-            await waitFor(() => {
+            await vi.waitFor(() => {
                 expect(mockScheduleToolCalls).toHaveBeenCalledWith([
                     expect.objectContaining({
                         name: 'save_memory',
@@ -769,7 +850,7 @@ describe('useGeminiStream', () => {
             await act(async () => {
                 await result.current.submitQuery('/help');
             });
-            await waitFor(() => {
+            await vi.waitFor(() => {
                 expect(mockHandleSlashCommand).toHaveBeenCalledWith('/help');
                 expect(mockScheduleToolCalls).not.toHaveBeenCalled();
                 expect(mockSendMessageStream).not.toHaveBeenCalled(); // No LLM call made
@@ -785,7 +866,7 @@ describe('useGeminiStream', () => {
             await act(async () => {
                 await result.current.submitQuery('/my-custom-command');
             });
-            await waitFor(() => {
+            await vi.waitFor(() => {
                 expect(mockHandleSlashCommand).toHaveBeenCalledWith('/my-custom-command');
                 expect(localMockSendMessageStream).not.toHaveBeenCalledWith('/my-custom-command', expect.anything(), expect.anything());
                 expect(localMockSendMessageStream).toHaveBeenCalledWith('This is the actual prompt from the command file.', expect.any(AbortSignal), expect.any(String));
@@ -802,7 +883,7 @@ describe('useGeminiStream', () => {
             await act(async () => {
                 await result.current.submitQuery('/emptycmd');
             });
-            await waitFor(() => {
+            await vi.waitFor(() => {
                 expect(mockHandleSlashCommand).toHaveBeenCalledWith('/emptycmd');
                 expect(localMockSendMessageStream).toHaveBeenCalledWith('', expect.any(AbortSignal), expect.any(String));
             });
@@ -812,7 +893,7 @@ describe('useGeminiStream', () => {
             await act(async () => {
                 await result.current.submitQuery('// This is a line comment');
             });
-            await waitFor(() => {
+            await vi.waitFor(() => {
                 expect(mockHandleSlashCommand).not.toHaveBeenCalled();
                 expect(localMockSendMessageStream).toHaveBeenCalledWith('// This is a line comment', expect.any(AbortSignal), expect.any(String));
             });
@@ -822,7 +903,7 @@ describe('useGeminiStream', () => {
             await act(async () => {
                 await result.current.submitQuery('/* This is a block comment */');
             });
-            await waitFor(() => {
+            await vi.waitFor(() => {
                 expect(mockHandleSlashCommand).not.toHaveBeenCalled();
                 expect(localMockSendMessageStream).toHaveBeenCalledWith('/* This is a block comment */', expect.any(AbortSignal), expect.any(String));
             });
@@ -832,7 +913,7 @@ describe('useGeminiStream', () => {
             await act(async () => {
                 await result.current.submitQuery('/about');
             });
-            await waitFor(() => {
+            await vi.waitFor(() => {
                 expect(mockHandleSlashCommand).not.toHaveBeenCalled();
             });
         });
@@ -871,7 +952,7 @@ describe('useGeminiStream', () => {
             let capturedOnComplete = null;
             mockUseReactToolScheduler.mockImplementation((onComplete) => {
                 capturedOnComplete = onComplete;
-                return [[], mockScheduleToolCalls, mockMarkToolsAsSubmitted];
+                return [[], mockScheduleToolCalls, mockMarkToolsAsSubmitted, vi.fn()];
             });
             renderHook(() => useGeminiStream(new MockedGeminiClientClass(mockConfig), [], mockAddItem, mockConfig, mockLoadedSettings, mockOnDebugMessage, mockHandleSlashCommand, false, () => 'vscode', () => { }, mockPerformMemoryRefresh, false, () => { }, () => { }, () => { }, () => { }, 80, 24));
             // Trigger the onComplete callback with the completed save_memory tool
@@ -880,7 +961,7 @@ describe('useGeminiStream', () => {
                     await capturedOnComplete([completedToolCall]);
                 }
             });
-            await waitFor(() => {
+            await vi.waitFor(() => {
                 expect(mockPerformMemoryRefresh).toHaveBeenCalledTimes(1);
             });
         });
@@ -908,7 +989,7 @@ describe('useGeminiStream', () => {
                 await result.current.submitQuery('test query');
             });
             // 3. Assertion
-            await waitFor(() => {
+            await vi.waitFor(() => {
                 expect(mockParseAndFormatApiError).toHaveBeenCalledWith('Rate limit exceeded', mockAuthType, undefined, 'gemini-2.5-pro', 'gemini-2.5-flash');
             });
         });
@@ -1129,8 +1210,8 @@ describe('useGeminiStream', () => {
             expect(mockOnConfirm).not.toHaveBeenCalled();
         });
         it('should handle errors gracefully when auto-approving tool calls', async () => {
-            const consoleSpy = vi
-                .spyOn(console, 'error')
+            const debuggerSpy = vi
+                .spyOn(debugLogger, 'warn')
                 .mockImplementation(() => { });
             const mockOnConfirmSuccess = vi.fn().mockResolvedValue(undefined);
             const mockOnConfirmError = vi
@@ -1206,8 +1287,8 @@ describe('useGeminiStream', () => {
             expect(mockOnConfirmSuccess).toHaveBeenCalledTimes(1);
             expect(mockOnConfirmError).toHaveBeenCalledTimes(1);
             // Error should be logged
-            expect(consoleSpy).toHaveBeenCalledWith('Failed to auto-approve tool call call2:', expect.any(Error));
-            consoleSpy.mockRestore();
+            expect(debuggerSpy).toHaveBeenCalledWith('Failed to auto-approve tool call call2:', expect.any(Error));
+            debuggerSpy.mockRestore();
         });
         it('should skip tool calls without confirmationDetails', async () => {
             const awaitingApprovalToolCalls = [
@@ -1363,7 +1444,7 @@ describe('useGeminiStream', () => {
                 await result.current.submitQuery('Generate long text');
             });
             // Check that the info message was added
-            await waitFor(() => {
+            await vi.waitFor(() => {
                 expect(mockAddItem).toHaveBeenCalledWith({
                     type: 'info',
                     text: '⚠️  Response truncated due to token limits.',
@@ -1392,7 +1473,7 @@ describe('useGeminiStream', () => {
                     await result.current.submitQuery('Test overflow');
                 });
                 // Check that the message was added without suggestion
-                await waitFor(() => {
+                await vi.waitFor(() => {
                     expect(mockAddItem).toHaveBeenCalledWith({
                         type: 'info',
                         text: `Sending this message (20 tokens) might exceed the remaining context window limit (80 tokens).`,
@@ -1417,7 +1498,7 @@ describe('useGeminiStream', () => {
                     await result.current.submitQuery('Test overflow');
                 });
                 // Check that the message was added with suggestion
-                await waitFor(() => {
+                await vi.waitFor(() => {
                     expect(mockAddItem).toHaveBeenCalledWith({
                         type: 'info',
                         text: `Sending this message (30 tokens) might exceed the remaining context window limit (70 tokens). Please try reducing the size of your message or use the \`/compress\` command to compress the chat history.`,
@@ -1443,7 +1524,7 @@ describe('useGeminiStream', () => {
                 await result.current.submitQuery('Test overflow');
             });
             // Check that onCancelSubmit was called
-            await waitFor(() => {
+            await vi.waitFor(() => {
                 expect(onCancelSubmitSpy).toHaveBeenCalled();
             });
         });
@@ -1553,7 +1634,7 @@ describe('useGeminiStream', () => {
                 await act(async () => {
                     await result.current.submitQuery(`Test ${reason}`);
                 });
-                await waitFor(() => {
+                await vi.waitFor(() => {
                     expect(mockAddItem).toHaveBeenCalledWith({
                         type: 'info',
                         text: message,
@@ -1626,7 +1707,7 @@ describe('useGeminiStream', () => {
                 await result.current.submitQuery('First query');
             });
             // Wait for the first response to complete
-            await waitFor(() => {
+            await vi.waitFor(() => {
                 expect(mockAddItem).toHaveBeenCalledWith(expect.objectContaining({
                     type: 'gemini',
                     text: 'Some response content',
@@ -1651,7 +1732,7 @@ describe('useGeminiStream', () => {
             // 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(() => {
+            await vi.waitFor(() => {
                 expect(mockAddItem).toHaveBeenCalledWith(expect.objectContaining({
                     type: 'gemini',
                     text: 'New response content',
@@ -1710,7 +1791,7 @@ describe('useGeminiStream', () => {
                 await result.current.submitQuery('Test query');
             });
             // Verify cancellation message was added
-            await waitFor(() => {
+            await vi.waitFor(() => {
                 expect(mockAddItem).toHaveBeenCalledWith(expect.objectContaining({
                     type: 'info',
                     text: 'User cancelled the request.',
@@ -1737,7 +1818,7 @@ describe('useGeminiStream', () => {
                 await result.current.submitQuery('Test query');
             });
             // Verify error message was added
-            await waitFor(() => {
+            await vi.waitFor(() => {
                 expect(mockAddItem).toHaveBeenCalledWith(expect.objectContaining({
                     type: 'error',
                 }), expect.any(Number));
@@ -1771,7 +1852,7 @@ describe('useGeminiStream', () => {
             await act(async () => {
                 await result.current.submitQuery('test query');
             });
-            await waitFor(() => {
+            await vi.waitFor(() => {
                 expect(result.current.loopDetectionConfirmationRequest).not.toBeNull();
                 expect(typeof result.current.loopDetectionConfirmationRequest?.onComplete).toBe('function');
             });
@@ -1807,7 +1888,7 @@ describe('useGeminiStream', () => {
                 await result.current.submitQuery('test query');
             });
             // Wait for confirmation request to be set
-            await waitFor(() => {
+            await vi.waitFor(() => {
                 expect(result.current.loopDetectionConfirmationRequest).not.toBeNull();
             });
             // Simulate user selecting "disable"
@@ -1826,7 +1907,7 @@ describe('useGeminiStream', () => {
                 text: 'Loop detection has been disabled for this session. Retrying request...',
             }, expect.any(Number));
             // Verify that the request was retried
-            await waitFor(() => {
+            await vi.waitFor(() => {
                 expect(mockSendMessageStream).toHaveBeenCalledTimes(2);
                 expect(mockSendMessageStream).toHaveBeenNthCalledWith(2, 'test query', expect.any(AbortSignal), expect.any(String));
             });
@@ -1850,7 +1931,7 @@ describe('useGeminiStream', () => {
                 await result.current.submitQuery('test query');
             });
             // Wait for confirmation request to be set
-            await waitFor(() => {
+            await vi.waitFor(() => {
                 expect(result.current.loopDetectionConfirmationRequest).not.toBeNull();
             });
             // Simulate user selecting "keep"
@@ -1883,7 +1964,7 @@ describe('useGeminiStream', () => {
             await act(async () => {
                 await result.current.submitQuery('first query');
             });
-            await waitFor(() => {
+            await vi.waitFor(() => {
                 expect(result.current.loopDetectionConfirmationRequest).not.toBeNull();
             });
             // Simulate user selecting "keep" for first request
@@ -1919,7 +2000,7 @@ describe('useGeminiStream', () => {
             await act(async () => {
                 await result.current.submitQuery('second query');
             });
-            await waitFor(() => {
+            await vi.waitFor(() => {
                 expect(result.current.loopDetectionConfirmationRequest).not.toBeNull();
             });
             // Simulate user selecting "disable" for second request
@@ -1935,7 +2016,7 @@ describe('useGeminiStream', () => {
                 text: 'Loop detection has been disabled for this session. Retrying request...',
             }, expect.any(Number));
             // Verify that the request was retried
-            await waitFor(() => {
+            await vi.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));
             });
@@ -1955,14 +2036,14 @@ describe('useGeminiStream', () => {
                 await result.current.submitQuery('test query');
             });
             // Verify that the content was added to history before the loop detection dialog
-            await waitFor(() => {
+            await vi.waitFor(() => {
                 expect(mockAddItem).toHaveBeenCalledWith(expect.objectContaining({
                     type: 'gemini',
                     text: 'Some response content',
                 }), expect.any(Number));
             });
             // Then verify loop detection confirmation request was set
-            await waitFor(() => {
+            await vi.waitFor(() => {
                 expect(result.current.loopDetectionConfirmationRequest).not.toBeNull();
             });
         });