@machina.ai/cell-cli 1.11.0-rc1 → 1.13.0-rc1

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

@@ -4,12 +4,14 @@
  * 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 { waitFor } from '../../test-utils/async.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 +177,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 +200,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,
@@ -235,6 +265,49 @@ describe('useGeminiStream', () => {
             client,
         };
     };
+    // Helper to create mock tool calls - reduces boilerplate
+    const createMockToolCall = (toolName, callId, confirmationType, mockOnConfirm, status = 'awaiting_approval') => ({
+        request: {
+            callId,
+            name: toolName,
+            args: {},
+            isClientInitiated: false,
+            prompt_id: 'prompt-id-1',
+        },
+        status: status,
+        responseSubmittedToGemini: false,
+        confirmationDetails: confirmationType === 'edit'
+            ? {
+                type: 'edit',
+                title: 'Confirm Edit',
+                onConfirm: mockOnConfirm,
+                fileName: 'file.txt',
+                filePath: '/test/file.txt',
+                fileDiff: 'fake diff',
+                originalContent: 'old',
+                newContent: 'new',
+            }
+            : {
+                type: 'info',
+                title: `${toolName} confirmation`,
+                onConfirm: mockOnConfirm,
+                prompt: `Execute ${toolName}?`,
+            },
+        tool: {
+            name: toolName,
+            displayName: toolName,
+            description: `${toolName} description`,
+            build: vi.fn(),
+        },
+        invocation: {
+            getDescription: () => 'Mock description',
+        },
+    });
+    // Helper to render hook with default parameters - reduces boilerplate
+    const renderHookWithDefaults = (options = {}) => {
+        const { shellModeActive = false, onCancelSubmit = () => { }, setShellInputFocused = () => { }, performMemoryRefresh = () => Promise.resolve(), onAuthError = () => { }, onEditorClose = () => { }, setModelSwitched = vi.fn(), modelSwitched = false, } = options;
+        return renderHook(() => useGeminiStream(new MockedGeminiClientClass(mockConfig), [], mockAddItem, mockConfig, mockLoadedSettings, mockOnDebugMessage, mockHandleSlashCommand, shellModeActive, () => 'vscode', onAuthError, performMemoryRefresh, modelSwitched, setModelSwitched, onEditorClose, onCancelSubmit, setShellInputFocused, 80, 24));
+    };
     it('should not submit tool responses if not all tool calls are completed', () => {
         const toolCalls = [
             {
@@ -341,7 +414,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
@@ -390,7 +463,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
@@ -473,7 +546,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
@@ -552,6 +625,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 +639,7 @@ describe('useGeminiStream', () => {
                 completedToolCalls,
                 mockScheduleToolCalls,
                 mockMarkToolsAsSubmitted,
+                vi.fn(), // setToolCallsForDisplay
             ];
         });
         act(() => {
@@ -696,18 +771,18 @@ describe('useGeminiStream', () => {
             // 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');
             // 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 +802,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 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', () => {
@@ -871,7 +996,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
@@ -917,62 +1042,8 @@ describe('useGeminiStream', () => {
         it('should auto-approve all pending tool calls when switching to YOLO mode', async () => {
             const mockOnConfirm = vi.fn().mockResolvedValue(undefined);
             const awaitingApprovalToolCalls = [
-                {
-                    request: {
-                        callId: 'call1',
-                        name: 'replace',
-                        args: { old_string: 'old', new_string: 'new' },
-                        isClientInitiated: false,
-                        prompt_id: 'prompt-id-1',
-                    },
-                    status: 'awaiting_approval',
-                    responseSubmittedToGemini: false,
-                    confirmationDetails: {
-                        type: 'edit',
-                        title: 'Confirm Edit',
-                        onConfirm: mockOnConfirm,
-                        fileName: 'file.txt',
-                        filePath: '/test/file.txt',
-                        fileDiff: 'fake diff',
-                        originalContent: 'old',
-                        newContent: 'new',
-                    },
-                    tool: {
-                        name: 'replace',
-                        displayName: 'replace',
-                        description: 'Replace text',
-                        build: vi.fn(),
-                    },
-                    invocation: {
-                        getDescription: () => 'Mock description',
-                    },
-                },
-                {
-                    request: {
-                        callId: 'call2',
-                        name: 'read_file',
-                        args: { path: '/test/file.txt' },
-                        isClientInitiated: false,
-                        prompt_id: 'prompt-id-1',
-                    },
-                    status: 'awaiting_approval',
-                    responseSubmittedToGemini: false,
-                    confirmationDetails: {
-                        type: 'info',
-                        title: 'Read File',
-                        onConfirm: mockOnConfirm,
-                        prompt: 'Read /test/file.txt?',
-                    },
-                    tool: {
-                        name: 'read_file',
-                        displayName: 'read_file',
-                        description: 'Read file',
-                        build: vi.fn(),
-                    },
-                    invocation: {
-                        getDescription: () => 'Mock description',
-                    },
-                },
+                createMockToolCall('replace', 'call1', 'edit', mockOnConfirm),
+                createMockToolCall('read_file', 'call2', 'info', mockOnConfirm),
             ];
             const { result } = renderTestHook(awaitingApprovalToolCalls);
             await act(async () => {
@@ -980,109 +1051,23 @@ describe('useGeminiStream', () => {
             });
             // Both tool calls should be auto-approved
             expect(mockOnConfirm).toHaveBeenCalledTimes(2);
-            expect(mockOnConfirm).toHaveBeenNthCalledWith(1, ToolConfirmationOutcome.ProceedOnce);
-            expect(mockOnConfirm).toHaveBeenNthCalledWith(2, ToolConfirmationOutcome.ProceedOnce);
+            expect(mockOnConfirm).toHaveBeenCalledWith(ToolConfirmationOutcome.ProceedOnce);
         });
         it('should only auto-approve edit tools when switching to AUTO_EDIT mode', async () => {
             const mockOnConfirmReplace = vi.fn().mockResolvedValue(undefined);
             const mockOnConfirmWrite = vi.fn().mockResolvedValue(undefined);
             const mockOnConfirmRead = vi.fn().mockResolvedValue(undefined);
             const awaitingApprovalToolCalls = [
-                {
-                    request: {
-                        callId: 'call1',
-                        name: 'replace',
-                        args: { old_string: 'old', new_string: 'new' },
-                        isClientInitiated: false,
-                        prompt_id: 'prompt-id-1',
-                    },
-                    status: 'awaiting_approval',
-                    responseSubmittedToGemini: false,
-                    confirmationDetails: {
-                        type: 'edit',
-                        title: 'Confirm Edit',
-                        onConfirm: mockOnConfirmReplace,
-                        fileName: 'file.txt',
-                        filePath: '/test/file.txt',
-                        fileDiff: 'fake diff',
-                        originalContent: 'old',
-                        newContent: 'new',
-                    },
-                    tool: {
-                        name: 'replace',
-                        displayName: 'replace',
-                        description: 'Replace text',
-                        build: vi.fn(),
-                    },
-                    invocation: {
-                        getDescription: () => 'Mock description',
-                    },
-                },
-                {
-                    request: {
-                        callId: 'call2',
-                        name: 'write_file',
-                        args: { path: '/test/new.txt', content: 'content' },
-                        isClientInitiated: false,
-                        prompt_id: 'prompt-id-1',
-                    },
-                    status: 'awaiting_approval',
-                    responseSubmittedToGemini: false,
-                    confirmationDetails: {
-                        type: 'edit',
-                        title: 'Confirm Edit',
-                        onConfirm: mockOnConfirmWrite,
-                        fileName: 'new.txt',
-                        filePath: '/test/new.txt',
-                        fileDiff: 'fake diff',
-                        originalContent: null,
-                        newContent: 'content',
-                    },
-                    tool: {
-                        name: 'write_file',
-                        displayName: 'write_file',
-                        description: 'Write file',
-                        build: vi.fn(),
-                    },
-                    invocation: {
-                        getDescription: () => 'Mock description',
-                    },
-                },
-                {
-                    request: {
-                        callId: 'call3',
-                        name: 'read_file',
-                        args: { path: '/test/file.txt' },
-                        isClientInitiated: false,
-                        prompt_id: 'prompt-id-1',
-                    },
-                    status: 'awaiting_approval',
-                    responseSubmittedToGemini: false,
-                    confirmationDetails: {
-                        type: 'info',
-                        title: 'Read File',
-                        onConfirm: mockOnConfirmRead,
-                        prompt: 'Read /test/file.txt?',
-                    },
-                    tool: {
-                        name: 'read_file',
-                        displayName: 'read_file',
-                        description: 'Read file',
-                        build: vi.fn(),
-                    },
-                    invocation: {
-                        getDescription: () => 'Mock description',
-                    },
-                },
+                createMockToolCall('replace', 'call1', 'edit', mockOnConfirmReplace),
+                createMockToolCall('write_file', 'call2', 'edit', mockOnConfirmWrite),
+                createMockToolCall('read_file', 'call3', 'info', mockOnConfirmRead),
             ];
             const { result } = renderTestHook(awaitingApprovalToolCalls);
             await act(async () => {
                 await result.current.handleApprovalModeChange(ApprovalMode.AUTO_EDIT);
             });
             // Only replace and write_file should be auto-approved
-            expect(mockOnConfirmReplace).toHaveBeenCalledTimes(1);
             expect(mockOnConfirmReplace).toHaveBeenCalledWith(ToolConfirmationOutcome.ProceedOnce);
-            expect(mockOnConfirmWrite).toHaveBeenCalledTimes(1);
             expect(mockOnConfirmWrite).toHaveBeenCalledWith(ToolConfirmationOutcome.ProceedOnce);
             // read_file should not be auto-approved
             expect(mockOnConfirmRead).not.toHaveBeenCalled();
@@ -1090,36 +1075,7 @@ describe('useGeminiStream', () => {
         it('should not auto-approve any tools when switching to REQUIRE_CONFIRMATION mode', async () => {
             const mockOnConfirm = vi.fn().mockResolvedValue(undefined);
             const awaitingApprovalToolCalls = [
-                {
-                    request: {
-                        callId: 'call1',
-                        name: 'replace',
-                        args: { old_string: 'old', new_string: 'new' },
-                        isClientInitiated: false,
-                        prompt_id: 'prompt-id-1',
-                    },
-                    status: 'awaiting_approval',
-                    responseSubmittedToGemini: false,
-                    confirmationDetails: {
-                        type: 'edit',
-                        title: 'Confirm Edit',
-                        onConfirm: mockOnConfirm,
-                        fileName: 'file.txt',
-                        filePath: '/test/file.txt',
-                        fileDiff: 'fake diff',
-                        originalContent: 'old',
-                        newContent: 'new',
-                    },
-                    tool: {
-                        name: 'replace',
-                        displayName: 'replace',
-                        description: 'Replace text',
-                        build: vi.fn(),
-                    },
-                    invocation: {
-                        getDescription: () => 'Mock description',
-                    },
-                },
+                createMockToolCall('replace', 'call1', 'edit', mockOnConfirm),
             ];
             const { result } = renderTestHook(awaitingApprovalToolCalls);
             await act(async () => {
@@ -1129,85 +1085,27 @@ 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
                 .fn()
                 .mockRejectedValue(new Error('Approval failed'));
             const awaitingApprovalToolCalls = [
-                {
-                    request: {
-                        callId: 'call1',
-                        name: 'replace',
-                        args: { old_string: 'old', new_string: 'new' },
-                        isClientInitiated: false,
-                        prompt_id: 'prompt-id-1',
-                    },
-                    status: 'awaiting_approval',
-                    responseSubmittedToGemini: false,
-                    confirmationDetails: {
-                        type: 'edit',
-                        title: 'Confirm Edit',
-                        onConfirm: mockOnConfirmSuccess,
-                        fileName: 'file.txt',
-                        filePath: '/test/file.txt',
-                        fileDiff: 'fake diff',
-                        originalContent: 'old',
-                        newContent: 'new',
-                    },
-                    tool: {
-                        name: 'replace',
-                        displayName: 'replace',
-                        description: 'Replace text',
-                        build: vi.fn(),
-                    },
-                    invocation: {
-                        getDescription: () => 'Mock description',
-                    },
-                },
-                {
-                    request: {
-                        callId: 'call2',
-                        name: 'write_file',
-                        args: { path: '/test/file.txt', content: 'content' },
-                        isClientInitiated: false,
-                        prompt_id: 'prompt-id-1',
-                    },
-                    status: 'awaiting_approval',
-                    responseSubmittedToGemini: false,
-                    confirmationDetails: {
-                        type: 'edit',
-                        title: 'Confirm Edit',
-                        onConfirm: mockOnConfirmError,
-                        fileName: 'file.txt',
-                        filePath: '/test/file.txt',
-                        fileDiff: 'fake diff',
-                        originalContent: null,
-                        newContent: 'content',
-                    },
-                    tool: {
-                        name: 'write_file',
-                        displayName: 'write_file',
-                        description: 'Write file',
-                        build: vi.fn(),
-                    },
-                    invocation: {
-                        getDescription: () => 'Mock description',
-                    },
-                },
+                createMockToolCall('replace', 'call1', 'edit', mockOnConfirmSuccess),
+                createMockToolCall('write_file', 'call2', 'edit', mockOnConfirmError),
             ];
             const { result } = renderTestHook(awaitingApprovalToolCalls);
             await act(async () => {
                 await result.current.handleApprovalModeChange(ApprovalMode.YOLO);
             });
             // Both confirmation methods should be called
-            expect(mockOnConfirmSuccess).toHaveBeenCalledTimes(1);
-            expect(mockOnConfirmError).toHaveBeenCalledTimes(1);
+            expect(mockOnConfirmSuccess).toHaveBeenCalled();
+            expect(mockOnConfirmError).toHaveBeenCalled();
             // 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 = [
@@ -1374,53 +1272,37 @@ describe('useGeminiStream', () => {
             beforeEach(() => {
                 vi.mocked(tokenLimit).mockReturnValue(100);
             });
-            it('should add message without suggestion when remaining tokens are > 75% of limit', async () => {
-                // Setup mock to return a stream with ContextWindowWillOverflow event
-                // Limit is 100, remaining is 80 (> 75)
-                mockSendMessageStream.mockReturnValue((async function* () {
-                    yield {
-                        type: ServerGeminiEventType.ContextWindowWillOverflow,
-                        value: {
-                            estimatedRequestTokenCount: 20,
-                            remainingTokenCount: 80,
-                        },
-                    };
-                })());
-                const { result } = renderHook(() => useGeminiStream(new MockedGeminiClientClass(mockConfig), [], mockAddItem, mockConfig, mockLoadedSettings, mockOnDebugMessage, mockHandleSlashCommand, false, () => 'vscode', () => { }, () => Promise.resolve(), false, () => { }, () => { }, () => { }, () => { }, 80, 24));
-                // Submit a query
-                await act(async () => {
-                    await result.current.submitQuery('Test overflow');
-                });
-                // Check that the message was added without suggestion
-                await waitFor(() => {
-                    expect(mockAddItem).toHaveBeenCalledWith({
-                        type: 'info',
-                        text: `Sending this message (20 tokens) might exceed the remaining context window limit (80 tokens).`,
-                    }, expect.any(Number));
-                });
-            });
-            it('should add message with suggestion when remaining tokens are < 75% of limit', async () => {
-                // Setup mock to return a stream with ContextWindowWillOverflow event
-                // Limit is 100, remaining is 70 (< 75)
+            it.each([
+                {
+                    name: 'without suggestion when remaining tokens are > 75% of limit',
+                    requestTokens: 20,
+                    remainingTokens: 80,
+                    expectedMessage: 'Sending this message (20 tokens) might exceed the remaining context window limit (80 tokens).',
+                },
+                {
+                    name: 'with suggestion when remaining tokens are < 75% of limit',
+                    requestTokens: 30,
+                    remainingTokens: 70,
+                    expectedMessage: '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.',
+                },
+            ])('should add message $name', async ({ requestTokens, remainingTokens, expectedMessage }) => {
                 mockSendMessageStream.mockReturnValue((async function* () {
                     yield {
                         type: ServerGeminiEventType.ContextWindowWillOverflow,
                         value: {
-                            estimatedRequestTokenCount: 30,
-                            remainingTokenCount: 70,
+                            estimatedRequestTokenCount: requestTokens,
+                            remainingTokenCount: remainingTokens,
                         },
                     };
                 })());
-                const { result } = renderHook(() => useGeminiStream(new MockedGeminiClientClass(mockConfig), [], mockAddItem, mockConfig, mockLoadedSettings, mockOnDebugMessage, mockHandleSlashCommand, false, () => 'vscode', () => { }, () => Promise.resolve(), false, () => { }, () => { }, () => { }, () => { }, 80, 24));
-                // Submit a query
+                const { result } = renderHookWithDefaults();
                 await act(async () => {
                     await result.current.submitQuery('Test overflow');
                 });
-                // Check that the message was added with suggestion
                 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.`,
+                        text: expectedMessage,
                     }, expect.any(Number));
                 });
             });
@@ -1447,112 +1329,71 @@ describe('useGeminiStream', () => {
                 expect(onCancelSubmitSpy).toHaveBeenCalled();
             });
         });
-        it('should not add message for STOP finish reason', async () => {
-            // Setup mock to return a stream with STOP finish reason
-            mockSendMessageStream.mockReturnValue((async function* () {
-                yield {
-                    type: ServerGeminiEventType.Content,
-                    value: 'Complete response',
-                };
-                yield {
-                    type: ServerGeminiEventType.Finished,
-                    value: { reason: 'STOP', usageMetadata: undefined },
-                };
-            })());
-            const { result } = renderHook(() => useGeminiStream(new MockedGeminiClientClass(mockConfig), [], mockAddItem, mockConfig, mockLoadedSettings, mockOnDebugMessage, mockHandleSlashCommand, false, () => 'vscode', () => { }, () => Promise.resolve(), false, () => { }, () => { }, () => { }, () => { }, 80, 24));
-            // Submit a query
-            await act(async () => {
-                await result.current.submitQuery('Test normal completion');
-            });
-            // Wait a bit to ensure no message is added
-            await new Promise((resolve) => setTimeout(resolve, 100));
-            // Check that no info message was added for STOP
-            const infoMessages = mockAddItem.mock.calls.filter((call) => call[0].type === 'info');
-            expect(infoMessages).toHaveLength(0);
-        });
-        it('should not add message for FINISH_REASON_UNSPECIFIED', async () => {
-            // Setup mock to return a stream with FINISH_REASON_UNSPECIFIED
+        it.each([
+            {
+                reason: 'STOP',
+                shouldAddMessage: false,
+            },
+            {
+                reason: 'FINISH_REASON_UNSPECIFIED',
+                shouldAddMessage: false,
+            },
+            {
+                reason: 'SAFETY',
+                message: '⚠️  Response stopped due to safety reasons.',
+            },
+            {
+                reason: 'RECITATION',
+                message: '⚠️  Response stopped due to recitation policy.',
+            },
+            {
+                reason: 'LANGUAGE',
+                message: '⚠️  Response stopped due to unsupported language.',
+            },
+            {
+                reason: 'BLOCKLIST',
+                message: '⚠️  Response stopped due to forbidden terms.',
+            },
+            {
+                reason: 'PROHIBITED_CONTENT',
+                message: '⚠️  Response stopped due to prohibited content.',
+            },
+            {
+                reason: 'SPII',
+                message: '⚠️  Response stopped due to sensitive personally identifiable information.',
+            },
+            {
+                reason: 'OTHER',
+                message: '⚠️  Response stopped for other reasons.',
+            },
+            {
+                reason: 'MALFORMED_FUNCTION_CALL',
+                message: '⚠️  Response stopped due to malformed function call.',
+            },
+            {
+                reason: 'IMAGE_SAFETY',
+                message: '⚠️  Response stopped due to image safety violations.',
+            },
+            {
+                reason: 'UNEXPECTED_TOOL_CALL',
+                message: '⚠️  Response stopped due to unexpected tool call.',
+            },
+        ])('should handle $reason finish reason correctly', async ({ reason, shouldAddMessage = true, message }) => {
             mockSendMessageStream.mockReturnValue((async function* () {
                 yield {
                     type: ServerGeminiEventType.Content,
-                    value: 'Response with unspecified finish',
+                    value: `Response for ${reason}`,
                 };
                 yield {
                     type: ServerGeminiEventType.Finished,
-                    value: {
-                        reason: 'FINISH_REASON_UNSPECIFIED',
-                        usageMetadata: undefined,
-                    },
+                    value: { reason, usageMetadata: undefined },
                 };
             })());
-            const { result } = renderHook(() => useGeminiStream(new MockedGeminiClientClass(mockConfig), [], mockAddItem, mockConfig, mockLoadedSettings, mockOnDebugMessage, mockHandleSlashCommand, false, () => 'vscode', () => { }, () => Promise.resolve(), false, () => { }, () => { }, () => { }, () => { }, 80, 24));
-            // Submit a query
+            const { result } = renderHookWithDefaults();
             await act(async () => {
-                await result.current.submitQuery('Test unspecified finish');
+                await result.current.submitQuery(`Test ${reason}`);
             });
-            // Wait a bit to ensure no message is added
-            await new Promise((resolve) => setTimeout(resolve, 100));
-            // Check that no info message was added
-            const infoMessages = mockAddItem.mock.calls.filter((call) => call[0].type === 'info');
-            expect(infoMessages).toHaveLength(0);
-        });
-        it('should add appropriate messages for other finish reasons', async () => {
-            const testCases = [
-                {
-                    reason: 'SAFETY',
-                    message: '⚠️  Response stopped due to safety reasons.',
-                },
-                {
-                    reason: 'RECITATION',
-                    message: '⚠️  Response stopped due to recitation policy.',
-                },
-                {
-                    reason: 'LANGUAGE',
-                    message: '⚠️  Response stopped due to unsupported language.',
-                },
-                {
-                    reason: 'BLOCKLIST',
-                    message: '⚠️  Response stopped due to forbidden terms.',
-                },
-                {
-                    reason: 'PROHIBITED_CONTENT',
-                    message: '⚠️  Response stopped due to prohibited content.',
-                },
-                {
-                    reason: 'SPII',
-                    message: '⚠️  Response stopped due to sensitive personally identifiable information.',
-                },
-                { reason: 'OTHER', message: '⚠️  Response stopped for other reasons.' },
-                {
-                    reason: 'MALFORMED_FUNCTION_CALL',
-                    message: '⚠️  Response stopped due to malformed function call.',
-                },
-                {
-                    reason: 'IMAGE_SAFETY',
-                    message: '⚠️  Response stopped due to image safety violations.',
-                },
-                {
-                    reason: 'UNEXPECTED_TOOL_CALL',
-                    message: '⚠️  Response stopped due to unexpected tool call.',
-                },
-            ];
-            for (const { reason, message } of testCases) {
-                // Reset mocks for each test case
-                mockAddItem.mockClear();
-                mockSendMessageStream.mockReturnValue((async function* () {
-                    yield {
-                        type: ServerGeminiEventType.Content,
-                        value: `Response for ${reason}`,
-                    };
-                    yield {
-                        type: ServerGeminiEventType.Finished,
-                        value: { reason, usageMetadata: undefined },
-                    };
-                })());
-                const { result } = renderHook(() => useGeminiStream(new MockedGeminiClientClass(mockConfig), [], mockAddItem, mockConfig, mockLoadedSettings, mockOnDebugMessage, mockHandleSlashCommand, false, () => 'vscode', () => { }, () => Promise.resolve(), false, () => { }, () => { }, () => { }, vi.fn(), 80, 24));
-                await act(async () => {
-                    await result.current.submitQuery(`Test ${reason}`);
-                });
+            if (shouldAddMessage) {
                 await waitFor(() => {
                     expect(mockAddItem).toHaveBeenCalledWith({
                         type: 'info',
@@ -1560,6 +1401,14 @@ describe('useGeminiStream', () => {
                     }, expect.any(Number));
                 });
             }
+            else {
+                // Verify state returns to idle without any info messages
+                await waitFor(() => {
+                    expect(result.current.streamingState).toBe(StreamingState.Idle);
+                });
+                const infoMessages = mockAddItem.mock.calls.filter((call) => call[0].type === 'info');
+                expect(infoMessages).toHaveLength(0);
+            }
         });
     });
     it('should process @include commands, adding user turn after processing to prevent race conditions', async () => {