@google/gemini-cli 0.12.0-preview.4 → 0.13.0-nightly.20251031.c89bc30d

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

@@ -6,6 +6,7 @@
 import { describe, it, expect, vi, beforeEach } from 'vitest';
 import { act } from 'react';
 import { renderHook } from '../../test-utils/render.js';
+import { waitFor } from '../../test-utils/async.js';
 import { useGeminiStream } from './useGeminiStream.js';
 import { useKeypress } from './useKeypress.js';
 import * as atCommandProcessor from './atCommandProcessor.js';
@@ -379,7 +380,7 @@ describe('useGeminiStream', () => {
                 await capturedOnComplete(completedToolCalls);
             }
         });
-        await vi.waitFor(() => {
+        await waitFor(() => {
             expect(mockMarkToolsAsSubmitted).toHaveBeenCalledTimes(1);
             expect(mockSendMessageStream).toHaveBeenCalledTimes(1);
         });
@@ -428,7 +429,7 @@ describe('useGeminiStream', () => {
                 await capturedOnComplete(cancelledToolCalls);
             }
         });
-        await vi.waitFor(() => {
+        await waitFor(() => {
             expect(mockMarkToolsAsSubmitted).toHaveBeenCalledWith(['1']);
             expect(client.addHistory).toHaveBeenCalledWith({
                 role: 'user',
@@ -511,7 +512,7 @@ describe('useGeminiStream', () => {
                 await capturedOnComplete(allCancelledTools);
             }
         });
-        await vi.waitFor(() => {
+        await waitFor(() => {
             // The tools should be marked as submitted locally
             expect(mockMarkToolsAsSubmitted).toHaveBeenCalledWith([
                 'cancel-1',
@@ -611,7 +612,7 @@ describe('useGeminiStream', () => {
             }
         });
         // 5. Wait for submitQuery to be called
-        await vi.waitFor(() => {
+        await waitFor(() => {
             expect(mockSendMessageStream).toHaveBeenCalledWith(toolCallResponseParts, expect.any(AbortSignal), 'prompt-id-4');
         });
         // 6. After submission, the state should remain Responding until the stream completes.
@@ -649,13 +650,13 @@ describe('useGeminiStream', () => {
                 result.current.submitQuery('test query');
             });
             // Wait for the first part of the response
-            await vi.waitFor(() => {
+            await waitFor(() => {
                 expect(result.current.streamingState).toBe(StreamingState.Responding);
             });
             // Simulate escape key press
             simulateEscapeKeyPress();
             // Verify cancellation message is added
-            await vi.waitFor(() => {
+            await waitFor(() => {
                 expect(mockAddItem).toHaveBeenCalledWith({
                     type: MessageType.INFO,
                     text: 'Request cancelled.',
@@ -721,17 +722,17 @@ describe('useGeminiStream', () => {
             await act(async () => {
                 result.current.submitQuery('long running query');
             });
-            await vi.waitFor(() => {
+            await waitFor(() => {
                 expect(result.current.streamingState).toBe(StreamingState.Responding);
             });
             // Cancel the request
             simulateEscapeKeyPress();
             // Allow the stream to continue
-            act(() => {
+            await act(async () => {
                 continueStream();
+                // Wait a bit to see if the second part is processed
+                await new Promise((resolve) => setTimeout(resolve, 50));
             });
-            // Wait a bit to see if the second part is processed
-            await new Promise((resolve) => setTimeout(resolve, 50));
             // The text should not have been updated with " Canceled"
             const lastCall = mockAddItem.mock.calls.find((call) => call[0].type === 'gemini');
             expect(lastCall?.[0].text).toBe('Initial');
@@ -809,7 +810,7 @@ describe('useGeminiStream', () => {
             // The imperative cancel function should be called on the scheduler
             expect(mockCancelAllToolCalls).toHaveBeenCalled();
             // A cancellation message should be added to history
-            await vi.waitFor(() => {
+            await waitFor(() => {
                 expect(mockAddItem).toHaveBeenCalledWith(expect.objectContaining({
                     text: 'Request cancelled.',
                 }), expect.any(Number));
@@ -830,7 +831,7 @@ describe('useGeminiStream', () => {
             await act(async () => {
                 await result.current.submitQuery('/memory add "test fact"');
             });
-            await vi.waitFor(() => {
+            await waitFor(() => {
                 expect(mockScheduleToolCalls).toHaveBeenCalledWith([
                     expect.objectContaining({
                         name: 'save_memory',
@@ -850,7 +851,7 @@ describe('useGeminiStream', () => {
             await act(async () => {
                 await result.current.submitQuery('/help');
             });
-            await vi.waitFor(() => {
+            await waitFor(() => {
                 expect(mockHandleSlashCommand).toHaveBeenCalledWith('/help');
                 expect(mockScheduleToolCalls).not.toHaveBeenCalled();
                 expect(mockSendMessageStream).not.toHaveBeenCalled(); // No LLM call made
@@ -866,7 +867,7 @@ describe('useGeminiStream', () => {
             await act(async () => {
                 await result.current.submitQuery('/my-custom-command');
             });
-            await vi.waitFor(() => {
+            await 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));
@@ -883,7 +884,7 @@ describe('useGeminiStream', () => {
             await act(async () => {
                 await result.current.submitQuery('/emptycmd');
             });
-            await vi.waitFor(() => {
+            await waitFor(() => {
                 expect(mockHandleSlashCommand).toHaveBeenCalledWith('/emptycmd');
                 expect(localMockSendMessageStream).toHaveBeenCalledWith('', expect.any(AbortSignal), expect.any(String));
             });
@@ -893,7 +894,7 @@ describe('useGeminiStream', () => {
             await act(async () => {
                 await result.current.submitQuery('// This is a line comment');
             });
-            await vi.waitFor(() => {
+            await waitFor(() => {
                 expect(mockHandleSlashCommand).not.toHaveBeenCalled();
                 expect(localMockSendMessageStream).toHaveBeenCalledWith('// This is a line comment', expect.any(AbortSignal), expect.any(String));
             });
@@ -903,7 +904,7 @@ describe('useGeminiStream', () => {
             await act(async () => {
                 await result.current.submitQuery('/* This is a block comment */');
             });
-            await vi.waitFor(() => {
+            await waitFor(() => {
                 expect(mockHandleSlashCommand).not.toHaveBeenCalled();
                 expect(localMockSendMessageStream).toHaveBeenCalledWith('/* This is a block comment */', expect.any(AbortSignal), expect.any(String));
             });
@@ -913,7 +914,7 @@ describe('useGeminiStream', () => {
             await act(async () => {
                 await result.current.submitQuery('/about');
             });
-            await vi.waitFor(() => {
+            await waitFor(() => {
                 expect(mockHandleSlashCommand).not.toHaveBeenCalled();
             });
         });
@@ -961,7 +962,7 @@ describe('useGeminiStream', () => {
                     await capturedOnComplete([completedToolCall]);
                 }
             });
-            await vi.waitFor(() => {
+            await waitFor(() => {
                 expect(mockPerformMemoryRefresh).toHaveBeenCalledTimes(1);
             });
         });
@@ -989,7 +990,7 @@ describe('useGeminiStream', () => {
                 await result.current.submitQuery('test query');
             });
             // 3. Assertion
-            await vi.waitFor(() => {
+            await waitFor(() => {
                 expect(mockParseAndFormatApiError).toHaveBeenCalledWith('Rate limit exceeded', mockAuthType, undefined, 'gemini-2.5-pro', 'gemini-2.5-flash');
             });
         });
@@ -1444,7 +1445,7 @@ describe('useGeminiStream', () => {
                 await result.current.submitQuery('Generate long text');
             });
             // Check that the info message was added
-            await vi.waitFor(() => {
+            await waitFor(() => {
                 expect(mockAddItem).toHaveBeenCalledWith({
                     type: 'info',
                     text: '⚠️  Response truncated due to token limits.',
@@ -1473,7 +1474,7 @@ describe('useGeminiStream', () => {
                     await result.current.submitQuery('Test overflow');
                 });
                 // Check that the message was added without suggestion
-                await vi.waitFor(() => {
+                await waitFor(() => {
                     expect(mockAddItem).toHaveBeenCalledWith({
                         type: 'info',
                         text: `Sending this message (20 tokens) might exceed the remaining context window limit (80 tokens).`,
@@ -1498,7 +1499,7 @@ describe('useGeminiStream', () => {
                     await result.current.submitQuery('Test overflow');
                 });
                 // Check that the message was added with suggestion
-                await vi.waitFor(() => {
+                await 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.`,
@@ -1524,7 +1525,7 @@ describe('useGeminiStream', () => {
                 await result.current.submitQuery('Test overflow');
             });
             // Check that onCancelSubmit was called
-            await vi.waitFor(() => {
+            await waitFor(() => {
                 expect(onCancelSubmitSpy).toHaveBeenCalled();
             });
         });
@@ -1634,7 +1635,7 @@ describe('useGeminiStream', () => {
                 await act(async () => {
                     await result.current.submitQuery(`Test ${reason}`);
                 });
-                await vi.waitFor(() => {
+                await waitFor(() => {
                     expect(mockAddItem).toHaveBeenCalledWith({
                         type: 'info',
                         text: message,
@@ -1707,7 +1708,7 @@ describe('useGeminiStream', () => {
                 await result.current.submitQuery('First query');
             });
             // Wait for the first response to complete
-            await vi.waitFor(() => {
+            await waitFor(() => {
                 expect(mockAddItem).toHaveBeenCalledWith(expect.objectContaining({
                     type: 'gemini',
                     text: 'Some response content',
@@ -1732,7 +1733,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 vi.waitFor(() => {
+            await waitFor(() => {
                 expect(mockAddItem).toHaveBeenCalledWith(expect.objectContaining({
                     type: 'gemini',
                     text: 'New response content',
@@ -1791,7 +1792,7 @@ describe('useGeminiStream', () => {
                 await result.current.submitQuery('Test query');
             });
             // Verify cancellation message was added
-            await vi.waitFor(() => {
+            await waitFor(() => {
                 expect(mockAddItem).toHaveBeenCalledWith(expect.objectContaining({
                     type: 'info',
                     text: 'User cancelled the request.',
@@ -1818,7 +1819,7 @@ describe('useGeminiStream', () => {
                 await result.current.submitQuery('Test query');
             });
             // Verify error message was added
-            await vi.waitFor(() => {
+            await waitFor(() => {
                 expect(mockAddItem).toHaveBeenCalledWith(expect.objectContaining({
                     type: 'error',
                 }), expect.any(Number));
@@ -1852,7 +1853,7 @@ describe('useGeminiStream', () => {
             await act(async () => {
                 await result.current.submitQuery('test query');
             });
-            await vi.waitFor(() => {
+            await waitFor(() => {
                 expect(result.current.loopDetectionConfirmationRequest).not.toBeNull();
                 expect(typeof result.current.loopDetectionConfirmationRequest?.onComplete).toBe('function');
             });
@@ -1888,7 +1889,7 @@ describe('useGeminiStream', () => {
                 await result.current.submitQuery('test query');
             });
             // Wait for confirmation request to be set
-            await vi.waitFor(() => {
+            await waitFor(() => {
                 expect(result.current.loopDetectionConfirmationRequest).not.toBeNull();
             });
             // Simulate user selecting "disable"
@@ -1907,7 +1908,7 @@ describe('useGeminiStream', () => {
                 text: 'Loop detection has been disabled for this session. Retrying request...',
             }, expect.any(Number));
             // Verify that the request was retried
-            await vi.waitFor(() => {
+            await waitFor(() => {
                 expect(mockSendMessageStream).toHaveBeenCalledTimes(2);
                 expect(mockSendMessageStream).toHaveBeenNthCalledWith(2, 'test query', expect.any(AbortSignal), expect.any(String));
             });
@@ -1931,7 +1932,7 @@ describe('useGeminiStream', () => {
                 await result.current.submitQuery('test query');
             });
             // Wait for confirmation request to be set
-            await vi.waitFor(() => {
+            await waitFor(() => {
                 expect(result.current.loopDetectionConfirmationRequest).not.toBeNull();
             });
             // Simulate user selecting "keep"
@@ -1964,7 +1965,7 @@ describe('useGeminiStream', () => {
             await act(async () => {
                 await result.current.submitQuery('first query');
             });
-            await vi.waitFor(() => {
+            await waitFor(() => {
                 expect(result.current.loopDetectionConfirmationRequest).not.toBeNull();
             });
             // Simulate user selecting "keep" for first request
@@ -2000,7 +2001,7 @@ describe('useGeminiStream', () => {
             await act(async () => {
                 await result.current.submitQuery('second query');
             });
-            await vi.waitFor(() => {
+            await waitFor(() => {
                 expect(result.current.loopDetectionConfirmationRequest).not.toBeNull();
             });
             // Simulate user selecting "disable" for second request
@@ -2016,7 +2017,7 @@ describe('useGeminiStream', () => {
                 text: 'Loop detection has been disabled for this session. Retrying request...',
             }, expect.any(Number));
             // Verify that the request was retried
-            await vi.waitFor(() => {
+            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));
             });
@@ -2036,14 +2037,14 @@ describe('useGeminiStream', () => {
                 await result.current.submitQuery('test query');
             });
             // Verify that the content was added to history before the loop detection dialog
-            await vi.waitFor(() => {
+            await waitFor(() => {
                 expect(mockAddItem).toHaveBeenCalledWith(expect.objectContaining({
                     type: 'gemini',
                     text: 'Some response content',
                 }), expect.any(Number));
             });
             // Then verify loop detection confirmation request was set
-            await vi.waitFor(() => {
+            await waitFor(() => {
                 expect(result.current.loopDetectionConfirmationRequest).not.toBeNull();
             });
         });