qwen-code/packages/cli/src/ui/hooks/useGeminiStream.test.tsx

/**
 * @license
 * Copyright 2025 Google LLC
 * SPDX-License-Identifier: Apache-2.0
 */

/* eslint-disable @typescript-eslint/no-explicit-any */
import type { Mock, MockInstance } from 'vitest';
import { describe, it, expect, vi, beforeEach } from 'vitest';
import { renderHook, act, waitFor } from '@testing-library/react';
import { useGeminiStream } from './useGeminiStream.js';
import { useKeypress } from './useKeypress.js';
import * as atCommandProcessor from './atCommandProcessor.js';
import type {
  TrackedToolCall,
  TrackedCompletedToolCall,
  TrackedExecutingToolCall,
  TrackedCancelledToolCall,
  TrackedWaitingToolCall,
} from './useReactToolScheduler.js';
import { useReactToolScheduler } from './useReactToolScheduler.js';
import type {
  Config,
  EditorType,
  GeminiClient,
  AnyToolInvocation,
} from '@qwen-code/qwen-code-core';
import {
  ApprovalMode,
  AuthType,
  GeminiEventType as ServerGeminiEventType,
  ToolErrorType,
  ToolConfirmationOutcome,
} from '@qwen-code/qwen-code-core';
import type { Part, PartListUnion } from '@google/genai';
import type { UseHistoryManagerReturn } from './useHistoryManager.js';
import type { HistoryItem, SlashCommandProcessorResult } from '../types.js';
import { MessageType, StreamingState } from '../types.js';
import type { LoadedSettings } from '../../config/settings.js';

// --- MOCKS ---
const mockSendMessageStream = vi
  .fn()
  .mockReturnValue((async function* () {})());
const mockStartChat = vi.fn();

const MockedGeminiClientClass = vi.hoisted(() =>
  vi.fn().mockImplementation(function (this: any, _config: any) {
    // _config
    this.startChat = mockStartChat;
    this.sendMessageStream = mockSendMessageStream;
    this.addHistory = vi.fn();
    this.getChatRecordingService = vi.fn().mockReturnValue({
      recordThought: vi.fn(),
      initialize: vi.fn(),
      recordMessage: vi.fn(),
      recordMessageTokens: vi.fn(),
      recordToolCalls: vi.fn(),
      getConversationFile: vi.fn(),
    });
  }),
);

const MockedUserPromptEvent = vi.hoisted(() =>
  vi.fn().mockImplementation(() => {}),
);
const MockedApiCancelEvent = vi.hoisted(() =>
  vi.fn().mockImplementation(() => {}),
);
const mockParseAndFormatApiError = vi.hoisted(() => vi.fn());
const mockLogApiCancel = vi.hoisted(() => vi.fn());

// Vision auto-switch mocks (hoisted)
const mockHandleVisionSwitch = vi.hoisted(() =>
  vi.fn().mockResolvedValue({ shouldProceed: true }),
);
const mockRestoreOriginalModel = vi.hoisted(() =>
  vi.fn().mockResolvedValue(undefined),
);

vi.mock('@qwen-code/qwen-code-core', async (importOriginal) => {
  const actualCoreModule = (await importOriginal()) as any;
  return {
    ...actualCoreModule,
    GitService: vi.fn(),
    GeminiClient: MockedGeminiClientClass,
    UserPromptEvent: MockedUserPromptEvent,
    ApiCancelEvent: MockedApiCancelEvent,
    parseAndFormatApiError: mockParseAndFormatApiError,
    logApiCancel: mockLogApiCancel,
  };
});

const mockUseReactToolScheduler = useReactToolScheduler as Mock;
vi.mock('./useReactToolScheduler.js', async (importOriginal) => {
  const actualSchedulerModule = (await importOriginal()) as any;
  return {
    ...(actualSchedulerModule || {}),
    useReactToolScheduler: vi.fn(),
  };
});

vi.mock('./useVisionAutoSwitch.js', () => ({
  useVisionAutoSwitch: vi.fn(() => ({
    handleVisionSwitch: mockHandleVisionSwitch,
    restoreOriginalModel: mockRestoreOriginalModel,
  })),
}));

vi.mock('./useKeypress.js', () => ({
  useKeypress: vi.fn(),
}));

vi.mock('./shellCommandProcessor.js', () => ({
  useShellCommandProcessor: vi.fn().mockReturnValue({
    handleShellCommand: vi.fn(),
  }),
}));

vi.mock('./atCommandProcessor.js');

vi.mock('../utils/markdownUtilities.js', () => ({
  findLastSafeSplitPoint: vi.fn((s: string) => s.length),
}));

vi.mock('./useStateAndRef.js', () => ({
  useStateAndRef: vi.fn((initial) => {
    let val = initial;
    const ref = { current: val };
    const setVal = vi.fn((updater) => {
      if (typeof updater === 'function') {
        val = updater(val);
      } else {
        val = updater;
      }
      ref.current = val;
    });
    return [val, ref, setVal];
  }),
}));

vi.mock('./useLogger.js', () => ({
  useLogger: vi.fn().mockReturnValue({
    logMessage: vi.fn().mockResolvedValue(undefined),
  }),
}));

const mockStartNewPrompt = vi.fn();
const mockAddUsage = vi.fn();
vi.mock('../contexts/SessionContext.js', () => ({
  useSessionStats: vi.fn(() => ({
    startNewPrompt: mockStartNewPrompt,
    addUsage: mockAddUsage,
    getPromptCount: vi.fn(() => 5),
    stats: {
      sessionId: 'test-session-id',
    },
  })),
}));

vi.mock('./slashCommandProcessor.js', () => ({
  handleSlashCommand: vi.fn().mockReturnValue(false),
}));

// --- END MOCKS ---

// --- Tests for useGeminiStream Hook ---
describe('useGeminiStream', () => {
  let mockAddItem: Mock;
  let mockConfig: Config;
  let mockOnDebugMessage: Mock;
  let mockHandleSlashCommand: Mock;
  let mockScheduleToolCalls: Mock;
  let mockCancelAllToolCalls: Mock;
  let mockMarkToolsAsSubmitted: Mock;
  let handleAtCommandSpy: MockInstance;

  beforeEach(() => {
    vi.clearAllMocks(); // Clear mocks before each test

    mockAddItem = vi.fn();
    // Define the mock for getGeminiClient
    const mockGetGeminiClient = vi.fn().mockImplementation(() => {
      // MockedGeminiClientClass is defined in the module scope by the previous change.
      // It will use the mockStartChat and mockSendMessageStream that are managed within beforeEach.
      const clientInstance = new MockedGeminiClientClass(mockConfig);
      return clientInstance;
    });

    const contentGeneratorConfig = {
      model: 'test-model',
      apiKey: 'test-key',
      vertexai: false,
      authType: AuthType.USE_GEMINI,
    };

    mockConfig = {
      apiKey: 'test-api-key',
      model: 'gemini-pro',
      sandbox: false,
      targetDir: '/test/dir',
      debugMode: false,
      question: undefined,
      fullContext: false,
      coreTools: [],
      toolDiscoveryCommand: undefined,
      toolCallCommand: undefined,
      mcpServerCommand: undefined,
      mcpServers: undefined,
      userAgent: 'test-agent',
      userMemory: '',
      geminiMdFileCount: 0,
      alwaysSkipModificationConfirmation: false,
      vertexai: false,
      showMemoryUsage: false,
      contextFileName: undefined,
      getToolRegistry: vi.fn(
        () => ({ getToolSchemaList: vi.fn(() => []) }) as any,
      ),
      getProjectRoot: vi.fn(() => '/test/dir'),
      getCheckpointingEnabled: vi.fn(() => false),
      getGeminiClient: mockGetGeminiClient,
      getApprovalMode: () => ApprovalMode.DEFAULT,
      getUsageStatisticsEnabled: () => true,
      getDebugMode: () => false,
      addHistory: vi.fn(),
      getSessionId() {
        return 'test-session-id';
      },
      setQuotaErrorOccurred: vi.fn(),
      getQuotaErrorOccurred: vi.fn(() => false),
      getModel: vi.fn(() => 'gemini-2.5-pro'),
      getContentGeneratorConfig: vi
        .fn()
        .mockReturnValue(contentGeneratorConfig),
      getMaxSessionTurns: vi.fn(() => 50),
      getUseSmartEdit: () => false,
    } as unknown as Config;
    mockOnDebugMessage = vi.fn();
    mockHandleSlashCommand = vi.fn().mockResolvedValue(false);

    // Mock return value for useReactToolScheduler
    mockScheduleToolCalls = vi.fn();
    mockCancelAllToolCalls = vi.fn();
    mockMarkToolsAsSubmitted = vi.fn();

    // Default mock for useReactToolScheduler to prevent toolCalls being undefined initially
    mockUseReactToolScheduler.mockReturnValue([
      [], // Default to empty array for toolCalls
      mockScheduleToolCalls,
      mockCancelAllToolCalls,
      mockMarkToolsAsSubmitted,
    ]);

    // Reset mocks for GeminiClient instance methods (startChat and sendMessageStream)
    // The GeminiClient constructor itself is mocked at the module level.
    mockStartChat.mockClear().mockResolvedValue({
      sendMessageStream: mockSendMessageStream,
    } as unknown as any); // GeminiChat -> any
    mockSendMessageStream
      .mockClear()
      .mockReturnValue((async function* () {})());
    handleAtCommandSpy = vi.spyOn(atCommandProcessor, 'handleAtCommand');
  });

  const mockLoadedSettings: LoadedSettings = {
    merged: { preferredEditor: 'vscode' },
    user: { path: '/user/settings.json', settings: {} },
    workspace: { path: '/workspace/.qwen/settings.json', settings: {} },
    errors: [],
    forScope: vi.fn(),
    setValue: vi.fn(),
  } as unknown as LoadedSettings;

  const renderTestHook = (
    initialToolCalls: TrackedToolCall[] = [],
    geminiClient?: any,
  ) => {
    let currentToolCalls = initialToolCalls;
    const setToolCalls = (newToolCalls: TrackedToolCall[]) => {
      currentToolCalls = newToolCalls;
    };

    mockUseReactToolScheduler.mockImplementation(() => [
      currentToolCalls,
      mockScheduleToolCalls,
      mockCancelAllToolCalls,
      mockMarkToolsAsSubmitted,
    ]);

    const client = geminiClient || mockConfig.getGeminiClient();

    const { result, rerender } = renderHook(
      (props: {
        client: any;
        history: HistoryItem[];
        addItem: UseHistoryManagerReturn['addItem'];
        config: Config;
        onDebugMessage: (message: string) => void;
        handleSlashCommand: (
          cmd: PartListUnion,
        ) => Promise<SlashCommandProcessorResult | false>;
        shellModeActive: boolean;
        loadedSettings: LoadedSettings;
        toolCalls?: TrackedToolCall[]; // Allow passing updated toolCalls
      }) => {
        // Update the mock's return value if new toolCalls are passed in props
        if (props.toolCalls) {
          setToolCalls(props.toolCalls);
        }
        return useGeminiStream(
          props.client,
          props.history,
          props.addItem,
          props.config,
          props.loadedSettings,
          props.onDebugMessage,
          props.handleSlashCommand,
          props.shellModeActive,
          () => 'vscode' as EditorType,
          () => {},
          () => Promise.resolve(),
          false,
          () => {},
          () => {},
          () => {},
          false, // visionModelPreviewEnabled
          () => {},
          80,
          24,
        );
      },
      {
        initialProps: {
          client,
          history: [],
          addItem: mockAddItem as unknown as UseHistoryManagerReturn['addItem'],
          config: mockConfig,
          onDebugMessage: mockOnDebugMessage,
          handleSlashCommand: mockHandleSlashCommand as unknown as (
            cmd: PartListUnion,
          ) => Promise<SlashCommandProcessorResult | false>,
          shellModeActive: false,
          loadedSettings: mockLoadedSettings,
          toolCalls: initialToolCalls,
        },
      },
    );
    return {
      result,
      rerender,
      mockMarkToolsAsSubmitted,
      mockSendMessageStream,
      client,
    };
  };

  it('should not submit tool responses if not all tool calls are completed', () => {
    const toolCalls: TrackedToolCall[] = [
      {
        request: {
          callId: 'call1',
          name: 'tool1',
          args: {},
          isClientInitiated: false,
          prompt_id: 'prompt-id-1',
        },
        status: 'success',
        responseSubmittedToGemini: false,
        response: {
          callId: 'call1',
          responseParts: [{ text: 'tool 1 response' }],
          error: undefined,
          errorType: undefined, // FIX: Added missing property
          resultDisplay: 'Tool 1 success display',
        },
        tool: {
          name: 'tool1',
          displayName: 'tool1',
          description: 'desc1',
          build: vi.fn(),
        } as any,
        invocation: {
          getDescription: () => `Mock description`,
        } as unknown as AnyToolInvocation,
        startTime: Date.now(),
        endTime: Date.now(),
      } as TrackedCompletedToolCall,
      {
        request: {
          callId: 'call2',
          name: 'tool2',
          args: {},
          prompt_id: 'prompt-id-1',
        },
        status: 'executing',
        responseSubmittedToGemini: false,
        tool: {
          name: 'tool2',
          displayName: 'tool2',
          description: 'desc2',
          build: vi.fn(),
        } as any,
        invocation: {
          getDescription: () => `Mock description`,
        } as unknown as AnyToolInvocation,
        startTime: Date.now(),
        liveOutput: '...',
      } as TrackedExecutingToolCall,
    ];

    const { mockMarkToolsAsSubmitted, mockSendMessageStream } =
      renderTestHook(toolCalls);

    // Effect for submitting tool responses depends on toolCalls and isResponding
    // isResponding is initially false, so the effect should run.

    expect(mockMarkToolsAsSubmitted).not.toHaveBeenCalled();
    expect(mockSendMessageStream).not.toHaveBeenCalled(); // submitQuery uses this
  });

  it('should submit tool responses when all tool calls are completed and ready', async () => {
    const toolCall1ResponseParts: Part[] = [{ text: 'tool 1 final response' }];
    const toolCall2ResponseParts: Part[] = [{ text: 'tool 2 final response' }];
    const completedToolCalls: TrackedToolCall[] = [
      {
        request: {
          callId: 'call1',
          name: 'tool1',
          args: {},
          isClientInitiated: false,
          prompt_id: 'prompt-id-2',
        },
        status: 'success',
        responseSubmittedToGemini: false,
        response: {
          callId: 'call1',
          responseParts: toolCall1ResponseParts,
          errorType: undefined, // FIX: Added missing property
        },
        tool: {
          displayName: 'MockTool',
        },
        invocation: {
          getDescription: () => `Mock description`,
        } as unknown as AnyToolInvocation,
      } as TrackedCompletedToolCall,
      {
        request: {
          callId: 'call2',
          name: 'tool2',
          args: {},
          isClientInitiated: false,
          prompt_id: 'prompt-id-2',
        },
        status: 'error',
        responseSubmittedToGemini: false,
        response: {
          callId: 'call2',
          responseParts: toolCall2ResponseParts,
          errorType: ToolErrorType.UNHANDLED_EXCEPTION, // FIX: Added missing property
        },
      } as TrackedCompletedToolCall, // Treat error as a form of completion for submission
    ];

    // Capture the onComplete callback
    let capturedOnComplete:
      | ((completedTools: TrackedToolCall[]) => Promise<void>)
      | null = null;

    mockUseReactToolScheduler.mockImplementation((onComplete) => {
      capturedOnComplete = onComplete;
      return [[], mockScheduleToolCalls, mockMarkToolsAsSubmitted];
    });

    renderHook(() =>
      useGeminiStream(
        new MockedGeminiClientClass(mockConfig),
        [],
        mockAddItem,
        mockConfig,
        mockLoadedSettings,
        mockOnDebugMessage,
        mockHandleSlashCommand,
        false,
        () => 'vscode' as EditorType,
        () => {},
        () => Promise.resolve(),
        false,
        () => {},
        () => {},
        () => {},
        false, // visionModelPreviewEnabled
        () => {},
        80,
        24,
      ),
    );

    // Trigger the onComplete callback with completed tools
    await act(async () => {
      if (capturedOnComplete) {
        await capturedOnComplete(completedToolCalls);
      }
    });

    await waitFor(() => {
      expect(mockMarkToolsAsSubmitted).toHaveBeenCalledTimes(1);
      expect(mockSendMessageStream).toHaveBeenCalledTimes(1);
    });

    const expectedMergedResponse = [
      ...toolCall1ResponseParts,
      ...toolCall2ResponseParts,
    ];
    expect(mockSendMessageStream).toHaveBeenCalledWith(
      expectedMergedResponse,
      expect.any(AbortSignal),
      'prompt-id-2',
      { isContinuation: true },
    );
  });

  it('should handle all tool calls being cancelled', async () => {
    const cancelledToolCalls: TrackedToolCall[] = [
      {
        request: {
          callId: '1',
          name: 'testTool',
          args: {},
          isClientInitiated: false,
          prompt_id: 'prompt-id-3',
        },
        status: 'cancelled',
        response: {
          callId: '1',
          responseParts: [{ text: 'cancelled' }],
          errorType: undefined, // FIX: Added missing property
        },
        responseSubmittedToGemini: false,
        tool: {
          displayName: 'mock tool',
        },
        invocation: {
          getDescription: () => `Mock description`,
        } as unknown as AnyToolInvocation,
      } as TrackedCancelledToolCall,
    ];
    const client = new MockedGeminiClientClass(mockConfig);

    // Capture the onComplete callback
    let capturedOnComplete:
      | ((completedTools: TrackedToolCall[]) => Promise<void>)
      | null = null;

    mockUseReactToolScheduler.mockImplementation((onComplete) => {
      capturedOnComplete = onComplete;
      return [[], mockScheduleToolCalls, mockMarkToolsAsSubmitted];
    });

    renderHook(() =>
      useGeminiStream(
        client,
        [],
        mockAddItem,
        mockConfig,
        mockLoadedSettings,
        mockOnDebugMessage,
        mockHandleSlashCommand,
        false,
        () => 'vscode' as EditorType,
        () => {},
        () => Promise.resolve(),
        false,
        () => {},
        () => {},
        () => {},
        false, // visionModelPreviewEnabled
        () => {},
        80,
        24,
      ),
    );

    // Trigger the onComplete callback with cancelled tools
    await act(async () => {
      if (capturedOnComplete) {
        await capturedOnComplete(cancelledToolCalls);
      }
    });

    await waitFor(() => {
      expect(mockMarkToolsAsSubmitted).toHaveBeenCalledWith(['1']);
      expect(client.addHistory).toHaveBeenCalledWith({
        role: 'user',
        parts: [{ text: 'cancelled' }],
      });
      // Ensure we do NOT call back to the API
      expect(mockSendMessageStream).not.toHaveBeenCalled();
    });
  });

  it('should group multiple cancelled tool call responses into a single history entry', async () => {
    const cancelledToolCall1: TrackedCancelledToolCall = {
      request: {
        callId: 'cancel-1',
        name: 'toolA',
        args: {},
        isClientInitiated: false,
        prompt_id: 'prompt-id-7',
      },
      tool: {
        name: 'toolA',
        displayName: 'toolA',
        description: 'descA',
        build: vi.fn(),
      } as any,
      invocation: {
        getDescription: () => `Mock description`,
      } as unknown as AnyToolInvocation,
      status: 'cancelled',
      response: {
        callId: 'cancel-1',
        responseParts: [
          { functionResponse: { name: 'toolA', id: 'cancel-1' } },
        ],
        resultDisplay: undefined,
        error: undefined,
        errorType: undefined, // FIX: Added missing property
      },
      responseSubmittedToGemini: false,
    };
    const cancelledToolCall2: TrackedCancelledToolCall = {
      request: {
        callId: 'cancel-2',
        name: 'toolB',
        args: {},
        isClientInitiated: false,
        prompt_id: 'prompt-id-8',
      },
      tool: {
        name: 'toolB',
        displayName: 'toolB',
        description: 'descB',
        build: vi.fn(),
      } as any,
      invocation: {
        getDescription: () => `Mock description`,
      } as unknown as AnyToolInvocation,
      status: 'cancelled',
      response: {
        callId: 'cancel-2',
        responseParts: [
          { functionResponse: { name: 'toolB', id: 'cancel-2' } },
        ],
        resultDisplay: undefined,
        error: undefined,
        errorType: undefined, // FIX: Added missing property
      },
      responseSubmittedToGemini: false,
    };
    const allCancelledTools = [cancelledToolCall1, cancelledToolCall2];
    const client = new MockedGeminiClientClass(mockConfig);

    let capturedOnComplete:
      | ((completedTools: TrackedToolCall[]) => Promise<void>)
      | null = null;

    mockUseReactToolScheduler.mockImplementation((onComplete) => {
      capturedOnComplete = onComplete;
      return [[], mockScheduleToolCalls, mockMarkToolsAsSubmitted];
    });

    renderHook(() =>
      useGeminiStream(
        client,
        [],
        mockAddItem,
        mockConfig,
        mockLoadedSettings,
        mockOnDebugMessage,
        mockHandleSlashCommand,
        false,
        () => 'vscode' as EditorType,
        () => {},
        () => Promise.resolve(),
        false,
        () => {},
        () => {},
        () => {},
        false, // visionModelPreviewEnabled
        () => {},
        80,
        24,
      ),
    );

    // Trigger the onComplete callback with multiple cancelled tools
    await act(async () => {
      if (capturedOnComplete) {
        await capturedOnComplete(allCancelledTools);
      }
    });

    await waitFor(() => {
      // The tools should be marked as submitted locally
      expect(mockMarkToolsAsSubmitted).toHaveBeenCalledWith([
        'cancel-1',
        'cancel-2',
      ]);

      // Crucially, addHistory should be called only ONCE
      expect(client.addHistory).toHaveBeenCalledTimes(1);

      // And that single call should contain BOTH function responses
      expect(client.addHistory).toHaveBeenCalledWith({
        role: 'user',
        parts: [
          ...(cancelledToolCall1.response.responseParts as Part[]),
          ...(cancelledToolCall2.response.responseParts as Part[]),
        ],
      });

      // No message should be sent back to the API for a turn with only cancellations
      expect(mockSendMessageStream).not.toHaveBeenCalled();
    });
  });

  it('should not flicker streaming state to Idle between tool completion and submission', async () => {
    const toolCallResponseParts: PartListUnion = [
      { text: 'tool 1 final response' },
    ];

    const initialToolCalls: TrackedToolCall[] = [
      {
        request: {
          callId: 'call1',
          name: 'tool1',
          args: {},
          isClientInitiated: false,
          prompt_id: 'prompt-id-4',
        },
        status: 'executing',
        responseSubmittedToGemini: false,
        tool: {
          name: 'tool1',
          displayName: 'tool1',
          description: 'desc',
          build: vi.fn(),
        } as any,
        invocation: {
          getDescription: () => `Mock description`,
        } as unknown as AnyToolInvocation,
        startTime: Date.now(),
      } as TrackedExecutingToolCall,
    ];

    const completedToolCalls: TrackedToolCall[] = [
      {
        ...(initialToolCalls[0] as TrackedExecutingToolCall),
        status: 'success',
        response: {
          callId: 'call1',
          responseParts: toolCallResponseParts,
          error: undefined,
          errorType: undefined, // FIX: Added missing property
          resultDisplay: 'Tool 1 success display',
        },
        endTime: Date.now(),
      } as TrackedCompletedToolCall,
    ];

    // Capture the onComplete callback
    let capturedOnComplete:
      | ((completedTools: TrackedToolCall[]) => Promise<void>)
      | null = null;
    let currentToolCalls = initialToolCalls;

    mockUseReactToolScheduler.mockImplementation((onComplete) => {
      capturedOnComplete = onComplete;
      return [
        currentToolCalls,
        mockScheduleToolCalls,
        mockMarkToolsAsSubmitted,
      ];
    });

    const { result, rerender } = renderHook(() =>
      useGeminiStream(
        new MockedGeminiClientClass(mockConfig),
        [],
        mockAddItem,
        mockConfig,
        mockLoadedSettings,
        mockOnDebugMessage,
        mockHandleSlashCommand,
        false,
        () => 'vscode' as EditorType,
        () => {},
        () => Promise.resolve(),
        false,
        () => {},
        () => {},
        () => {},
        false, // visionModelPreviewEnabled
        () => {},
        80,
        24,
      ),
    );

    // 1. Initial state should be Responding because a tool is executing.
    expect(result.current.streamingState).toBe(StreamingState.Responding);

    // 2. Update the tool calls to completed state and rerender
    currentToolCalls = completedToolCalls;
    mockUseReactToolScheduler.mockImplementation((onComplete) => {
      capturedOnComplete = onComplete;
      return [
        completedToolCalls,
        mockScheduleToolCalls,
        mockMarkToolsAsSubmitted,
      ];
    });

    act(() => {
      rerender();
    });

    // 3. The state should *still* be Responding, not Idle.
    // This is because the completed tool's response has not been submitted yet.
    expect(result.current.streamingState).toBe(StreamingState.Responding);

    // 4. Trigger the onComplete callback to simulate tool completion
    await act(async () => {
      if (capturedOnComplete) {
        await capturedOnComplete(completedToolCalls);
      }
    });

    // 5. Wait for submitQuery to be called
    await waitFor(() => {
      expect(mockSendMessageStream).toHaveBeenCalledWith(
        toolCallResponseParts,
        expect.any(AbortSignal),
        'prompt-id-4',
        { isContinuation: true },
      );
    });

    // 6. After submission, the state should remain Responding until the stream completes.
    expect(result.current.streamingState).toBe(StreamingState.Responding);
  });

  describe('User Cancellation', () => {
    let keypressCallback: (key: any) => void;
    const mockUseKeypress = useKeypress as Mock;

    beforeEach(() => {
      // Capture the callback passed to useKeypress
      mockUseKeypress.mockImplementation((callback, options) => {
        if (options.isActive) {
          keypressCallback = callback;
        } else {
          keypressCallback = () => {};
        }
      });
    });

    const simulateEscapeKeyPress = () => {
      act(() => {
        keypressCallback({ name: 'escape' });
      });
    };

    it('should cancel an in-progress stream when escape is pressed', async () => {
      const mockStream = (async function* () {
        yield { type: 'content', value: 'Part 1' };
        // Keep the stream open
        await new Promise(() => {});
      })();
      mockSendMessageStream.mockReturnValue(mockStream);

      const { result } = renderTestHook();

      // Start a query
      await act(async () => {
        result.current.submitQuery('test query');
      });

      // Wait for the first part of the response
      await waitFor(() => {
        expect(result.current.streamingState).toBe(StreamingState.Responding);
      });

      // Simulate escape key press
      simulateEscapeKeyPress();

      // Verify cancellation message is added
      await waitFor(() => {
        expect(mockAddItem).toHaveBeenCalledWith(
          {
            type: MessageType.INFO,
            text: 'Request cancelled.',
          },
          expect.any(Number),
        );
      });

      // Verify state is reset
      expect(result.current.streamingState).toBe(StreamingState.Idle);
    });

    it('should call onCancelSubmit handler when escape is pressed', async () => {
      const cancelSubmitSpy = vi.fn();
      const mockStream = (async function* () {
        yield { type: 'content', value: 'Part 1' };
        // Keep the stream open
        await new Promise(() => {});
      })();
      mockSendMessageStream.mockReturnValue(mockStream);

      const { result } = renderHook(() =>
        useGeminiStream(
          mockConfig.getGeminiClient(),
          [],
          mockAddItem,
          mockConfig,
          mockLoadedSettings,
          mockOnDebugMessage,
          mockHandleSlashCommand,
          false,
          () => 'vscode' as EditorType,
          () => {},
          () => Promise.resolve(),
          false,
          () => {},
          () => {},
          cancelSubmitSpy,
          false, // visionModelPreviewEnabled
          () => {},
          80,
          24,
        ),
      );

      // Start a query
      await act(async () => {
        result.current.submitQuery('test query');
      });

      simulateEscapeKeyPress();

      expect(cancelSubmitSpy).toHaveBeenCalled();
    });

    it('should call setShellInputFocused(false) when escape is pressed', async () => {
      const setShellInputFocusedSpy = vi.fn();
      const mockStream = (async function* () {
        yield { type: 'content', value: 'Part 1' };
        await new Promise(() => {}); // Keep stream open
      })();
      mockSendMessageStream.mockReturnValue(mockStream);

      const { result } = renderHook(() =>
        useGeminiStream(
          mockConfig.getGeminiClient(),
          [],
          mockAddItem,
          mockConfig,
          mockLoadedSettings,
          mockOnDebugMessage,
          mockHandleSlashCommand,
          false,
          () => 'vscode' as EditorType,
          () => {},
          () => Promise.resolve(),
          false,
          () => {},
          () => {},
          vi.fn(),
          false,
          setShellInputFocusedSpy, // Pass the spy here
          80,
          24,
        ),
      );

      // Start a query
      await act(async () => {
        result.current.submitQuery('test query');
      });

      simulateEscapeKeyPress();

      expect(setShellInputFocusedSpy).toHaveBeenCalledWith(false);
    });

    it('should not do anything if escape is pressed when not responding', () => {
      const { result } = renderTestHook();

      expect(result.current.streamingState).toBe(StreamingState.Idle);

      // Simulate escape key press
      simulateEscapeKeyPress();

      // No change should happen, no cancellation message
      expect(mockAddItem).not.toHaveBeenCalledWith(
        expect.objectContaining({
          text: 'Request cancelled.',
        }),
        expect.any(Number),
      );
    });

    it('should prevent further processing after cancellation', async () => {
      let continueStream: () => void;
      const streamPromise = new Promise<void>((resolve) => {
        continueStream = resolve;
      });

      const mockStream = (async function* () {
        yield { type: 'content', value: 'Initial' };
        await streamPromise; // Wait until we manually continue
        yield { type: 'content', value: ' Canceled' };
      })();
      mockSendMessageStream.mockReturnValue(mockStream);

      const { result } = renderTestHook();

      await act(async () => {
        result.current.submitQuery('long running query');
      });

      await waitFor(() => {
        expect(result.current.streamingState).toBe(StreamingState.Responding);
      });

      // Cancel the request
      simulateEscapeKeyPress();

      // Allow the stream to continue
      act(() => {
        continueStream();
      });

      // 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 () => {
      const toolCalls: TrackedToolCall[] = [
        {
          request: { callId: 'call1', name: 'tool1', args: {} },
          status: 'executing',
          responseSubmittedToGemini: false,
          tool: {
            name: 'tool1',
            description: 'desc1',
            build: vi.fn().mockImplementation((_) => ({
              getDescription: () => `Mock description`,
            })),
          } as any,
          invocation: {
            getDescription: () => `Mock description`,
          },
          startTime: Date.now(),
          liveOutput: '...',
        } as TrackedExecutingToolCall,
      ];

      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();
    });
  });

  describe('Slash Command Handling', () => {
    it('should schedule a tool call when the command processor returns a schedule_tool action', async () => {
      const clientToolRequest: SlashCommandProcessorResult = {
        type: 'schedule_tool',
        toolName: 'save_memory',
        toolArgs: { fact: 'test fact' },
      };
      mockHandleSlashCommand.mockResolvedValue(clientToolRequest);

      const { result } = renderTestHook();

      await act(async () => {
        await result.current.submitQuery('/memory add "test fact"');
      });

      await waitFor(() => {
        expect(mockScheduleToolCalls).toHaveBeenCalledWith(
          [
            expect.objectContaining({
              name: 'save_memory',
              args: { fact: 'test fact' },
              isClientInitiated: true,
            }),
          ],
          expect.any(AbortSignal),
        );
        expect(mockSendMessageStream).not.toHaveBeenCalled();
      });
    });

    it('should stop processing and not call Gemini when a command is handled without a tool call', async () => {
      const uiOnlyCommandResult: SlashCommandProcessorResult = {
        type: 'handled',
      };
      mockHandleSlashCommand.mockResolvedValue(uiOnlyCommandResult);

      const { result } = renderTestHook();

      await act(async () => {
        await result.current.submitQuery('/help');
      });

      await waitFor(() => {
        expect(mockHandleSlashCommand).toHaveBeenCalledWith('/help');
        expect(mockScheduleToolCalls).not.toHaveBeenCalled();
        expect(mockSendMessageStream).not.toHaveBeenCalled(); // No LLM call made
      });
    });

    it('should call Gemini with prompt content when slash command returns a `submit_prompt` action', async () => {
      const customCommandResult: SlashCommandProcessorResult = {
        type: 'submit_prompt',
        content: 'This is the actual prompt from the command file.',
      };
      mockHandleSlashCommand.mockResolvedValue(customCommandResult);

      const { result, mockSendMessageStream: localMockSendMessageStream } =
        renderTestHook();

      await act(async () => {
        await result.current.submitQuery('/my-custom-command');
      });

      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),
          undefined,
        );

        expect(mockScheduleToolCalls).not.toHaveBeenCalled();
      });
    });

    it('should correctly handle a submit_prompt action with empty content', async () => {
      const emptyPromptResult: SlashCommandProcessorResult = {
        type: 'submit_prompt',
        content: '',
      };
      mockHandleSlashCommand.mockResolvedValue(emptyPromptResult);

      const { result, mockSendMessageStream: localMockSendMessageStream } =
        renderTestHook();

      await act(async () => {
        await result.current.submitQuery('/emptycmd');
      });

      await waitFor(() => {
        expect(mockHandleSlashCommand).toHaveBeenCalledWith('/emptycmd');
        expect(localMockSendMessageStream).toHaveBeenCalledWith(
          '',
          expect.any(AbortSignal),
          expect.any(String),
          undefined,
        );
      });
    });

    it('should not call handleSlashCommand for line comments', async () => {
      const { result, mockSendMessageStream: localMockSendMessageStream } =
        renderTestHook();

      await act(async () => {
        await result.current.submitQuery('// This is a line comment');
      });

      await waitFor(() => {
        expect(mockHandleSlashCommand).not.toHaveBeenCalled();
        expect(localMockSendMessageStream).toHaveBeenCalledWith(
          '// This is a line comment',
          expect.any(AbortSignal),
          expect.any(String),
          undefined,
        );
      });
    });

    it('should not call handleSlashCommand for block comments', async () => {
      const { result, mockSendMessageStream: localMockSendMessageStream } =
        renderTestHook();

      await act(async () => {
        await result.current.submitQuery('/* This is a block comment */');
      });

      await waitFor(() => {
        expect(mockHandleSlashCommand).not.toHaveBeenCalled();
        expect(localMockSendMessageStream).toHaveBeenCalledWith(
          '/* This is a block comment */',
          expect.any(AbortSignal),
          expect.any(String),
          undefined,
        );
      });
    });
  });

  describe('Memory Refresh on save_memory', () => {
    it('should call performMemoryRefresh when a save_memory tool call completes successfully', async () => {
      const mockPerformMemoryRefresh = vi.fn();
      const completedToolCall: TrackedCompletedToolCall = {
        request: {
          callId: 'save-mem-call-1',
          name: 'save_memory',
          args: { fact: 'test' },
          isClientInitiated: true,
          prompt_id: 'prompt-id-6',
        },
        status: 'success',
        responseSubmittedToGemini: false,
        response: {
          callId: 'save-mem-call-1',
          responseParts: [{ text: 'Memory saved' }],
          resultDisplay: 'Success: Memory saved',
          error: undefined,
          errorType: undefined, // FIX: Added missing property
        },
        tool: {
          name: 'save_memory',
          displayName: 'save_memory',
          description: 'Saves memory',
          build: vi.fn(),
        } as any,
        invocation: {
          getDescription: () => `Mock description`,
        } as unknown as AnyToolInvocation,
      };

      // Capture the onComplete callback
      let capturedOnComplete:
        | ((completedTools: TrackedToolCall[]) => Promise<void>)
        | null = null;

      mockUseReactToolScheduler.mockImplementation((onComplete) => {
        capturedOnComplete = onComplete;
        return [[], mockScheduleToolCalls, mockMarkToolsAsSubmitted];
      });

      renderHook(() =>
        useGeminiStream(
          new MockedGeminiClientClass(mockConfig),
          [],
          mockAddItem,
          mockConfig,
          mockLoadedSettings,
          mockOnDebugMessage,
          mockHandleSlashCommand,
          false,
          () => 'vscode' as EditorType,
          () => {},
          mockPerformMemoryRefresh,
          false,
          () => {},
          () => {},
          () => {},
          false, // visionModelPreviewEnabled
          () => {},
          80,
          24,
        ),
      );

      // Trigger the onComplete callback with the completed save_memory tool
      await act(async () => {
        if (capturedOnComplete) {
          await capturedOnComplete([completedToolCall]);
        }
      });

      await waitFor(() => {
        expect(mockPerformMemoryRefresh).toHaveBeenCalledTimes(1);
      });
    });
  });

  describe('Error Handling', () => {
    it('should call parseAndFormatApiError with the correct authType on stream initialization failure', async () => {
      // 1. Setup
      const mockError = new Error('Rate limit exceeded');
      const mockAuthType = AuthType.USE_VERTEX_AI;
      mockParseAndFormatApiError.mockClear();
      mockSendMessageStream.mockReturnValue(
        (async function* () {
          yield { type: 'content', value: '' };
          throw mockError;
        })(),
      );

      const testConfig = {
        ...mockConfig,
        getContentGeneratorConfig: vi.fn(() => ({
          authType: mockAuthType,
        })),
        getModel: vi.fn(() => 'gemini-2.5-pro'),
      } as unknown as Config;

      const { result } = renderHook(() =>
        useGeminiStream(
          new MockedGeminiClientClass(testConfig),
          [],
          mockAddItem,
          testConfig,
          mockLoadedSettings,
          mockOnDebugMessage,
          mockHandleSlashCommand,
          false,
          () => 'vscode' as EditorType,
          () => {},
          () => Promise.resolve(),
          false,
          () => {},
          () => {},
          () => {},
          false, // visionModelPreviewEnabled
          () => {},
          80,
          24,
        ),
      );

      // 2. Action
      await act(async () => {
        await result.current.submitQuery('test query');
      });

      // 3. Assertion
      await waitFor(() => {
        expect(mockParseAndFormatApiError).toHaveBeenCalledWith(
          'Rate limit exceeded',
          mockAuthType,
        );
      });
    });
  });

  describe('handleApprovalModeChange', () => {
    it('should auto-approve all pending tool calls when switching to YOLO mode', async () => {
      const mockOnConfirm = vi.fn().mockResolvedValue(undefined);
      const awaitingApprovalToolCalls: TrackedToolCall[] = [
        {
          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: {
            onConfirm: mockOnConfirm,
            onCancel: vi.fn(),
            message: 'Replace text?',
            displayedText: 'Replace old with new',
          },
          tool: {
            name: 'replace',
            displayName: 'replace',
            description: 'Replace text',
            build: vi.fn(),
          } as any,
          invocation: {
            getDescription: () => 'Mock description',
          } as unknown as AnyToolInvocation,
        } as unknown as TrackedWaitingToolCall,
        {
          request: {
            callId: 'call2',
            name: 'read_file',
            args: { path: '/test/file.txt' },
            isClientInitiated: false,
            prompt_id: 'prompt-id-1',
          },
          status: 'awaiting_approval',
          responseSubmittedToGemini: false,
          confirmationDetails: {
            onConfirm: mockOnConfirm,
            onCancel: vi.fn(),
            message: 'Read file?',
            displayedText: 'Read /test/file.txt',
          },
          tool: {
            name: 'read_file',
            displayName: 'read_file',
            description: 'Read file',
            build: vi.fn(),
          } as any,
          invocation: {
            getDescription: () => 'Mock description',
          } as unknown as AnyToolInvocation,
        } as unknown as TrackedWaitingToolCall,
      ];

      const { result } = renderTestHook(awaitingApprovalToolCalls);

      await act(async () => {
        await result.current.handleApprovalModeChange(ApprovalMode.YOLO);
      });

      // Both tool calls should be auto-approved
      expect(mockOnConfirm).toHaveBeenCalledTimes(2);
      expect(mockOnConfirm).toHaveBeenNthCalledWith(
        1,
        ToolConfirmationOutcome.ProceedOnce,
      );
      expect(mockOnConfirm).toHaveBeenNthCalledWith(
        2,
        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: TrackedToolCall[] = [
        {
          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: {
            onConfirm: mockOnConfirmReplace,
            onCancel: vi.fn(),
            message: 'Replace text?',
            displayedText: 'Replace old with new',
          },
          tool: {
            name: 'replace',
            displayName: 'replace',
            description: 'Replace text',
            build: vi.fn(),
          } as any,
          invocation: {
            getDescription: () => 'Mock description',
          } as unknown as AnyToolInvocation,
        } as unknown as TrackedWaitingToolCall,
        {
          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: {
            onConfirm: mockOnConfirmWrite,
            onCancel: vi.fn(),
            message: 'Write file?',
            displayedText: 'Write to /test/new.txt',
          },
          tool: {
            name: 'write_file',
            displayName: 'write_file',
            description: 'Write file',
            build: vi.fn(),
          } as any,
          invocation: {
            getDescription: () => 'Mock description',
          } as unknown as AnyToolInvocation,
        } as unknown as TrackedWaitingToolCall,
        {
          request: {
            callId: 'call3',
            name: 'read_file',
            args: { path: '/test/file.txt' },
            isClientInitiated: false,
            prompt_id: 'prompt-id-1',
          },
          status: 'awaiting_approval',
          responseSubmittedToGemini: false,
          confirmationDetails: {
            onConfirm: mockOnConfirmRead,
            onCancel: vi.fn(),
            message: 'Read file?',
            displayedText: 'Read /test/file.txt',
          },
          tool: {
            name: 'read_file',
            displayName: 'read_file',
            description: 'Read file',
            build: vi.fn(),
          } as any,
          invocation: {
            getDescription: () => 'Mock description',
          } as unknown as AnyToolInvocation,
        } as unknown as TrackedWaitingToolCall,
      ];

      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();
    });

    it('should not auto-approve any tools when switching to REQUIRE_CONFIRMATION mode', async () => {
      const mockOnConfirm = vi.fn().mockResolvedValue(undefined);
      const awaitingApprovalToolCalls: TrackedToolCall[] = [
        {
          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: {
            onConfirm: mockOnConfirm,
            onCancel: vi.fn(),
            message: 'Replace text?',
            displayedText: 'Replace old with new',
          },
          tool: {
            name: 'replace',
            displayName: 'replace',
            description: 'Replace text',
            build: vi.fn(),
          } as any,
          invocation: {
            getDescription: () => 'Mock description',
          } as unknown as AnyToolInvocation,
        } as unknown as TrackedWaitingToolCall,
      ];

      const { result } = renderTestHook(awaitingApprovalToolCalls);

      await act(async () => {
        await result.current.handleApprovalModeChange(ApprovalMode.DEFAULT);
      });

      // No tools should be auto-approved
      expect(mockOnConfirm).not.toHaveBeenCalled();
    });

    it('should handle errors gracefully when auto-approving tool calls', async () => {
      const consoleSpy = vi
        .spyOn(console, 'error')
        .mockImplementation(() => {});
      const mockOnConfirmSuccess = vi.fn().mockResolvedValue(undefined);
      const mockOnConfirmError = vi
        .fn()
        .mockRejectedValue(new Error('Approval failed'));

      const awaitingApprovalToolCalls: TrackedToolCall[] = [
        {
          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: {
            onConfirm: mockOnConfirmSuccess,
            onCancel: vi.fn(),
            message: 'Replace text?',
            displayedText: 'Replace old with new',
          },
          tool: {
            name: 'replace',
            displayName: 'replace',
            description: 'Replace text',
            build: vi.fn(),
          } as any,
          invocation: {
            getDescription: () => 'Mock description',
          } as unknown as AnyToolInvocation,
        } as unknown as TrackedWaitingToolCall,
        {
          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: {
            onConfirm: mockOnConfirmError,
            onCancel: vi.fn(),
            message: 'Write file?',
            displayedText: 'Write to /test/file.txt',
          },
          tool: {
            name: 'write_file',
            displayName: 'write_file',
            description: 'Write file',
            build: vi.fn(),
          } as any,
          invocation: {
            getDescription: () => 'Mock description',
          } as unknown as AnyToolInvocation,
        } as unknown as TrackedWaitingToolCall,
      ];

      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);

      // Error should be logged
      expect(consoleSpy).toHaveBeenCalledWith(
        'Failed to auto-approve tool call call2:',
        expect.any(Error),
      );

      consoleSpy.mockRestore();
    });

    it('should skip tool calls without confirmationDetails', async () => {
      const awaitingApprovalToolCalls: TrackedToolCall[] = [
        {
          request: {
            callId: 'call1',
            name: 'replace',
            args: { old_string: 'old', new_string: 'new' },
            isClientInitiated: false,
            prompt_id: 'prompt-id-1',
          },
          status: 'awaiting_approval',
          responseSubmittedToGemini: false,
          // No confirmationDetails
          tool: {
            name: 'replace',
            displayName: 'replace',
            description: 'Replace text',
            build: vi.fn(),
          } as any,
          invocation: {
            getDescription: () => 'Mock description',
          } as unknown as AnyToolInvocation,
        } as unknown as TrackedWaitingToolCall,
      ];

      const { result } = renderTestHook(awaitingApprovalToolCalls);

      // Should not throw an error
      await act(async () => {
        await result.current.handleApprovalModeChange(ApprovalMode.YOLO);
      });
    });

    it('should skip tool calls without onConfirm method in confirmationDetails', async () => {
      const awaitingApprovalToolCalls: TrackedToolCall[] = [
        {
          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: {
            onCancel: vi.fn(),
            message: 'Replace text?',
            displayedText: 'Replace old with new',
            // No onConfirm method
          } as any,
          tool: {
            name: 'replace',
            displayName: 'replace',
            description: 'Replace text',
            build: vi.fn(),
          } as any,
          invocation: {
            getDescription: () => 'Mock description',
          } as unknown as AnyToolInvocation,
        } as unknown as TrackedWaitingToolCall,
      ];

      const { result } = renderTestHook(awaitingApprovalToolCalls);

      // Should not throw an error
      await act(async () => {
        await result.current.handleApprovalModeChange(ApprovalMode.YOLO);
      });
    });

    it('should only process tool calls with awaiting_approval status', async () => {
      const mockOnConfirmAwaiting = vi.fn().mockResolvedValue(undefined);
      const mockOnConfirmExecuting = vi.fn().mockResolvedValue(undefined);

      const mixedStatusToolCalls: TrackedToolCall[] = [
        {
          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: {
            onConfirm: mockOnConfirmAwaiting,
            onCancel: vi.fn(),
            message: 'Replace text?',
            displayedText: 'Replace old with new',
          },
          tool: {
            name: 'replace',
            displayName: 'replace',
            description: 'Replace text',
            build: vi.fn(),
          } as any,
          invocation: {
            getDescription: () => 'Mock description',
          } as unknown as AnyToolInvocation,
        } as unknown as TrackedWaitingToolCall,
        {
          request: {
            callId: 'call2',
            name: 'write_file',
            args: { path: '/test/file.txt', content: 'content' },
            isClientInitiated: false,
            prompt_id: 'prompt-id-1',
          },
          status: 'executing',
          responseSubmittedToGemini: false,
          confirmationDetails: {
            onConfirm: mockOnConfirmExecuting,
            onCancel: vi.fn(),
            message: 'Write file?',
            displayedText: 'Write to /test/file.txt',
          },
          tool: {
            name: 'write_file',
            displayName: 'write_file',
            description: 'Write file',
            build: vi.fn(),
          } as any,
          invocation: {
            getDescription: () => 'Mock description',
          } as unknown as AnyToolInvocation,
          startTime: Date.now(),
          liveOutput: 'Writing...',
        } as TrackedExecutingToolCall,
      ];

      const { result } = renderTestHook(mixedStatusToolCalls);

      await act(async () => {
        await result.current.handleApprovalModeChange(ApprovalMode.YOLO);
      });

      // Only the awaiting_approval tool should be processed
      expect(mockOnConfirmAwaiting).toHaveBeenCalledTimes(1);
      expect(mockOnConfirmExecuting).not.toHaveBeenCalled();
    });
  });

  describe('handleFinishedEvent', () => {
    it('should add info message for MAX_TOKENS finish reason', async () => {
      // Setup mock to return a stream with MAX_TOKENS finish reason
      mockSendMessageStream.mockReturnValue(
        (async function* () {
          yield {
            type: ServerGeminiEventType.Content,
            value: 'This is a truncated response...',
          };
          yield {
            type: ServerGeminiEventType.Finished,
            value: { reason: 'MAX_TOKENS', usageMetadata: undefined },
          };
        })(),
      );

      const { result } = renderHook(() =>
        useGeminiStream(
          new MockedGeminiClientClass(mockConfig),
          [],
          mockAddItem,
          mockConfig,
          mockLoadedSettings,
          mockOnDebugMessage,
          mockHandleSlashCommand,
          false,
          () => 'vscode' as EditorType,
          () => {},
          () => Promise.resolve(),
          false,
          () => {},
          () => {},
          () => {},
          false, // visionModelPreviewEnabled
          () => {},
          80,
          24,
        ),
      );

      // Submit a query
      await act(async () => {
        await result.current.submitQuery('Generate long text');
      });

      // Check that the info message was added
      await waitFor(() => {
        expect(mockAddItem).toHaveBeenCalledWith(
          {
            type: 'info',
            text: '⚠️  Response truncated due to token limits.',
          },
          expect.any(Number),
        );
      });
    });

    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' as EditorType,
          () => {},
          () => Promise.resolve(),
          false,
          () => {},
          () => {},
          () => {},
          false, // visionModelPreviewEnabled
          () => {},
          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
      mockSendMessageStream.mockReturnValue(
        (async function* () {
          yield {
            type: ServerGeminiEventType.Content,
            value: 'Response with unspecified finish',
          };
          yield {
            type: ServerGeminiEventType.Finished,
            value: {
              reason: 'FINISH_REASON_UNSPECIFIED',
              usageMetadata: undefined,
            },
          };
        })(),
      );

      const { result } = renderHook(() =>
        useGeminiStream(
          new MockedGeminiClientClass(mockConfig),
          [],
          mockAddItem,
          mockConfig,
          mockLoadedSettings,
          mockOnDebugMessage,
          mockHandleSlashCommand,
          false,
          () => 'vscode' as EditorType,
          () => {},
          () => Promise.resolve(),
          false,
          () => {},
          () => {},
          () => {},
          false, // visionModelPreviewEnabled
          () => {},
          80,
          24,
        ),
      );

      // Submit a query
      await act(async () => {
        await result.current.submitQuery('Test unspecified finish');
      });

      // 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' as EditorType,
            () => {},
            () => Promise.resolve(),
            false,
            () => {},
            () => {},
            () => {},
            false, // visionModelPreviewEnabled
            vi.fn(),
            80,
            24,
          ),
        );

        await act(async () => {
          await result.current.submitQuery(`Test ${reason}`);
        });

        await waitFor(() => {
          expect(mockAddItem).toHaveBeenCalledWith(
            {
              type: 'info',
              text: message,
            },
            expect.any(Number),
          );
        });
      }
    });
  });

  it('should process @include commands, adding user turn after processing to prevent race conditions', async () => {
    const rawQuery = '@include file.txt Summarize this.';
    const processedQueryParts = [
      { text: 'Summarize this with content from @file.txt' },
      { text: 'File content...' },
    ];
    const userMessageTimestamp = Date.now();
    vi.spyOn(Date, 'now').mockReturnValue(userMessageTimestamp);

    handleAtCommandSpy.mockResolvedValue({
      processedQuery: processedQueryParts,
      shouldProceed: true,
    });

    const { result } = renderHook(() =>
      useGeminiStream(
        mockConfig.getGeminiClient() as GeminiClient,
        [],
        mockAddItem,
        mockConfig,
        mockLoadedSettings,
        mockOnDebugMessage,
        mockHandleSlashCommand,
        false, // shellModeActive
        vi.fn(), // getPreferredEditor
        vi.fn(), // onAuthError
        vi.fn(), // performMemoryRefresh
        false, // modelSwitched
        vi.fn(), // setModelSwitched
        vi.fn(), // onEditorClose
        vi.fn(), // onCancelSubmit
        false, // visionModelPreviewEnabled
        vi.fn(), // setShellInputFocused
        80, // terminalWidth
        24, // terminalHeight
      ),
    );

    await act(async () => {
      await result.current.submitQuery(rawQuery);
    });

    expect(handleAtCommandSpy).toHaveBeenCalledWith(
      expect.objectContaining({
        query: rawQuery,
      }),
    );

    expect(mockAddItem).toHaveBeenCalledWith(
      {
        type: MessageType.USER,
        text: rawQuery,
      },
      userMessageTimestamp,
    );

    // FIX: The expectation now matches the actual call signature.
    expect(mockSendMessageStream).toHaveBeenCalledWith(
      processedQueryParts, // Argument 1: The parts array directly
      expect.any(AbortSignal), // Argument 2: An AbortSignal
      expect.any(String), // Argument 3: The prompt_id string
      undefined, // Argument 4: Options (undefined for normal prompts)
    );
  });

  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: { reason: 'STOP', usageMetadata: undefined },
          };
        })(),
      );

      const { result } = renderHook(() =>
        useGeminiStream(
          new MockedGeminiClientClass(mockConfig),
          [],
          mockAddItem,
          mockConfig,
          mockLoadedSettings,
          mockOnDebugMessage,
          mockHandleSlashCommand,
          false,
          () => 'vscode' as EditorType,
          () => {},
          () => Promise.resolve(),
          false,
          () => {},
          () => {},
          () => {},
          false, // visionModelPreviewEnabled
          () => {},
          80,
          24,
        ),
      );

      // 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: { reason: 'STOP', usageMetadata: undefined },
          };
        })(),
      );

      // 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 accumulate streamed thought descriptions', async () => {
      mockSendMessageStream.mockReturnValue(
        (async function* () {
          yield {
            type: ServerGeminiEventType.Thought,
            value: { subject: '', description: 'thinking ' },
          };
          yield {
            type: ServerGeminiEventType.Thought,
            value: { subject: '', description: 'more' },
          };
          yield {
            type: ServerGeminiEventType.Finished,
            value: { reason: 'STOP', usageMetadata: undefined },
          };
        })(),
      );

      const { result } = renderHook(() =>
        useGeminiStream(
          new MockedGeminiClientClass(mockConfig),
          [],
          mockAddItem,
          mockConfig,
          mockLoadedSettings,
          mockOnDebugMessage,
          mockHandleSlashCommand,
          false,
          () => 'vscode' as EditorType,
          () => {},
          () => Promise.resolve(),
          false,
          () => {},
          () => {},
          () => {},
          false, // visionModelPreviewEnabled
          () => {},
          80,
          24,
        ),
      );

      await act(async () => {
        await result.current.submitQuery('Streamed thought');
      });

      await waitFor(() => {
        expect(result.current.thought?.description).toBe('thinking more');
      });
    });

    it('should memoize pendingHistoryItems', () => {
      mockUseReactToolScheduler.mockReturnValue([
        [],
        mockScheduleToolCalls,
        mockCancelAllToolCalls,
        mockMarkToolsAsSubmitted,
      ]);

      const { result, rerender } = renderHook(() =>
        useGeminiStream(
          mockConfig.getGeminiClient(),
          [],
          mockAddItem,
          mockConfig,
          mockLoadedSettings,
          mockOnDebugMessage,
          mockHandleSlashCommand,
          false,
          () => 'vscode' as EditorType,
          () => {},
          () => Promise.resolve(),
          false,
          () => {},
          () => {},
          () => {},
          false, // visionModelPreviewEnabled
          vi.fn(), // setShellInputFocused
          80,
          24,
        ),
      );

      const firstResult = result.current.pendingHistoryItems;
      rerender();
      const secondResult = result.current.pendingHistoryItems;

      expect(firstResult).toStrictEqual(secondResult);

      const newToolCalls: TrackedToolCall[] = [
        {
          request: { callId: 'call1', name: 'tool1', args: {} },
          status: 'executing',
          tool: {
            name: 'tool1',
            displayName: 'tool1',
            description: 'desc1',
            build: vi.fn(),
          },
          invocation: {
            getDescription: () => 'Mock description',
          },
        } as unknown as TrackedExecutingToolCall,
      ];

      mockUseReactToolScheduler.mockReturnValue([
        newToolCalls,
        mockScheduleToolCalls,
        mockCancelAllToolCalls,
        mockMarkToolsAsSubmitted,
      ]);

      rerender();
      const thirdResult = result.current.pendingHistoryItems;

      expect(thirdResult).not.toStrictEqual(secondResult);
    });

    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,
          mockConfig,
          mockLoadedSettings,
          mockOnDebugMessage,
          mockHandleSlashCommand,
          false,
          () => 'vscode' as EditorType,
          () => {},
          () => Promise.resolve(),
          false,
          () => {},
          () => {},
          () => {},
          false, // visionModelPreviewEnabled
          () => {},
          80,
          24,
        ),
      );

      // 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,
          mockConfig,
          mockLoadedSettings,
          mockOnDebugMessage,
          mockHandleSlashCommand,
          false,
          () => 'vscode' as EditorType,
          () => {},
          () => Promise.resolve(),
          false,
          () => {},
          () => {},
          () => {},
          false, // visionModelPreviewEnabled
          () => {},
          80,
          24,
        ),
      );

      // 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),
      );
    });
  });

  describe('Concurrent Execution Prevention', () => {
    it('should prevent concurrent submitQuery calls', async () => {
      let resolveFirstCall!: () => void;
      let resolveSecondCall!: () => void;

      const firstCallPromise = new Promise<void>((resolve) => {
        resolveFirstCall = resolve;
      });

      const secondCallPromise = new Promise<void>((resolve) => {
        resolveSecondCall = resolve;
      });

      // Mock a long-running stream for the first call
      const firstStream = (async function* () {
        yield {
          type: ServerGeminiEventType.Content,
          value: 'First call content',
        };
        await firstCallPromise; // Wait until we manually resolve
        yield { type: ServerGeminiEventType.Finished, value: 'STOP' };
      })();

      // Mock a stream for the second call (should not be used)
      const secondStream = (async function* () {
        yield {
          type: ServerGeminiEventType.Content,
          value: 'Second call content',
        };
        await secondCallPromise;
        yield { type: ServerGeminiEventType.Finished, value: 'STOP' };
      })();

      let callCount = 0;
      mockSendMessageStream.mockImplementation(() => {
        callCount++;
        if (callCount === 1) {
          return firstStream;
        } else {
          return secondStream;
        }
      });

      const { result } = renderTestHook();

      // Start first call
      const firstCallResult = act(async () => {
        await result.current.submitQuery('First query');
      });

      // Wait a bit to ensure first call has started
      await new Promise((resolve) => setTimeout(resolve, 10));

      // Try to start second call while first is still running
      const secondCallResult = act(async () => {
        await result.current.submitQuery('Second query');
      });

      // Resolve both calls
      resolveFirstCall();
      resolveSecondCall();

      await Promise.all([firstCallResult, secondCallResult]);

      // Verify only one call was made to sendMessageStream
      expect(mockSendMessageStream).toHaveBeenCalledTimes(1);
      expect(mockSendMessageStream).toHaveBeenCalledWith(
        'First query',
        expect.any(AbortSignal),
        expect.any(String),
        undefined,
      );

      // Verify only the first query was added to history
      const userMessages = mockAddItem.mock.calls.filter(
        (call) => call[0].type === MessageType.USER,
      );
      expect(userMessages).toHaveLength(1);
      expect(userMessages[0][0].text).toBe('First query');
    });

    it('should allow subsequent calls after first call completes', async () => {
      // Mock streams that complete immediately
      mockSendMessageStream
        .mockReturnValueOnce(
          (async function* () {
            yield {
              type: ServerGeminiEventType.Content,
              value: 'First response',
            };
            yield { type: ServerGeminiEventType.Finished, value: 'STOP' };
          })(),
        )
        .mockReturnValueOnce(
          (async function* () {
            yield {
              type: ServerGeminiEventType.Content,
              value: 'Second response',
            };
            yield { type: ServerGeminiEventType.Finished, value: 'STOP' };
          })(),
        );

      const { result } = renderTestHook();

      // First call
      await act(async () => {
        await result.current.submitQuery('First query');
      });

      // Second call after first completes
      await act(async () => {
        await result.current.submitQuery('Second query');
      });

      // Both calls should have been made
      expect(mockSendMessageStream).toHaveBeenCalledTimes(2);
      expect(mockSendMessageStream).toHaveBeenNthCalledWith(
        1,
        'First query',
        expect.any(AbortSignal),
        expect.any(String),
        undefined,
      );
      expect(mockSendMessageStream).toHaveBeenNthCalledWith(
        2,
        'Second query',
        expect.any(AbortSignal),
        expect.any(String),
        undefined,
      );
    });

    it('should reset execution flag even when query preparation fails', async () => {
      const { result } = renderTestHook();

      // First call with empty query (should fail in preparation)
      await act(async () => {
        await result.current.submitQuery('   '); // Empty trimmed query
      });

      // Second call should work normally
      await act(async () => {
        await result.current.submitQuery('Second query');
      });

      // Verify that only the second call was made (empty query is filtered out)
      expect(mockSendMessageStream).toHaveBeenCalledTimes(1);
      expect(mockSendMessageStream).toHaveBeenCalledWith(
        'Second query',
        expect.any(AbortSignal),
        expect.any(String),
        undefined,
      );
    });
  });

  // --- New tests focused on recent modifications ---
  describe('Vision Auto Switch Integration', () => {
    it('should call handleVisionSwitch and proceed to send when allowed', async () => {
      mockHandleVisionSwitch.mockResolvedValueOnce({ shouldProceed: true });
      mockSendMessageStream.mockReturnValue(
        (async function* () {
          yield { type: ServerGeminiEventType.Content, value: 'ok' };
          yield { type: ServerGeminiEventType.Finished, value: 'STOP' };
        })(),
      );

      const { result } = renderHook(() =>
        useGeminiStream(
          new MockedGeminiClientClass(mockConfig),
          [],
          mockAddItem,
          mockConfig,
          mockLoadedSettings,
          mockOnDebugMessage,
          mockHandleSlashCommand,
          false,
          () => 'vscode' as EditorType,
          () => {},
          () => Promise.resolve(),
          false,
          () => {},
          () => {},
          () => {},
          false, // visionModelPreviewEnabled
          vi.fn(), // setShellInputFocused
          80,
          24,
        ),
      );

      await act(async () => {
        await result.current.submitQuery('image prompt');
      });

      await waitFor(() => {
        expect(mockHandleVisionSwitch).toHaveBeenCalled();
        expect(mockSendMessageStream).toHaveBeenCalled();
      });
    });

    it('should gate submission when handleVisionSwitch returns shouldProceed=false', async () => {
      mockHandleVisionSwitch.mockResolvedValueOnce({ shouldProceed: false });

      const { result } = renderHook(() =>
        useGeminiStream(
          new MockedGeminiClientClass(mockConfig),
          [],
          mockAddItem,
          mockConfig,
          mockLoadedSettings,
          mockOnDebugMessage,
          mockHandleSlashCommand,
          false,
          () => 'vscode' as EditorType,
          () => {},
          () => Promise.resolve(),
          false,
          () => {},
          () => {},
          () => {},
          false, // visionModelPreviewEnabled
          vi.fn(), // setShellInputFocused
          80,
          24,
        ),
      );

      await act(async () => {
        await result.current.submitQuery('vision-gated');
      });

      // No call to API, no restoreOriginalModel needed since no override occurred
      expect(mockSendMessageStream).not.toHaveBeenCalled();
      expect(mockRestoreOriginalModel).not.toHaveBeenCalled();

      // Next call allowed (flag reset path)
      mockHandleVisionSwitch.mockResolvedValueOnce({ shouldProceed: true });
      mockSendMessageStream.mockReturnValue(
        (async function* () {
          yield { type: ServerGeminiEventType.Content, value: 'ok' };
          yield { type: ServerGeminiEventType.Finished, value: 'STOP' };
        })(),
      );
      await act(async () => {
        await result.current.submitQuery('after-gate');
      });
      await waitFor(() => {
        expect(mockSendMessageStream).toHaveBeenCalled();
      });
    });
  });

  describe('Model restore on completion and errors', () => {
    it('should restore model after successful stream completion', async () => {
      mockSendMessageStream.mockReturnValue(
        (async function* () {
          yield { type: ServerGeminiEventType.Content, value: 'content' };
          yield { type: ServerGeminiEventType.Finished, value: 'STOP' };
        })(),
      );

      const { result } = renderHook(() =>
        useGeminiStream(
          new MockedGeminiClientClass(mockConfig),
          [],
          mockAddItem,
          mockConfig,
          mockLoadedSettings,
          mockOnDebugMessage,
          mockHandleSlashCommand,
          false,
          () => 'vscode' as EditorType,
          () => {},
          () => Promise.resolve(),
          false,
          () => {},
          () => {},
          () => {},
          false, // visionModelPreviewEnabled
          vi.fn(), // setShellInputFocused
          80,
          24,
        ),
      );

      await act(async () => {
        await result.current.submitQuery('restore-success');
      });

      await waitFor(() => {
        expect(mockRestoreOriginalModel).toHaveBeenCalledTimes(1);
      });
    });

    it('should restore model when an error occurs during streaming', async () => {
      const testError = new Error('stream failure');
      mockSendMessageStream.mockReturnValue(
        (async function* () {
          yield { type: ServerGeminiEventType.Content, value: 'content' };
          throw testError;
        })(),
      );

      const { result } = renderHook(() =>
        useGeminiStream(
          new MockedGeminiClientClass(mockConfig),
          [],
          mockAddItem,
          mockConfig,
          mockLoadedSettings,
          mockOnDebugMessage,
          mockHandleSlashCommand,
          false,
          () => 'vscode' as EditorType,
          () => {},
          () => Promise.resolve(),
          false,
          () => {},
          () => {},
          () => {},
          false, // visionModelPreviewEnabled
          vi.fn(), // setShellInputFocused
          80,
          24,
        ),
      );

      await act(async () => {
        await result.current.submitQuery('restore-error');
      });

      await waitFor(() => {
        expect(mockRestoreOriginalModel).toHaveBeenCalledTimes(1);
      });
    });
  });

  describe('Loop Detection Confirmation', () => {
    beforeEach(() => {
      // Add mock for getLoopDetectionService to the config
      const mockLoopDetectionService = {
        disableForSession: vi.fn(),
      };
      mockConfig.getGeminiClient = vi.fn().mockReturnValue({
        ...new MockedGeminiClientClass(mockConfig),
        getLoopDetectionService: () => mockLoopDetectionService,
      });
    });

    it('should set loopDetectionConfirmationRequest when LoopDetected event is received', async () => {
      mockSendMessageStream.mockReturnValue(
        (async function* () {
          yield {
            type: ServerGeminiEventType.Content,
            value: 'Some content',
          };
          yield {
            type: ServerGeminiEventType.LoopDetected,
          };
        })(),
      );

      const { result } = renderTestHook();

      await act(async () => {
        await result.current.submitQuery('test query');
      });

      await waitFor(() => {
        expect(result.current.loopDetectionConfirmationRequest).not.toBeNull();
        expect(
          typeof result.current.loopDetectionConfirmationRequest?.onComplete,
        ).toBe('function');
      });
    });

    it('should disable loop detection and show message when user selects "disable"', async () => {
      const mockLoopDetectionService = {
        disableForSession: vi.fn(),
      };
      const mockClient = {
        ...new MockedGeminiClientClass(mockConfig),
        getLoopDetectionService: () => mockLoopDetectionService,
      };
      mockConfig.getGeminiClient = vi.fn().mockReturnValue(mockClient);

      mockSendMessageStream.mockReturnValueOnce(
        (async function* () {
          yield {
            type: ServerGeminiEventType.LoopDetected,
          };
        })(),
      );

      const { result } = renderTestHook([], mockClient);

      await act(async () => {
        await result.current.submitQuery('test query');
      });

      // Wait for confirmation request to be set
      await waitFor(() => {
        expect(result.current.loopDetectionConfirmationRequest).not.toBeNull();
      });

      // Simulate user selecting "disable"
      await act(async () => {
        result.current.loopDetectionConfirmationRequest?.onComplete({
          userSelection: 'disable',
        });
      });

      // Verify loop detection was disabled
      expect(mockLoopDetectionService.disableForSession).toHaveBeenCalledTimes(
        1,
      );

      // Verify confirmation request was cleared
      expect(result.current.loopDetectionConfirmationRequest).toBeNull();

      // Verify appropriate message was added
      expect(mockAddItem).toHaveBeenCalledWith(
        {
          type: 'info',
          text: 'Loop detection has been disabled for this session. Please try your request again.',
        },
        expect.any(Number),
      );
    });

    it('should keep loop detection enabled and show message when user selects "keep"', async () => {
      const mockLoopDetectionService = {
        disableForSession: vi.fn(),
      };
      const mockClient = {
        ...new MockedGeminiClientClass(mockConfig),
        getLoopDetectionService: () => mockLoopDetectionService,
      };
      mockConfig.getGeminiClient = vi.fn().mockReturnValue(mockClient);

      mockSendMessageStream.mockReturnValue(
        (async function* () {
          yield {
            type: ServerGeminiEventType.LoopDetected,
          };
        })(),
      );

      const { result } = renderTestHook();

      await act(async () => {
        await result.current.submitQuery('test query');
      });

      // Wait for confirmation request to be set
      await waitFor(() => {
        expect(result.current.loopDetectionConfirmationRequest).not.toBeNull();
      });

      // Simulate user selecting "keep"
      await act(async () => {
        result.current.loopDetectionConfirmationRequest?.onComplete({
          userSelection: 'keep',
        });
      });

      // Verify loop detection was NOT disabled
      expect(mockLoopDetectionService.disableForSession).not.toHaveBeenCalled();

      // Verify confirmation request was cleared
      expect(result.current.loopDetectionConfirmationRequest).toBeNull();

      // Verify appropriate message was added
      expect(mockAddItem).toHaveBeenCalledWith(
        {
          type: 'info',
          text: 'A potential loop was detected. This can happen due to repetitive tool calls or other model behavior. The request has been halted.',
        },
        expect.any(Number),
      );
    });

    it('should handle multiple loop detection events properly', async () => {
      const { result } = renderTestHook();

      // First loop detection - set up fresh mock for first call
      mockSendMessageStream.mockReturnValueOnce(
        (async function* () {
          yield {
            type: ServerGeminiEventType.LoopDetected,
          };
        })(),
      );

      // First loop detection
      await act(async () => {
        await result.current.submitQuery('first query');
      });

      await waitFor(() => {
        expect(result.current.loopDetectionConfirmationRequest).not.toBeNull();
      });

      // Simulate user selecting "keep" for first request
      await act(async () => {
        result.current.loopDetectionConfirmationRequest?.onComplete({
          userSelection: 'keep',
        });
      });

      expect(result.current.loopDetectionConfirmationRequest).toBeNull();

      // Verify first message was added
      expect(mockAddItem).toHaveBeenCalledWith(
        {
          type: 'info',
          text: 'A potential loop was detected. This can happen due to repetitive tool calls or other model behavior. The request has been halted.',
        },
        expect.any(Number),
      );

      // Second loop detection - set up fresh mock for second call
      mockSendMessageStream.mockReturnValueOnce(
        (async function* () {
          yield {
            type: ServerGeminiEventType.LoopDetected,
          };
        })(),
      );

      // Second loop detection
      await act(async () => {
        await result.current.submitQuery('second query');
      });

      await waitFor(() => {
        expect(result.current.loopDetectionConfirmationRequest).not.toBeNull();
      });

      // Simulate user selecting "disable" for second request
      await act(async () => {
        result.current.loopDetectionConfirmationRequest?.onComplete({
          userSelection: 'disable',
        });
      });

      expect(result.current.loopDetectionConfirmationRequest).toBeNull();

      // Verify second message was added
      expect(mockAddItem).toHaveBeenCalledWith(
        {
          type: 'info',
          text: 'Loop detection has been disabled for this session. Please try your request again.',
        },
        expect.any(Number),
      );
    });

    it('should process LoopDetected event after moving pending history to history', async () => {
      mockSendMessageStream.mockReturnValue(
        (async function* () {
          yield {
            type: ServerGeminiEventType.Content,
            value: 'Some response content',
          };
          yield {
            type: ServerGeminiEventType.LoopDetected,
          };
        })(),
      );

      const { result } = renderTestHook();

      await act(async () => {
        await result.current.submitQuery('test query');
      });

      // Verify that the content was added to history before the loop detection dialog
      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 waitFor(() => {
        expect(result.current.loopDetectionConfirmationRequest).not.toBeNull();
      });
    });
  });
});