Download packages/cli/src/ui/hooks/useGeminiStream.test.tsx from SaylorTwift/gemini-cli: direct link, hf CLI and curl.
- Browser
- Download file 140 kB
-
https://huggingface.co/SaylorTwift/gemini-cli/resolve/main/packages/cli/src/ui/hooks/useGeminiStream.test.tsx
- Command line
-
hf download hf://SaylorTwift/gemini-cli/packages/cli/src/ui/hooks/useGeminiStream.test.tsx
-
curl -L -o useGeminiStream.test.tsx https://huggingface.co/SaylorTwift/gemini-cli/resolve/main/packages/cli/src/ui/hooks/useGeminiStream.test.tsx
140 kB
| /** | |
| * @license | |
| * Copyright 2025 Google LLC | |
| * SPDX-License-Identifier: Apache-2.0 | |
| */ | |
| /* eslint-disable @typescript-eslint/no-explicit-any */ | |
| import { | |
| describe, | |
| it, | |
| expect, | |
| vi, | |
| beforeEach, | |
| type Mock, | |
| type MockInstance, | |
| } from 'vitest'; | |
| import { act } from 'react'; | |
| import { renderHookWithProviders } from '../../test-utils/render.js'; | |
| import { waitFor } from '../../test-utils/async.js'; | |
| import { useGeminiStream } from './useGeminiStream.js'; | |
| import { useKeypress } from './useKeypress.js'; | |
| import * as atCommandProcessor from './atCommandProcessor.js'; | |
| import { | |
| useToolScheduler, | |
| type TrackedToolCall, | |
| type TrackedCompletedToolCall, | |
| type TrackedExecutingToolCall, | |
| type TrackedCancelledToolCall, | |
| type TrackedWaitingToolCall, | |
| } from './useToolScheduler.js'; | |
| import type { | |
| Config, | |
| EditorType, | |
| AnyToolInvocation, | |
| AnyDeclarativeTool, | |
| SpanMetadata, | |
| CompletedToolCall, | |
| ToolCallRequestInfo, | |
| } from '@google/gemini-cli-core'; | |
| import { | |
| CoreToolCallStatus, | |
| ApprovalMode, | |
| AuthType, | |
| GeminiEventType as ServerGeminiEventType, | |
| ToolErrorType, | |
| ToolConfirmationOutcome, | |
| MessageBusType, | |
| tokenLimit, | |
| debugLogger, | |
| coreEvents, | |
| CoreEvent, | |
| SHELL_TOOL_NAME, | |
| MCPDiscoveryState, | |
| GeminiCliOperation, | |
| getPlanModeExitMessage, | |
| UPDATE_TOPIC_TOOL_NAME, | |
| TRUE_EMPTY_RESPONSE_MESSAGE, | |
| } from '@google/gemini-cli-core'; | |
| import type { Part, PartListUnion } from '@google/genai'; | |
| import type { UseHistoryManagerReturn } from './useHistoryManager.js'; | |
| import type { | |
| SlashCommandProcessorResult, | |
| HistoryItemWithoutId, | |
| HistoryItem, | |
| } from '../types.js'; | |
| import { MessageType, StreamingState } from '../types.js'; | |
| import type { LoadedSettings } from '../../config/settings.js'; | |
| import { findLastSafeSplitPoint } from '../utils/markdownUtilities.js'; | |
| import { theme } from '../semantic-colors.js'; | |
| // --- MOCKS --- | |
| const mockSendMessageStream = vi | |
| .fn() | |
| .mockReturnValue((async function* () {})()); | |
| const mockStartChat = vi.fn(); | |
| const mockMessageBus = { | |
| publish: vi.fn(), | |
| subscribe: vi.fn(), | |
| unsubscribe: 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(); | |
| let mockHistory: any[] = []; | |
| this.getHistory = vi.fn().mockImplementation(() => mockHistory); | |
| this.setHistory = vi.fn().mockImplementation((newHistory: any[]) => { | |
| mockHistory = [...newHistory]; | |
| }); | |
| this.generateContent = vi.fn().mockResolvedValue({ | |
| candidates: [ | |
| { content: { parts: [{ text: 'Got it. Focusing on tests only.' }] } }, | |
| ], | |
| }); | |
| this.getCurrentSequenceModel = vi.fn().mockReturnValue('test-model'); | |
| this.getChat = vi.fn().mockReturnValue({ | |
| recordCompletedToolCalls: 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(), | |
| }); | |
| this.getCurrentSequenceModel = vi | |
| .fn() | |
| .mockReturnValue('gemini-2.0-flash-exp'); | |
| }), | |
| ); | |
| const MockedUserPromptEvent = vi.hoisted(() => | |
| vi.fn().mockImplementation(() => {}), | |
| ); | |
| const mockParseAndFormatApiError = vi.hoisted(() => vi.fn()); | |
| const mockIsBackgroundExecutionData = vi.hoisted( | |
| () => | |
| (data: unknown): data is { pid?: number } => { | |
| if (typeof data !== 'object' || data === null) { | |
| return false; | |
| } | |
| const value = data as { | |
| pid?: unknown; | |
| command?: unknown; | |
| initialOutput?: unknown; | |
| }; | |
| return ( | |
| (value.pid === undefined || typeof value.pid === 'number') && | |
| (value.command === undefined || typeof value.command === 'string') && | |
| (value.initialOutput === undefined || | |
| typeof value.initialOutput === 'string') | |
| ); | |
| }, | |
| ); | |
| const MockValidationRequiredError = vi.hoisted( | |
| () => | |
| class extends Error { | |
| userHandled = false; | |
| }, | |
| ); | |
| const mockRunInDevTraceSpan = vi.hoisted(() => | |
| vi.fn(async (opts, fn) => { | |
| const metadata: SpanMetadata = { | |
| name: opts.operation, | |
| attributes: opts.attributes || {}, | |
| }; | |
| return await fn({ | |
| metadata, | |
| }); | |
| }), | |
| ); | |
| vi.mock('@google/gemini-cli-core', async (importOriginal) => { | |
| const actualCoreModule = (await importOriginal()) as any; | |
| return { | |
| ...actualCoreModule, | |
| isBackgroundExecutionData: mockIsBackgroundExecutionData, | |
| GitService: vi.fn(), | |
| GeminiClient: MockedGeminiClientClass, | |
| UserPromptEvent: MockedUserPromptEvent, | |
| ValidationRequiredError: MockValidationRequiredError, | |
| parseAndFormatApiError: mockParseAndFormatApiError, | |
| tokenLimit: vi.fn().mockReturnValue(100), // Mock tokenLimit | |
| recordToolCallInteractions: vi.fn().mockResolvedValue(undefined), | |
| getCodeAssistServer: vi.fn().mockReturnValue(undefined), | |
| runInDevTraceSpan: mockRunInDevTraceSpan, | |
| }; | |
| }); | |
| const mockUseToolScheduler = useToolScheduler as Mock; | |
| vi.mock('./useToolScheduler.js', async (importOriginal) => { | |
| const actualSchedulerModule = (await importOriginal()) as any; | |
| return { | |
| ...(actualSchedulerModule || {}), | |
| useToolScheduler: vi.fn(), | |
| }; | |
| }); | |
| vi.mock('./useKeypress.js', () => ({ | |
| useKeypress: vi.fn(), | |
| })); | |
| vi.mock('./useExecutionLifecycle.js', () => ({ | |
| useExecutionLifecycle: vi.fn().mockReturnValue({ | |
| handleShellCommand: vi.fn(), | |
| activeShellPtyId: null, | |
| lastShellOutputTime: 0, | |
| backgroundTaskCount: 0, | |
| isBackgroundTaskVisible: false, | |
| toggleBackgroundTasks: vi.fn(), | |
| backgroundCurrentExecution: vi.fn(), | |
| backgroundTasks: new Map(), | |
| dismissBackgroundTask: vi.fn(), | |
| registerBackgroundTask: vi.fn(), | |
| }), | |
| })); | |
| vi.mock('./atCommandProcessor.js'); | |
| vi.mock('../utils/markdownUtilities.js', () => ({ | |
| findLastSafeSplitPoint: vi.fn((s: string) => s.length), | |
| })); | |
| vi.mock('./useStateAndRef.js', async () => { | |
| const React = await vi.importActual<typeof import('react')>('react'); | |
| return { | |
| useStateAndRef: vi.fn((initial) => { | |
| // Keep the heavyweight test file lightweight, but still let | |
| // `isResponding` participate in real rerenders. | |
| if (initial === false) { | |
| const [state, setState] = React.useState(initial); | |
| const ref = React.useRef(initial); | |
| const setStateInternal = ( | |
| updater: typeof initial | ((prev: typeof initial) => typeof initial), | |
| ) => { | |
| const nextValue = | |
| typeof updater === 'function' | |
| ? (updater as (prev: typeof initial) => typeof initial)( | |
| ref.current, | |
| ) | |
| : updater; | |
| ref.current = nextValue; | |
| setState(nextValue); | |
| }; | |
| return [state, ref, setStateInternal]; | |
| } | |
| 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', async (importOriginal) => { | |
| const actual = (await importOriginal()) as any; | |
| return { | |
| ...actual, | |
| useSessionStats: vi.fn(() => ({ | |
| startNewPrompt: mockStartNewPrompt, | |
| addUsage: mockAddUsage, | |
| getPromptCount: vi.fn(() => 5), | |
| })), | |
| }; | |
| }); | |
| vi.mock('./slashCommandProcessor.js', () => ({ | |
| handleSlashCommand: vi.fn().mockReturnValue(false), | |
| })); | |
| vi.mock('./useAlternateBuffer.js', () => ({ | |
| useAlternateBuffer: vi.fn(() => false), | |
| })); | |
| // --- END MOCKS --- | |
| // --- Tests for useGeminiStream Hook --- | |
| describe('useGeminiStream', () => { | |
| let mockAddItem = vi.fn(); | |
| let mockOnDebugMessage = vi.fn(); | |
| let mockHandleSlashCommand = vi.fn().mockResolvedValue(false); | |
| let mockScheduleToolCalls: Mock; | |
| let mockCancelAllToolCalls: Mock; | |
| let mockMarkToolsAsSubmitted: Mock; | |
| let handleAtCommandSpy: MockInstance; | |
| const emptyHistory: HistoryItem[] = []; | |
| let capturedOnComplete: | |
| | ((tools: CompletedToolCall[]) => Promise<void>) | |
| | null = null; | |
| const mockGetPreferredEditor = vi.fn(() => 'vscode' as EditorType); | |
| const mockOnAuthError = vi.fn(); | |
| const mockPerformMemoryRefresh = vi.fn(() => Promise.resolve()); | |
| const mockSetModelSwitchedFromQuotaError = vi.fn(); | |
| const mockOnCancelSubmit = vi.fn(); | |
| const mockSetShellInputFocused = vi.fn(); | |
| const mockGetGeminiClient = vi.fn().mockImplementation(() => { | |
| const clientInstance = new MockedGeminiClientClass(mockConfig); | |
| return clientInstance; | |
| }); | |
| const mockMcpClientManager = { | |
| getDiscoveryState: vi.fn().mockReturnValue(MCPDiscoveryState.COMPLETED), | |
| getMcpServerCount: vi.fn().mockReturnValue(0), | |
| }; | |
| const mockConfig: Config = { | |
| apiKey: 'test-api-key', | |
| model: 'gemini-pro', | |
| sandbox: false, | |
| targetDir: '/test/dir', | |
| debugMode: false, | |
| question: undefined, | |
| coreTools: [], | |
| toolDiscoveryCommand: undefined, | |
| toolCallCommand: undefined, | |
| mcpServerCommand: undefined, | |
| mcpServers: undefined, | |
| userAgent: 'test-agent', | |
| userMemory: '', | |
| geminiMdFileCount: 0, | |
| alwaysSkipModificationConfirmation: false, | |
| vertexai: false, | |
| showMemoryUsage: false, | |
| contextFileName: undefined, | |
| storage: { | |
| getProjectTempDir: vi.fn(() => '/test/temp'), | |
| getProjectTempCheckpointsDir: vi.fn(() => '/test/temp/checkpoints'), | |
| } as any, | |
| getToolRegistry: vi.fn( | |
| () => ({ getToolSchemaList: vi.fn(() => []) }) as any, | |
| ), | |
| getProjectRoot: vi.fn(() => '/test/dir'), | |
| getCheckpointingEnabled: vi.fn(() => false), | |
| getGeminiClient: mockGetGeminiClient, | |
| getMcpClientManager: () => mockMcpClientManager as any, | |
| getApprovalMode: vi.fn(() => ApprovalMode.DEFAULT), | |
| getUsageStatisticsEnabled: () => true, | |
| getDebugMode: () => false, | |
| addHistory: vi.fn(), | |
| getSessionId: vi.fn(() => 'test-session-id'), | |
| setQuotaErrorOccurred: vi.fn(), | |
| resetBillingTurnState: vi.fn(), | |
| getQuotaErrorOccurred: vi.fn(() => false), | |
| getModel: vi.fn(() => 'gemini-2.5-pro'), | |
| getContentGeneratorConfig: vi.fn(() => ({ | |
| model: 'test-model', | |
| apiKey: 'test-key', | |
| vertexai: false, | |
| authType: AuthType.USE_GEMINI, | |
| })), | |
| getContentGenerator: vi.fn(), | |
| isInteractive: () => false, | |
| getExperiments: () => {}, | |
| getMaxSessionTurns: vi.fn(() => 100), | |
| getGlobalMemory: vi.fn(() => ''), | |
| getUserMemory: vi.fn(() => ''), | |
| getMessageBus: vi.fn(() => mockMessageBus), | |
| getBaseLlmClient: vi.fn(() => ({ | |
| generateContent: vi.fn().mockResolvedValue({ | |
| candidates: [ | |
| { content: { parts: [{ text: 'Got it. Focusing on tests only.' }] } }, | |
| ], | |
| }), | |
| })), | |
| getIdeMode: vi.fn(() => false), | |
| getEnableHooks: vi.fn(() => false), | |
| } as unknown as Config; | |
| beforeEach(() => { | |
| vi.clearAllMocks(); // Clear mocks before each test | |
| mockAddItem = vi.fn(); | |
| mockOnDebugMessage = vi.fn(); | |
| mockHandleSlashCommand = vi.fn().mockResolvedValue(false); | |
| // Mock return value for useReactToolScheduler | |
| mockScheduleToolCalls = vi.fn(); | |
| mockCancelAllToolCalls = vi.fn(); | |
| mockMarkToolsAsSubmitted = vi.fn(); | |
| // Reset properties of mockConfig if needed | |
| (mockConfig.getCheckpointingEnabled as Mock).mockReturnValue(false); | |
| (mockConfig.getApprovalMode as Mock).mockReturnValue(ApprovalMode.DEFAULT); | |
| // Default mock for useReactToolScheduler to prevent toolCalls being undefined initially | |
| mockUseToolScheduler.mockReturnValue([ | |
| [], // Default to empty array for toolCalls | |
| mockScheduleToolCalls, | |
| mockMarkToolsAsSubmitted, | |
| vi.fn(), // setToolCallsForDisplay | |
| mockCancelAllToolCalls, | |
| 0, // lastToolOutputTime | |
| ]); | |
| // 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'); | |
| vi.spyOn(coreEvents, 'emitFeedback'); | |
| }); | |
| const mockLoadedSettings: LoadedSettings = { | |
| merged: { | |
| preferredEditor: 'vscode', | |
| ui: { errorVerbosity: 'full' }, | |
| }, | |
| user: { path: '/user/settings.json', settings: {} }, | |
| workspace: { path: '/workspace/.gemini/settings.json', settings: {} }, | |
| errors: [], | |
| forScope: vi.fn(), | |
| setValue: vi.fn(), | |
| } as unknown as LoadedSettings; | |
| const renderTestHook = async ( | |
| initialToolCalls: TrackedToolCall[] = [], | |
| geminiClient?: any, | |
| loadedSettings: LoadedSettings = mockLoadedSettings, | |
| ) => { | |
| const client = geminiClient || mockConfig.getGeminiClient(); | |
| let lastToolCalls = initialToolCalls; | |
| const initialProps = { | |
| client, | |
| history: emptyHistory, | |
| addItem: mockAddItem as unknown as UseHistoryManagerReturn['addItem'], | |
| config: mockConfig, | |
| onDebugMessage: mockOnDebugMessage, | |
| handleSlashCommand: mockHandleSlashCommand as unknown as ( | |
| cmd: PartListUnion, | |
| ) => Promise<SlashCommandProcessorResult | false>, | |
| shellModeActive: false, | |
| loadedSettings, | |
| toolCalls: initialToolCalls, | |
| }; | |
| mockUseToolScheduler.mockImplementation((onComplete) => { | |
| capturedOnComplete = onComplete; | |
| return [ | |
| lastToolCalls, | |
| mockScheduleToolCalls, | |
| mockMarkToolsAsSubmitted, | |
| ( | |
| updater: | |
| | TrackedToolCall[] | |
| | ((prev: TrackedToolCall[]) => TrackedToolCall[]), | |
| ) => { | |
| lastToolCalls = | |
| typeof updater === 'function' ? updater(lastToolCalls) : updater; | |
| rerender({ ...initialProps, toolCalls: lastToolCalls }); | |
| }, | |
| (signal: AbortSignal) => { | |
| mockCancelAllToolCalls(signal); | |
| lastToolCalls = lastToolCalls.map((tc) => { | |
| if ( | |
| tc.status === CoreToolCallStatus.AwaitingApproval || | |
| tc.status === CoreToolCallStatus.Executing || | |
| tc.status === CoreToolCallStatus.Scheduled || | |
| tc.status === CoreToolCallStatus.Validating | |
| ) { | |
| return { | |
| ...tc, | |
| status: CoreToolCallStatus.Cancelled, | |
| response: { | |
| callId: tc.request.callId, | |
| responseParts: [], | |
| resultDisplay: 'Request cancelled.', | |
| }, | |
| responseSubmittedToGemini: true, | |
| } as any as TrackedCancelledToolCall; | |
| } | |
| return tc; | |
| }); | |
| rerender({ ...initialProps, toolCalls: lastToolCalls }); | |
| }, | |
| 0, | |
| ]; | |
| }); | |
| const { result, rerender } = await renderHookWithProviders( | |
| (props: typeof initialProps) => | |
| useGeminiStream( | |
| props.client, | |
| props.history, | |
| props.addItem, | |
| props.config, | |
| props.loadedSettings, | |
| props.onDebugMessage, | |
| props.handleSlashCommand, | |
| props.shellModeActive, | |
| mockGetPreferredEditor, | |
| mockOnAuthError, | |
| mockPerformMemoryRefresh, | |
| false, | |
| mockSetModelSwitchedFromQuotaError, | |
| mockOnCancelSubmit, | |
| mockSetShellInputFocused, | |
| 80, | |
| 24, | |
| ), | |
| { | |
| initialProps, | |
| }, | |
| ); | |
| return { | |
| result, | |
| rerender, | |
| mockMarkToolsAsSubmitted, | |
| mockSendMessageStream, | |
| client, | |
| }; | |
| }; | |
| // Helper to create mock tool calls - reduces boilerplate | |
| const createMockToolCall = ( | |
| toolName: string, | |
| callId: string, | |
| confirmationType: 'edit' | 'info', | |
| status: TrackedToolCall['status'] = CoreToolCallStatus.AwaitingApproval, | |
| mockOnConfirm: Mock = vi.fn(), | |
| ): TrackedWaitingToolCall => ({ | |
| request: { | |
| callId, | |
| name: toolName, | |
| args: {}, | |
| isClientInitiated: false, | |
| prompt_id: 'prompt-id-1', | |
| }, | |
| status: status as CoreToolCallStatus.AwaitingApproval, | |
| responseSubmittedToGemini: false, | |
| confirmationDetails: | |
| confirmationType === 'edit' | |
| ? { | |
| type: 'edit', | |
| title: 'Confirm Edit', | |
| fileName: 'file.txt', | |
| filePath: '/test/file.txt', | |
| fileDiff: 'fake diff', | |
| originalContent: 'old', | |
| newContent: 'new', | |
| onConfirm: mockOnConfirm, | |
| } | |
| : { | |
| type: 'info', | |
| title: `${toolName} confirmation`, | |
| prompt: `Execute ${toolName}?`, | |
| onConfirm: mockOnConfirm, | |
| }, | |
| tool: { | |
| name: toolName, | |
| displayName: toolName, | |
| description: `${toolName} description`, | |
| build: vi.fn(), | |
| } as any, | |
| invocation: { | |
| getDescription: () => 'Mock description', | |
| } as unknown as AnyToolInvocation, | |
| correlationId: `corr-${callId}`, | |
| }); | |
| // Helper to render hook with default parameters - reduces boilerplate | |
| const renderHookWithDefaults = async ( | |
| options: { | |
| shellModeActive?: boolean; | |
| onCancelSubmit?: () => void; | |
| setShellInputFocused?: (focused: boolean) => void; | |
| performMemoryRefresh?: () => Promise<void>; | |
| onAuthError?: () => void; | |
| setModelSwitched?: Mock; | |
| modelSwitched?: boolean; | |
| } = {}, | |
| ) => { | |
| const { | |
| shellModeActive = false, | |
| onCancelSubmit = () => {}, | |
| setShellInputFocused = () => {}, | |
| performMemoryRefresh = () => Promise.resolve(), | |
| onAuthError = () => {}, | |
| setModelSwitched = vi.fn(), | |
| modelSwitched = false, | |
| } = options; | |
| return renderHookWithProviders(() => | |
| useGeminiStream( | |
| new MockedGeminiClientClass(mockConfig), | |
| [], | |
| mockAddItem, | |
| mockConfig, | |
| mockLoadedSettings, | |
| mockOnDebugMessage, | |
| mockHandleSlashCommand, | |
| shellModeActive, | |
| () => 'vscode' as EditorType, | |
| onAuthError, | |
| performMemoryRefresh, | |
| modelSwitched, | |
| setModelSwitched, | |
| onCancelSubmit, | |
| setShellInputFocused, | |
| 80, | |
| 24, | |
| ), | |
| ); | |
| }; | |
| it('should not submit tool responses if not all tool calls are completed', async () => { | |
| const toolCalls: TrackedToolCall[] = [ | |
| { | |
| request: { | |
| callId: 'call1', | |
| name: 'tool1', | |
| args: {}, | |
| isClientInitiated: false, | |
| prompt_id: 'prompt-id-1', | |
| }, | |
| status: CoreToolCallStatus.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: CoreToolCallStatus.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 } = | |
| await 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 expose activePtyId for non-shell executing tools that report an execution ID', async () => { | |
| const remoteExecutingTool: TrackedExecutingToolCall = { | |
| request: { | |
| callId: 'remote-call-1', | |
| name: 'remote_agent_call', | |
| args: {}, | |
| isClientInitiated: false, | |
| prompt_id: 'prompt-id-remote', | |
| }, | |
| status: CoreToolCallStatus.Executing, | |
| responseSubmittedToGemini: false, | |
| tool: { | |
| name: 'remote_agent_call', | |
| displayName: 'Remote Agent', | |
| description: 'Remote agent execution', | |
| build: vi.fn(), | |
| } as any, | |
| invocation: { | |
| getDescription: () => 'Calling remote agent', | |
| } as unknown as AnyToolInvocation, | |
| startTime: Date.now(), | |
| liveOutput: 'working...', | |
| pid: 4242, | |
| }; | |
| const { result } = await renderTestHook([remoteExecutingTool]); | |
| expect(result.current.activePtyId).toBe(4242); | |
| }); | |
| 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: CoreToolCallStatus.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: CoreToolCallStatus.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; | |
| mockUseToolScheduler.mockImplementation((onComplete) => { | |
| capturedOnComplete = onComplete; | |
| return [ | |
| [], | |
| mockScheduleToolCalls, | |
| mockMarkToolsAsSubmitted, | |
| vi.fn(), | |
| mockCancelAllToolCalls, | |
| 0, | |
| ]; | |
| }); | |
| mockSendMessageStream.mockReturnValue( | |
| (async function* () { | |
| yield { | |
| type: ServerGeminiEventType.Content, | |
| value: 'Visible response text', | |
| }; | |
| })(), | |
| ); | |
| await renderHookWithProviders(() => | |
| useGeminiStream( | |
| new MockedGeminiClientClass(mockConfig), | |
| [], | |
| mockAddItem, | |
| mockConfig, | |
| mockLoadedSettings, | |
| mockOnDebugMessage, | |
| mockHandleSlashCommand, | |
| false, | |
| () => 'vscode' as EditorType, | |
| () => {}, | |
| () => Promise.resolve(), | |
| false, | |
| () => {}, | |
| () => {}, | |
| () => {}, | |
| 80, | |
| 24, | |
| ), | |
| ); | |
| // Trigger the onComplete callback with completed tools | |
| await act(async () => { | |
| if (capturedOnComplete) { | |
| // Wait a tick for refs to be set up | |
| await new Promise((resolve) => setTimeout(resolve, 0)); | |
| 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', | |
| undefined, | |
| expectedMergedResponse, | |
| ); | |
| }); | |
| it('should inject steering hint prompt for continuation', async () => { | |
| const toolCallResponseParts: Part[] = [{ text: 'tool final response' }]; | |
| const completedToolCalls: TrackedToolCall[] = [ | |
| { | |
| request: { | |
| callId: 'call1', | |
| name: 'tool1', | |
| args: {}, | |
| isClientInitiated: false, | |
| prompt_id: 'prompt-id-ack', | |
| }, | |
| status: 'success', | |
| responseSubmittedToGemini: false, | |
| response: { | |
| callId: 'call1', | |
| responseParts: toolCallResponseParts, | |
| errorType: undefined, | |
| }, | |
| tool: { | |
| displayName: 'MockTool', | |
| }, | |
| invocation: { | |
| getDescription: () => `Mock description`, | |
| } as unknown as AnyToolInvocation, | |
| } as TrackedCompletedToolCall, | |
| ]; | |
| mockSendMessageStream.mockReturnValue( | |
| (async function* () { | |
| yield { | |
| type: ServerGeminiEventType.Content, | |
| value: 'Applied the requested adjustment.', | |
| }; | |
| })(), | |
| ); | |
| let capturedOnComplete: | |
| | ((completedTools: TrackedToolCall[]) => Promise<void>) | |
| | null = null; | |
| mockUseToolScheduler.mockImplementation((onComplete) => { | |
| capturedOnComplete = onComplete; | |
| return [ | |
| [], | |
| mockScheduleToolCalls, | |
| mockMarkToolsAsSubmitted, | |
| vi.fn(), | |
| mockCancelAllToolCalls, | |
| 0, | |
| ]; | |
| }); | |
| await renderHookWithProviders(() => | |
| useGeminiStream( | |
| new MockedGeminiClientClass(mockConfig), | |
| [], | |
| mockAddItem, | |
| mockConfig, | |
| mockLoadedSettings, | |
| mockOnDebugMessage, | |
| mockHandleSlashCommand, | |
| false, | |
| () => 'vscode' as EditorType, | |
| () => {}, | |
| () => Promise.resolve(), | |
| false, | |
| () => {}, | |
| () => {}, | |
| () => {}, | |
| 80, | |
| 24, | |
| undefined, | |
| () => 'focus on tests only', | |
| ), | |
| ); | |
| await act(async () => { | |
| if (capturedOnComplete) { | |
| await new Promise((resolve) => setTimeout(resolve, 0)); | |
| await capturedOnComplete(completedToolCalls); | |
| } | |
| }); | |
| await waitFor(() => { | |
| expect(mockSendMessageStream).toHaveBeenCalledTimes(1); | |
| }); | |
| const sentParts = mockSendMessageStream.mock.calls[0][0] as Part[]; | |
| const injectedHintPart = sentParts[0] as { text?: string }; | |
| expect(injectedHintPart.text).toContain('User steering update:'); | |
| expect(injectedHintPart.text).toContain( | |
| '<user_input>\nfocus on tests only\n</user_input>', | |
| ); | |
| expect(injectedHintPart.text).toContain( | |
| 'Classify it as ADD_TASK, MODIFY_TASK, CANCEL_TASK, or EXTRA_CONTEXT.', | |
| ); | |
| expect(injectedHintPart.text).toContain( | |
| 'Do not cancel/skip tasks unless the user explicitly cancels them.', | |
| ); | |
| expect(mockRunInDevTraceSpan).toHaveBeenCalledWith( | |
| expect.objectContaining({ | |
| operation: GeminiCliOperation.SystemPrompt, | |
| }), | |
| expect.any(Function), | |
| ); | |
| const spanArgs = mockRunInDevTraceSpan.mock.calls[0]; | |
| const fn = spanArgs[1]; | |
| const metadata = { attributes: {} }; | |
| await act(async () => { | |
| await fn({ metadata }); | |
| }); | |
| expect(metadata).toMatchObject({ | |
| input: sentParts, | |
| }); | |
| }); | |
| it('should auto-nudge the model when tool execution succeeds but model stream is empty', async () => { | |
| const toolCallResponseParts: Part[] = [{ text: 'tool final response' }]; | |
| const completedToolCalls: TrackedToolCall[] = [ | |
| { | |
| request: { | |
| callId: 'call1', | |
| name: 'tool1', | |
| args: {}, | |
| isClientInitiated: false, | |
| prompt_id: 'prompt-id-ack', | |
| }, | |
| status: 'success', | |
| responseSubmittedToGemini: false, | |
| response: { | |
| callId: 'call1', | |
| responseParts: toolCallResponseParts, | |
| errorType: undefined, | |
| }, | |
| tool: { | |
| displayName: 'MockTool', | |
| }, | |
| invocation: { | |
| getDescription: () => `Mock description`, | |
| } as unknown as AnyToolInvocation, | |
| } as TrackedCompletedToolCall, | |
| ]; | |
| let callCount = 0; | |
| mockSendMessageStream.mockImplementation(() => { | |
| callCount += 1; | |
| if (callCount === 1) { | |
| return (async function* () {})(); | |
| } else { | |
| return (async function* () { | |
| yield { | |
| type: ServerGeminiEventType.Content, | |
| value: | |
| 'I have analyzed the empty response. Here is the final answer.', | |
| }; | |
| })(); | |
| } | |
| }); | |
| let capturedOnComplete: | |
| | ((completedTools: TrackedToolCall[]) => Promise<void>) | |
| | null = null; | |
| mockUseToolScheduler.mockImplementation((onComplete) => { | |
| capturedOnComplete = onComplete; | |
| return [ | |
| [], | |
| mockScheduleToolCalls, | |
| mockMarkToolsAsSubmitted, | |
| vi.fn(), | |
| mockCancelAllToolCalls, | |
| 0, | |
| ]; | |
| }); | |
| await renderHookWithProviders(() => | |
| useGeminiStream( | |
| new MockedGeminiClientClass(mockConfig), | |
| [], | |
| mockAddItem, | |
| mockConfig, | |
| mockLoadedSettings, | |
| mockOnDebugMessage, | |
| mockHandleSlashCommand, | |
| false, | |
| () => 'vscode' as EditorType, | |
| () => {}, | |
| () => Promise.resolve(), | |
| false, | |
| () => {}, | |
| () => {}, | |
| () => {}, | |
| 80, | |
| 24, | |
| undefined, | |
| () => 'focus on tests only', | |
| ), | |
| ); | |
| await act(async () => { | |
| if (capturedOnComplete) { | |
| await new Promise((resolve) => setTimeout(resolve, 0)); | |
| await capturedOnComplete(completedToolCalls); | |
| } | |
| }); | |
| await waitFor(() => { | |
| expect(mockSendMessageStream).toHaveBeenCalledTimes(2); | |
| }); | |
| const sentParts = mockSendMessageStream.mock.calls[1][0] as Part[]; | |
| expect(sentParts[0].text).toContain( | |
| '[System: You successfully executed a tool but returned an empty response. Please analyze the tool output and explain your progress or final answer.]', | |
| ); | |
| }); | |
| it('should handle all tool calls being cancelled by rolling back the history', async () => { | |
| const cancelledToolCalls: TrackedToolCall[] = [ | |
| { | |
| request: { | |
| callId: 'topic1', | |
| name: UPDATE_TOPIC_TOOL_NAME, | |
| args: {}, | |
| isClientInitiated: false, | |
| prompt_id: 'prompt-id-3', | |
| }, | |
| status: CoreToolCallStatus.Success, | |
| response: { | |
| callId: 'topic1', | |
| responseParts: [ | |
| { | |
| functionResponse: { | |
| name: UPDATE_TOPIC_TOOL_NAME, | |
| id: 'topic1', | |
| response: {}, | |
| }, | |
| }, | |
| ], | |
| }, | |
| tool: { displayName: 'Update Topic Context' }, | |
| invocation: { getDescription: () => 'Updating topic' }, | |
| } as any, | |
| { | |
| request: { | |
| callId: '1', | |
| name: 'testTool', | |
| args: {}, | |
| isClientInitiated: false, | |
| prompt_id: 'prompt-id-3', | |
| }, | |
| status: CoreToolCallStatus.Cancelled, | |
| response: { | |
| callId: '1', | |
| responseParts: [{ text: CoreToolCallStatus.Cancelled }], | |
| errorType: undefined, // FIX: Added missing property | |
| }, | |
| responseSubmittedToGemini: false, | |
| tool: { | |
| displayName: 'mock tool', | |
| }, | |
| invocation: { | |
| getDescription: () => `Mock description`, | |
| }, | |
| } as any, | |
| ]; | |
| const client = new MockedGeminiClientClass(mockConfig); | |
| client.setHistory([{ role: 'user', parts: [{ text: 'User prompt' }] }]); | |
| // Capture the onComplete callback | |
| let capturedOnComplete: | |
| | ((completedTools: TrackedToolCall[]) => Promise<void>) | |
| | null = null; | |
| mockUseToolScheduler.mockImplementation((onComplete) => { | |
| capturedOnComplete = onComplete; | |
| return [ | |
| [], | |
| mockScheduleToolCalls, | |
| mockMarkToolsAsSubmitted, | |
| vi.fn(), | |
| mockCancelAllToolCalls, | |
| 0, | |
| ]; | |
| }); | |
| const { result } = await renderHookWithProviders(() => | |
| useGeminiStream( | |
| client, | |
| [], | |
| mockAddItem, | |
| mockConfig, | |
| mockLoadedSettings, | |
| mockOnDebugMessage, | |
| mockHandleSlashCommand, | |
| false, | |
| () => 'vscode' as EditorType, | |
| () => {}, | |
| () => Promise.resolve(), | |
| false, | |
| () => {}, | |
| () => {}, | |
| () => {}, | |
| 80, | |
| 24, | |
| ), | |
| ); | |
| // Call submitQuery to populate the user turn and set historyLengthAfterUserPromptRef | |
| await act(async () => { | |
| // eslint-disable-next-line @typescript-eslint/no-floating-promises | |
| result.current.submitQuery('User prompt'); | |
| }); | |
| // Model issues a functionCall request, which appends to history | |
| client.setHistory([ | |
| { role: 'user', parts: [{ text: 'User prompt' }] }, | |
| { | |
| role: 'model', | |
| parts: [{ functionCall: { name: 'testTool', args: {} } }], | |
| }, | |
| ]); | |
| // Trigger the onComplete callback with cancelled tools | |
| await act(async () => { | |
| if (capturedOnComplete) { | |
| // Wait a tick for refs to be set up | |
| await new Promise((resolve) => setTimeout(resolve, 0)); | |
| await capturedOnComplete(cancelledToolCalls); | |
| } | |
| }); | |
| await waitFor(() => { | |
| expect(mockMarkToolsAsSubmitted).toHaveBeenCalledWith(['topic1', '1']); | |
| // Should NOT have appended cancellations via addHistory | |
| expect(client.addHistory).not.toHaveBeenCalled(); | |
| // Should have rolled history back to pre-model length (1) | |
| expect(client.getHistory().length).toBe(1); | |
| // Ensure we do NOT call back to the API a second time (only the initial user turn was sent) | |
| expect(mockSendMessageStream).toHaveBeenCalledTimes(1); | |
| }); | |
| }); | |
| it('should record tool responses in history when the model was switched due to a quota error', async () => { | |
| // Regression test: returning early on a quota-triggered model switch | |
| // without recording the responses leaves the already-recorded | |
| // functionCall unpaired, which corrupts all subsequent requests. | |
| const responseParts: Part[] = [ | |
| { | |
| functionResponse: { | |
| name: 'testTool', | |
| id: 'call1', | |
| response: { output: 'tool result' }, | |
| }, | |
| }, | |
| ]; | |
| const completedToolCalls: TrackedToolCall[] = [ | |
| { | |
| request: { | |
| callId: 'call1', | |
| name: 'testTool', | |
| args: {}, | |
| isClientInitiated: false, | |
| prompt_id: 'prompt-id-quota', | |
| }, | |
| status: CoreToolCallStatus.Success, | |
| responseSubmittedToGemini: false, | |
| response: { | |
| callId: 'call1', | |
| responseParts, | |
| errorType: undefined, | |
| }, | |
| tool: { displayName: 'MockTool' }, | |
| invocation: { | |
| getDescription: () => `Mock description`, | |
| } as unknown as AnyToolInvocation, | |
| } as TrackedCompletedToolCall, | |
| ]; | |
| const client = new MockedGeminiClientClass(mockConfig); | |
| const mockConsumeUserHint = vi.fn(() => 'switch to the nprd database'); | |
| let capturedOnComplete: | |
| | ((completedTools: TrackedToolCall[]) => Promise<void>) | |
| | null = null; | |
| mockUseToolScheduler.mockImplementation((onComplete) => { | |
| capturedOnComplete = onComplete; | |
| return [ | |
| [], | |
| mockScheduleToolCalls, | |
| mockMarkToolsAsSubmitted, | |
| vi.fn(), | |
| mockCancelAllToolCalls, | |
| 0, | |
| ]; | |
| }); | |
| await renderHookWithProviders(() => | |
| useGeminiStream( | |
| client, | |
| [], | |
| mockAddItem, | |
| mockConfig, | |
| mockLoadedSettings, | |
| mockOnDebugMessage, | |
| mockHandleSlashCommand, | |
| false, | |
| () => 'vscode' as EditorType, | |
| () => {}, | |
| () => Promise.resolve(), | |
| true, // modelSwitchedFromQuotaError | |
| () => {}, | |
| () => {}, | |
| () => {}, | |
| 80, | |
| 24, | |
| false, | |
| mockConsumeUserHint, | |
| ), | |
| ); | |
| await act(async () => { | |
| if (capturedOnComplete) { | |
| await new Promise((resolve) => setTimeout(resolve, 0)); | |
| await capturedOnComplete(completedToolCalls); | |
| } | |
| }); | |
| await waitFor(() => { | |
| expect(mockMarkToolsAsSubmitted).toHaveBeenCalledWith(['call1']); | |
| // The tool response must be paired with its functionCall in history, | |
| // with no steering-hint text ahead of it... | |
| expect(client.addHistory).toHaveBeenCalledWith({ | |
| role: 'user', | |
| parts: responseParts, | |
| }); | |
| // ...the turn must NOT auto-continue on the fallback model... | |
| expect(mockSendMessageStream).not.toHaveBeenCalled(); | |
| // ...and the pending hint is left for the next real submit. | |
| expect(mockConsumeUserHint).not.toHaveBeenCalled(); | |
| }); | |
| }); | |
| it('should NOT stop responding when only update_topic is called', async () => { | |
| const topicToolCalls: TrackedToolCall[] = [ | |
| { | |
| request: { | |
| callId: 'topic1', | |
| name: UPDATE_TOPIC_TOOL_NAME, | |
| args: {}, | |
| isClientInitiated: false, | |
| prompt_id: 'prompt-id-3', | |
| }, | |
| status: CoreToolCallStatus.Success, | |
| response: { | |
| callId: 'topic1', | |
| responseParts: [ | |
| { | |
| functionResponse: { | |
| name: UPDATE_TOPIC_TOOL_NAME, | |
| id: 'topic1', | |
| response: {}, | |
| }, | |
| }, | |
| ], | |
| }, | |
| tool: { displayName: 'Update Topic Context' }, | |
| invocation: { getDescription: () => 'Updating topic' }, | |
| } as any, | |
| ]; | |
| const client = new MockedGeminiClientClass(mockConfig); | |
| // Capture the onComplete callback | |
| let capturedOnComplete: | |
| | ((completedTools: TrackedToolCall[]) => Promise<void>) | |
| | null = null; | |
| mockUseToolScheduler.mockImplementation((onComplete) => { | |
| capturedOnComplete = onComplete; | |
| return [ | |
| topicToolCalls, | |
| vi.fn(), | |
| mockMarkToolsAsSubmitted, | |
| vi.fn(), | |
| vi.fn(), | |
| 0, | |
| ]; | |
| }); | |
| await renderHookWithProviders(() => | |
| useGeminiStream( | |
| client, | |
| [], | |
| mockAddItem, | |
| mockConfig, | |
| mockLoadedSettings, | |
| mockOnDebugMessage, | |
| mockHandleSlashCommand, | |
| false, | |
| () => 'vscode' as EditorType, | |
| () => {}, | |
| () => Promise.resolve(), | |
| false, | |
| () => {}, | |
| () => {}, | |
| () => {}, | |
| 80, | |
| 24, | |
| ), | |
| ); | |
| // Trigger the onComplete callback with the topic tool | |
| await act(async () => { | |
| if (capturedOnComplete) { | |
| await capturedOnComplete(topicToolCalls); | |
| } | |
| }); | |
| await waitFor(() => { | |
| // The streaming state should still be Responding because we didn't cancel anything important | |
| // and we expect a continuation. | |
| expect(mockMarkToolsAsSubmitted).toHaveBeenCalledWith(['topic1']); | |
| // Should HAVE called back to the API for continuation | |
| expect(mockSendMessageStream).toHaveBeenCalled(); | |
| }); | |
| }); | |
| it('should stop agent execution immediately when a tool call returns STOP_EXECUTION error', async () => { | |
| const stopExecutionToolCalls: TrackedCompletedToolCall[] = [ | |
| { | |
| request: { | |
| callId: 'stop-call', | |
| name: 'stopTool', | |
| args: {}, | |
| isClientInitiated: false, | |
| prompt_id: 'prompt-id-stop', | |
| }, | |
| status: CoreToolCallStatus.Error, | |
| response: { | |
| callId: 'stop-call', | |
| responseParts: [{ text: 'error occurred' }], | |
| errorType: ToolErrorType.STOP_EXECUTION, | |
| error: new Error('Stop reason from hook'), | |
| resultDisplay: undefined, | |
| }, | |
| responseSubmittedToGemini: false, | |
| tool: { | |
| displayName: 'stop tool', | |
| }, | |
| invocation: { | |
| getDescription: () => `Mock description`, | |
| } as unknown as AnyToolInvocation, | |
| } as unknown as TrackedCompletedToolCall, | |
| ]; | |
| const client = new MockedGeminiClientClass(mockConfig); | |
| const { result } = await renderTestHook([], client); | |
| // Trigger the onComplete callback with STOP_EXECUTION tool | |
| await act(async () => { | |
| if (capturedOnComplete) { | |
| await capturedOnComplete(stopExecutionToolCalls); | |
| } | |
| }); | |
| await waitFor(() => { | |
| expect(mockMarkToolsAsSubmitted).toHaveBeenCalledWith(['stop-call']); | |
| // Should add an info message to history | |
| expect(mockAddItem).toHaveBeenCalledWith( | |
| expect.objectContaining({ | |
| type: MessageType.INFO, | |
| text: expect.stringContaining( | |
| 'Agent execution stopped: Stop reason from hook', | |
| ), | |
| }), | |
| ); | |
| // Ensure we do NOT call back to the API | |
| expect(mockSendMessageStream).not.toHaveBeenCalled(); | |
| // Streaming state should be Idle | |
| expect(result.current.streamingState).toBe(StreamingState.Idle); | |
| }); | |
| const infoTexts = mockAddItem.mock.calls.map( | |
| ([item]) => (item as { text?: string }).text ?? '', | |
| ); | |
| expect( | |
| infoTexts.some((text) => | |
| text.includes( | |
| 'Some internal tool attempts failed before this final error', | |
| ), | |
| ), | |
| ).toBe(false); | |
| expect( | |
| infoTexts.some((text) => | |
| text.includes('This request failed. Press F12 for diagnostics'), | |
| ), | |
| ).toBe(false); | |
| }); | |
| it('should add a compact suppressed-error note before STOP_EXECUTION terminal info in low verbosity mode', async () => { | |
| const stopExecutionToolCalls: TrackedCompletedToolCall[] = [ | |
| { | |
| request: { | |
| callId: 'stop-call', | |
| name: 'stopTool', | |
| args: {}, | |
| isClientInitiated: false, | |
| prompt_id: 'prompt-id-stop', | |
| }, | |
| status: CoreToolCallStatus.Error, | |
| response: { | |
| callId: 'stop-call', | |
| responseParts: [{ text: 'error occurred' }], | |
| errorType: ToolErrorType.STOP_EXECUTION, | |
| error: new Error('Stop reason from hook'), | |
| resultDisplay: undefined, | |
| }, | |
| responseSubmittedToGemini: false, | |
| tool: { | |
| displayName: 'stop tool', | |
| }, | |
| invocation: { | |
| getDescription: () => `Mock description`, | |
| } as unknown as AnyToolInvocation, | |
| } as unknown as TrackedCompletedToolCall, | |
| ]; | |
| const lowVerbositySettings = { | |
| // eslint-disable-next-line @typescript-eslint/no-misused-spread | |
| ...mockLoadedSettings, | |
| merged: { | |
| ...mockLoadedSettings.merged, | |
| ui: { errorVerbosity: 'low' }, | |
| }, | |
| } as LoadedSettings; | |
| const client = new MockedGeminiClientClass(mockConfig); | |
| const { result } = await renderTestHook([], client, lowVerbositySettings); | |
| await act(async () => { | |
| if (capturedOnComplete) { | |
| await capturedOnComplete(stopExecutionToolCalls); | |
| } | |
| }); | |
| await waitFor(() => { | |
| expect(mockMarkToolsAsSubmitted).toHaveBeenCalledWith(['stop-call']); | |
| expect(mockSendMessageStream).not.toHaveBeenCalled(); | |
| expect(result.current.streamingState).toBe(StreamingState.Idle); | |
| }); | |
| const infoTexts = mockAddItem.mock.calls.map( | |
| ([item]) => (item as { text?: string }).text ?? '', | |
| ); | |
| const noteIndex = infoTexts.findIndex((text) => | |
| text.includes( | |
| 'Some internal tool attempts failed before this final error', | |
| ), | |
| ); | |
| const stopIndex = infoTexts.findIndex((text) => | |
| text.includes('Agent execution stopped: Stop reason from hook'), | |
| ); | |
| const failureHintIndex = infoTexts.findIndex((text) => | |
| text.includes('This request failed. Press F12 for diagnostics'), | |
| ); | |
| expect(noteIndex).toBeGreaterThanOrEqual(0); | |
| expect(stopIndex).toBeGreaterThanOrEqual(0); | |
| // The failure hint should NOT be present if the suppressed error note was shown | |
| expect(failureHintIndex).toBe(-1); | |
| expect(noteIndex).toBeLessThan(stopIndex); | |
| }); | |
| it('should rollback multiple cancelled tool calls rather than appending them to history', 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: CoreToolCallStatus.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: CoreToolCallStatus.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); | |
| client.setHistory([{ role: 'user', parts: [{ text: 'User prompt' }] }]); | |
| let capturedOnComplete: | |
| | ((completedTools: TrackedToolCall[]) => Promise<void>) | |
| | null = null; | |
| mockUseToolScheduler.mockImplementation((onComplete) => { | |
| capturedOnComplete = onComplete; | |
| return [ | |
| [], | |
| mockScheduleToolCalls, | |
| mockMarkToolsAsSubmitted, | |
| vi.fn(), | |
| mockCancelAllToolCalls, | |
| 0, | |
| ]; | |
| }); | |
| const { result } = await renderHookWithProviders(() => | |
| useGeminiStream( | |
| client, | |
| [], | |
| mockAddItem, | |
| mockConfig, | |
| mockLoadedSettings, | |
| mockOnDebugMessage, | |
| mockHandleSlashCommand, | |
| false, | |
| () => 'vscode' as EditorType, | |
| () => {}, | |
| () => Promise.resolve(), | |
| false, | |
| () => {}, | |
| () => {}, | |
| () => {}, | |
| 80, | |
| 24, | |
| ), | |
| ); | |
| // Call submitQuery to populate the user turn and set historyLengthAfterUserPromptRef | |
| await act(async () => { | |
| // eslint-disable-next-line @typescript-eslint/no-floating-promises | |
| result.current.submitQuery('User prompt'); | |
| }); | |
| // Model issues model turns, which appends to history | |
| client.setHistory([ | |
| { role: 'user', parts: [{ text: 'User prompt' }] }, | |
| { role: 'model', parts: [{ functionCall: { name: 'toolA', args: {} } }] }, | |
| ]); | |
| // Trigger the onComplete callback with multiple cancelled tools | |
| await act(async () => { | |
| if (capturedOnComplete) { | |
| // Wait a tick for refs to be set up | |
| await new Promise((resolve) => setTimeout(resolve, 0)); | |
| await capturedOnComplete(allCancelledTools); | |
| } | |
| }); | |
| await waitFor(() => { | |
| // The tools should be marked as submitted locally | |
| expect(mockMarkToolsAsSubmitted).toHaveBeenCalledWith([ | |
| 'cancel-1', | |
| 'cancel-2', | |
| ]); | |
| // Crucially, addHistory should NOT be called | |
| expect(client.addHistory).not.toHaveBeenCalled(); | |
| // Instead, history should be rolled back to pre-model length (1) | |
| expect(client.getHistory().length).toBe(1); | |
| // No message should be sent back to the API for a turn with only cancellations | |
| expect(mockSendMessageStream).toHaveBeenCalledTimes(1); | |
| }); | |
| }); | |
| 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: CoreToolCallStatus.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: CoreToolCallStatus.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; | |
| mockUseToolScheduler.mockImplementation((onComplete) => { | |
| capturedOnComplete = onComplete; | |
| return [ | |
| currentToolCalls, | |
| mockScheduleToolCalls, | |
| mockMarkToolsAsSubmitted, | |
| vi.fn(), // setToolCallsForDisplay | |
| mockCancelAllToolCalls, | |
| 0, | |
| ]; | |
| }); | |
| const { result, rerender } = await renderHookWithProviders(() => | |
| useGeminiStream( | |
| new MockedGeminiClientClass(mockConfig), | |
| [], | |
| mockAddItem, | |
| mockConfig, | |
| mockLoadedSettings, | |
| mockOnDebugMessage, | |
| mockHandleSlashCommand, | |
| false, | |
| () => 'vscode' as EditorType, | |
| () => {}, | |
| () => Promise.resolve(), | |
| false, | |
| () => {}, | |
| () => {}, | |
| () => {}, | |
| 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; | |
| mockUseToolScheduler.mockImplementation((onComplete) => { | |
| capturedOnComplete = onComplete; | |
| return [ | |
| completedToolCalls, | |
| mockScheduleToolCalls, | |
| mockMarkToolsAsSubmitted, | |
| vi.fn(), // setToolCallsForDisplay | |
| mockCancelAllToolCalls, | |
| 0, | |
| ]; | |
| }); | |
| 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) { | |
| // Wait a tick for refs to be set up | |
| await new Promise((resolve) => setTimeout(resolve, 0)); | |
| await capturedOnComplete(completedToolCalls); | |
| } | |
| }); | |
| // 5. Wait for submitQuery to be called | |
| await waitFor(() => { | |
| expect(mockSendMessageStream).toHaveBeenCalledWith( | |
| toolCallResponseParts, | |
| expect.any(AbortSignal), | |
| 'prompt-id-4', | |
| undefined, | |
| toolCallResponseParts, | |
| ); | |
| }); | |
| // 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 } = await renderTestHook(); | |
| // Start a query | |
| await act(async () => { | |
| // eslint-disable-next-line @typescript-eslint/no-floating-promises | |
| 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.', | |
| }); | |
| }); | |
| // 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 } = await renderHookWithProviders(() => | |
| useGeminiStream( | |
| mockConfig.getGeminiClient(), | |
| [], | |
| mockAddItem, | |
| mockConfig, | |
| mockLoadedSettings, | |
| mockOnDebugMessage, | |
| mockHandleSlashCommand, | |
| false, | |
| () => 'vscode' as EditorType, | |
| () => {}, | |
| () => Promise.resolve(), | |
| false, | |
| () => {}, | |
| cancelSubmitSpy, | |
| () => {}, | |
| 80, | |
| 24, | |
| ), | |
| ); | |
| // Start a query | |
| await act(async () => { | |
| // eslint-disable-next-line @typescript-eslint/no-floating-promises | |
| result.current.submitQuery('test query'); | |
| }); | |
| simulateEscapeKeyPress(); | |
| expect(cancelSubmitSpy).toHaveBeenCalledWith(false, false); | |
| }); | |
| 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 } = await renderHookWithProviders(() => | |
| useGeminiStream( | |
| mockConfig.getGeminiClient(), | |
| [], | |
| mockAddItem, | |
| mockConfig, | |
| mockLoadedSettings, | |
| mockOnDebugMessage, | |
| mockHandleSlashCommand, | |
| false, | |
| () => 'vscode' as EditorType, | |
| () => {}, | |
| () => Promise.resolve(), | |
| false, | |
| () => {}, | |
| vi.fn(), | |
| setShellInputFocusedSpy, // Pass the spy here | |
| 80, | |
| 24, | |
| ), | |
| ); | |
| // Start a query | |
| await act(async () => { | |
| // eslint-disable-next-line @typescript-eslint/no-floating-promises | |
| result.current.submitQuery('test query'); | |
| }); | |
| simulateEscapeKeyPress(); | |
| expect(setShellInputFocusedSpy).toHaveBeenCalledWith(false); | |
| }); | |
| it('should not do anything if escape is pressed when not responding', async () => { | |
| const { result } = await 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.', | |
| }), | |
| ); | |
| }); | |
| 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 } = await renderTestHook(); | |
| await act(async () => { | |
| // eslint-disable-next-line @typescript-eslint/no-floating-promises | |
| result.current.submitQuery('long running query'); | |
| }); | |
| await waitFor(() => { | |
| expect(result.current.streamingState).toBe(StreamingState.Responding); | |
| }); | |
| // Cancel the request | |
| simulateEscapeKeyPress(); | |
| // Allow the stream to continue | |
| await act(async () => { | |
| 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 cancel if a tool call is in progress', async () => { | |
| const toolCalls: TrackedToolCall[] = [ | |
| { | |
| request: { callId: 'call1', name: 'tool1', args: {} }, | |
| status: CoreToolCallStatus.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 { result } = await renderTestHook(toolCalls); | |
| // State is `Responding` because a tool is running | |
| expect(result.current.streamingState).toBe(StreamingState.Responding); | |
| // Try to cancel | |
| simulateEscapeKeyPress(); | |
| // The cancel function should be called | |
| expect(mockCancelAllToolCalls).toHaveBeenCalled(); | |
| }); | |
| it('should transition to Idle state when cancelled while a tool call is in progress and completes', async () => { | |
| const toolCalls: TrackedToolCall[] = [ | |
| { | |
| request: { callId: 'call1', name: 'tool1', args: {} }, | |
| status: CoreToolCallStatus.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 { result } = await renderTestHook(toolCalls); | |
| // State is `Responding` because a tool is running | |
| expect(result.current.streamingState).toBe(StreamingState.Responding); | |
| // Try to cancel | |
| simulateEscapeKeyPress(); | |
| // Trigger the onComplete callback with the cancelled tool call | |
| await act(async () => { | |
| if (capturedOnComplete) { | |
| await capturedOnComplete([ | |
| { | |
| ...toolCalls[0], | |
| status: CoreToolCallStatus.Cancelled, | |
| response: { | |
| callId: 'call1', | |
| responseParts: [], | |
| }, | |
| } as any, | |
| ]); | |
| } | |
| }); | |
| // The final state should be idle because the cancelled tool call was marked as submitted | |
| expect(result.current.streamingState).toBe(StreamingState.Idle); | |
| }); | |
| it('should append cancelled tool responses to history when cancelled while a tool call is in progress and completes with response parts', async () => { | |
| const toolCalls: TrackedToolCall[] = [ | |
| { | |
| request: { callId: 'call1', name: 'tool1', args: {} }, | |
| status: CoreToolCallStatus.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 { result, client } = await renderTestHook(toolCalls); | |
| // State is `Responding` because a tool is running | |
| expect(result.current.streamingState).toBe(StreamingState.Responding); | |
| // Try to cancel | |
| simulateEscapeKeyPress(); | |
| const expectedResponseParts = [ | |
| { | |
| functionResponse: { | |
| name: 'tool1', | |
| id: 'call1', | |
| response: { error: 'cancelled' }, | |
| }, | |
| }, | |
| ]; | |
| // Trigger the onComplete callback with the cancelled tool call having non-empty response parts | |
| await act(async () => { | |
| if (capturedOnComplete) { | |
| await capturedOnComplete([ | |
| { | |
| ...toolCalls[0], | |
| status: CoreToolCallStatus.Cancelled, | |
| response: { | |
| callId: 'call1', | |
| responseParts: expectedResponseParts, | |
| }, | |
| } as any, | |
| ]); | |
| } | |
| }); | |
| // Assert that addHistory was called with the combined response parts | |
| expect(client.addHistory).toHaveBeenCalledWith({ | |
| role: 'user', | |
| parts: expectedResponseParts, | |
| }); | |
| // The final state should be idle because the cancelled tool call was marked as submitted | |
| expect(result.current.streamingState).toBe(StreamingState.Idle); | |
| }); | |
| it('should cancel a request when a tool is awaiting confirmation', async () => { | |
| const mockOnConfirm = vi.fn().mockResolvedValue(undefined); | |
| const toolCalls: TrackedToolCall[] = [ | |
| { | |
| request: { | |
| callId: 'confirm-call', | |
| name: 'some_tool', | |
| args: {}, | |
| isClientInitiated: false, | |
| prompt_id: 'prompt-id-1', | |
| }, | |
| status: CoreToolCallStatus.AwaitingApproval, | |
| responseSubmittedToGemini: false, | |
| tool: { | |
| name: 'some_tool', | |
| description: 'a tool', | |
| build: vi.fn().mockImplementation((_) => ({ | |
| getDescription: () => `Mock description`, | |
| })), | |
| } as any, | |
| invocation: { | |
| getDescription: () => `Mock description`, | |
| } as unknown as AnyToolInvocation, | |
| confirmationDetails: { | |
| type: 'edit', | |
| title: 'Confirm Edit', | |
| onConfirm: mockOnConfirm, | |
| fileName: 'file.txt', | |
| filePath: '/test/file.txt', | |
| fileDiff: 'fake diff', | |
| originalContent: 'old', | |
| newContent: 'new', | |
| }, | |
| } as TrackedWaitingToolCall, | |
| ]; | |
| const { result } = await renderTestHook(toolCalls); | |
| // State is `WaitingForConfirmation` because a tool is awaiting approval | |
| expect(result.current.streamingState).toBe( | |
| StreamingState.WaitingForConfirmation, | |
| ); | |
| // Try to cancel | |
| simulateEscapeKeyPress(); | |
| // The imperative cancel function should be called on the scheduler | |
| expect(mockCancelAllToolCalls).toHaveBeenCalled(); | |
| // A cancellation message should be added to history | |
| await waitFor(() => { | |
| expect(mockAddItem).toHaveBeenCalledWith( | |
| expect.objectContaining({ | |
| text: 'Request cancelled.', | |
| }), | |
| ); | |
| }); | |
| // The final state should be idle | |
| expect(result.current.streamingState).toBe(StreamingState.Idle); | |
| }); | |
| }); | |
| describe('Retry Handling', () => { | |
| it('should ignore retryStatus updates when not responding', async () => { | |
| const { result } = await renderHookWithDefaults(); | |
| const retryPayload = { | |
| model: 'gemini-2.5-pro', | |
| attempt: 2, | |
| maxAttempts: 3, | |
| delayMs: 1000, | |
| }; | |
| await act(async () => { | |
| coreEvents.emit(CoreEvent.RetryAttempt, retryPayload); | |
| }); | |
| expect(result.current.retryStatus).toBeNull(); | |
| }); | |
| it('should reset retryStatus when isResponding becomes false', async () => { | |
| const { result } = await renderTestHook(); | |
| const retryPayload = { | |
| model: 'gemini-2.5-pro', | |
| attempt: 2, | |
| maxAttempts: 3, | |
| delayMs: 1000, | |
| }; | |
| // Start a query to make isResponding true | |
| const mockStream = (async function* () { | |
| yield { type: ServerGeminiEventType.Content, value: 'Part 1' }; | |
| await new Promise(() => {}); // Keep stream open | |
| })(); | |
| mockSendMessageStream.mockReturnValue(mockStream); | |
| await act(async () => { | |
| // eslint-disable-next-line @typescript-eslint/no-floating-promises | |
| result.current.submitQuery('test query'); | |
| }); | |
| await waitFor(() => { | |
| expect(result.current.streamingState).toBe(StreamingState.Responding); | |
| }); | |
| // Emit retry event | |
| await act(async () => { | |
| coreEvents.emit(CoreEvent.RetryAttempt, retryPayload); | |
| }); | |
| expect(result.current.retryStatus).toEqual(retryPayload); | |
| // Cancel to make isResponding false | |
| await act(async () => { | |
| result.current.cancelOngoingRequest(); | |
| }); | |
| expect(result.current.retryStatus).toBeNull(); | |
| }); | |
| it('should ignore late retry events after cancellation', async () => { | |
| const { result } = await renderTestHook(); | |
| const retryPayload = { | |
| model: 'gemini-2.5-pro', | |
| attempt: 2, | |
| maxAttempts: 3, | |
| delayMs: 1000, | |
| }; | |
| const lateRetryPayload = { | |
| model: 'gemini-2.5-pro', | |
| attempt: 3, | |
| maxAttempts: 3, | |
| delayMs: 2000, | |
| }; | |
| const mockStream = (async function* () { | |
| yield { type: ServerGeminiEventType.Content, value: 'Part 1' }; | |
| await new Promise(() => {}); // Keep stream open | |
| })(); | |
| mockSendMessageStream.mockReturnValue(mockStream); | |
| await act(async () => { | |
| // eslint-disable-next-line @typescript-eslint/no-floating-promises | |
| result.current.submitQuery('test query'); | |
| }); | |
| await waitFor(() => { | |
| expect(result.current.streamingState).toBe(StreamingState.Responding); | |
| }); | |
| await act(async () => { | |
| coreEvents.emit(CoreEvent.RetryAttempt, retryPayload); | |
| }); | |
| expect(result.current.retryStatus).toEqual(retryPayload); | |
| await act(async () => { | |
| result.current.cancelOngoingRequest(); | |
| }); | |
| await waitFor(() => { | |
| expect(result.current.retryStatus).toBeNull(); | |
| }); | |
| await act(async () => { | |
| coreEvents.emit(CoreEvent.RetryAttempt, lateRetryPayload); | |
| }); | |
| expect(result.current.retryStatus).toBeNull(); | |
| }); | |
| }); | |
| 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: 'activate_skill', | |
| toolArgs: { name: 'test-skill' }, | |
| }; | |
| mockHandleSlashCommand.mockResolvedValue(clientToolRequest); | |
| const { result } = await renderTestHook(); | |
| await act(async () => { | |
| await result.current.submitQuery('/memory show'); | |
| }); | |
| await waitFor(() => { | |
| expect(mockScheduleToolCalls).toHaveBeenCalledWith( | |
| [ | |
| expect.objectContaining({ | |
| name: 'activate_skill', | |
| args: { name: 'test-skill' }, | |
| 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 } = await 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 } = | |
| await 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, | |
| '/my-custom-command', | |
| ); | |
| 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 } = | |
| await 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, | |
| '/emptycmd', | |
| ); | |
| }); | |
| }); | |
| it('should not call handleSlashCommand for line comments', async () => { | |
| const { result, mockSendMessageStream: localMockSendMessageStream } = | |
| await 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, | |
| '// This is a line comment', | |
| ); | |
| }); | |
| }); | |
| it('should not call handleSlashCommand for block comments', async () => { | |
| const { result, mockSendMessageStream: localMockSendMessageStream } = | |
| await 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, | |
| '/* This is a block comment */', | |
| ); | |
| }); | |
| }); | |
| it('should not call handleSlashCommand is shell mode is active', async () => { | |
| const { result } = await renderHookWithProviders(() => | |
| useGeminiStream( | |
| new MockedGeminiClientClass(mockConfig), | |
| [], | |
| mockAddItem, | |
| mockConfig, | |
| mockLoadedSettings, | |
| () => {}, | |
| mockHandleSlashCommand, | |
| true, | |
| () => 'vscode' as EditorType, | |
| () => {}, | |
| () => Promise.resolve(), | |
| false, | |
| () => {}, | |
| () => {}, | |
| () => {}, | |
| 80, | |
| 24, | |
| ), | |
| ); | |
| await act(async () => { | |
| await result.current.submitQuery('/about'); | |
| }); | |
| await waitFor(() => { | |
| expect(mockHandleSlashCommand).not.toHaveBeenCalled(); | |
| }); | |
| }); | |
| it('should record client-initiated tool calls in GeminiChat history', async () => { | |
| const { result, client: mockGeminiClient } = await renderTestHook(); | |
| mockHandleSlashCommand.mockResolvedValue({ | |
| type: 'schedule_tool', | |
| toolName: 'activate_skill', | |
| toolArgs: { name: 'test-skill' }, | |
| }); | |
| await act(async () => { | |
| await result.current.submitQuery('/test-skill'); | |
| }); | |
| // Simulate tool completion | |
| const completedTool = { | |
| request: { | |
| callId: 'test-call-id', | |
| name: 'activate_skill', | |
| args: { name: 'test-skill' }, | |
| isClientInitiated: true, | |
| }, | |
| status: CoreToolCallStatus.Success, | |
| invocation: { | |
| getDescription: () => 'Activating skill test-skill', | |
| }, | |
| tool: { | |
| isOutputMarkdown: true, | |
| }, | |
| response: { | |
| responseParts: [ | |
| { | |
| functionResponse: { | |
| name: 'activate_skill', | |
| response: { content: 'skill instructions' }, | |
| }, | |
| }, | |
| ], | |
| }, | |
| } as unknown as TrackedCompletedToolCall; | |
| await act(async () => { | |
| if (capturedOnComplete) { | |
| await capturedOnComplete([completedTool]); | |
| } | |
| }); | |
| // Verify that the tool call and response were added to GeminiChat history | |
| expect(mockGeminiClient.addHistory).toHaveBeenCalledWith({ | |
| role: 'model', | |
| parts: [ | |
| { | |
| functionCall: { | |
| name: 'activate_skill', | |
| args: { name: 'test-skill' }, | |
| }, | |
| }, | |
| ], | |
| }); | |
| expect(mockGeminiClient.addHistory).toHaveBeenCalledWith({ | |
| role: 'user', | |
| parts: completedTool.response.responseParts, | |
| }); | |
| }); | |
| it('should NOT record other client-initiated tool calls in history', async () => { | |
| const { result, client: mockGeminiClient } = await renderTestHook(); | |
| mockHandleSlashCommand.mockResolvedValue({ | |
| type: 'schedule_tool', | |
| toolName: 'write_todos', | |
| toolArgs: { todos: [] }, | |
| }); | |
| await act(async () => { | |
| await result.current.submitQuery('/todos'); | |
| }); | |
| // Simulate tool completion | |
| const completedTool = { | |
| request: { | |
| callId: 'test-call-id', | |
| name: 'write_todos', | |
| args: { todos: [] }, | |
| isClientInitiated: true, | |
| }, | |
| status: CoreToolCallStatus.Success, | |
| invocation: { | |
| getDescription: () => 'Saving memory', | |
| }, | |
| tool: { | |
| isOutputMarkdown: true, | |
| }, | |
| response: { | |
| responseParts: [ | |
| { | |
| functionResponse: { | |
| name: 'write_todos', | |
| response: { success: true }, | |
| }, | |
| }, | |
| ], | |
| }, | |
| } as unknown as TrackedCompletedToolCall; | |
| await act(async () => { | |
| if (capturedOnComplete) { | |
| await capturedOnComplete([completedTool]); | |
| } | |
| }); | |
| // Verify that addHistory was NOT called | |
| expect(mockGeminiClient.addHistory).not.toHaveBeenCalled(); | |
| }); | |
| }); | |
| 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.LOGIN_WITH_GOOGLE; | |
| mockParseAndFormatApiError.mockClear(); | |
| mockSendMessageStream.mockReturnValue( | |
| (async function* () { | |
| yield { type: 'content', value: '' }; | |
| throw mockError; | |
| })(), | |
| ); | |
| const testConfig = { | |
| // eslint-disable-next-line @typescript-eslint/no-misused-spread | |
| ...mockConfig, | |
| getContentGenerator: vi.fn(), | |
| getContentGeneratorConfig: vi.fn(() => ({ | |
| authType: mockAuthType, | |
| })), | |
| getModel: vi.fn(() => 'gemini-2.5-pro'), | |
| } as unknown as Config; | |
| const { result } = await renderHookWithProviders(() => | |
| useGeminiStream( | |
| new MockedGeminiClientClass(testConfig), | |
| [], | |
| mockAddItem, | |
| testConfig, | |
| mockLoadedSettings, | |
| mockOnDebugMessage, | |
| mockHandleSlashCommand, | |
| false, | |
| () => 'vscode' as EditorType, | |
| () => {}, | |
| () => Promise.resolve(), | |
| false, | |
| () => {}, | |
| () => {}, | |
| () => {}, | |
| 80, | |
| 24, | |
| ), | |
| ); | |
| // 2. Action | |
| await act(async () => { | |
| await result.current.submitQuery('test query'); | |
| }); | |
| // 3. Assertion | |
| await waitFor(() => { | |
| expect(mockParseAndFormatApiError).toHaveBeenCalledWith( | |
| 'Rate limit exceeded', | |
| mockAuthType, | |
| undefined, | |
| 'gemini-2.5-pro', | |
| 'gemini-2.5-flash', | |
| ); | |
| }); | |
| }); | |
| it('should use TRUE_EMPTY_RESPONSE_MESSAGE when receiving an invalid stream event of type NO_RESPONSE_TEXT', async () => { | |
| mockSendMessageStream.mockClear(); | |
| mockSendMessageStream.mockReturnValue( | |
| (async function* () { | |
| yield { | |
| type: ServerGeminiEventType.InvalidStream, | |
| value: { | |
| type: 'NO_RESPONSE_TEXT', | |
| message: 'empty response text', | |
| }, | |
| }; | |
| })(), | |
| ); | |
| const { result } = await renderTestHook(); | |
| await act(async () => { | |
| await result.current.submitQuery('test query'); | |
| }); | |
| await waitFor(() => { | |
| expect(mockAddItem).toHaveBeenCalledWith( | |
| expect.objectContaining({ | |
| type: MessageType.ERROR, | |
| text: TRUE_EMPTY_RESPONSE_MESSAGE, | |
| }), | |
| expect.any(Number), | |
| ); | |
| }); | |
| }); | |
| it('should use the event message when receiving a non-NO_RESPONSE_TEXT invalid stream event', async () => { | |
| mockSendMessageStream.mockClear(); | |
| mockSendMessageStream.mockReturnValue( | |
| (async function* () { | |
| yield { | |
| type: ServerGeminiEventType.InvalidStream, | |
| value: { | |
| type: 'MALFORMED_FUNCTION_CALL', | |
| message: 'Custom malformed function call message', | |
| }, | |
| }; | |
| })(), | |
| ); | |
| const { result } = await renderTestHook(); | |
| await act(async () => { | |
| await result.current.submitQuery('test query'); | |
| }); | |
| await waitFor(() => { | |
| expect(mockAddItem).toHaveBeenCalledWith( | |
| expect.objectContaining({ | |
| type: MessageType.ERROR, | |
| text: 'Custom malformed function call message', | |
| }), | |
| expect.any(Number), | |
| ); | |
| }); | |
| }); | |
| }); | |
| describe('handleApprovalModeChange', () => { | |
| it('should auto-approve all pending tool calls when switching to YOLO mode', async () => { | |
| const awaitingApprovalToolCalls: TrackedToolCall[] = [ | |
| createMockToolCall('replace', 'call1', 'edit'), | |
| createMockToolCall('read_file', 'call2', 'info'), | |
| ]; | |
| const { result } = await renderTestHook(awaitingApprovalToolCalls); | |
| await act(async () => { | |
| await result.current.handleApprovalModeChange(ApprovalMode.YOLO); | |
| }); | |
| // Both tool calls should be auto-approved | |
| expect(mockMessageBus.publish).toHaveBeenCalledTimes(2); | |
| expect(mockMessageBus.publish).toHaveBeenCalledWith( | |
| expect.objectContaining({ | |
| type: MessageBusType.TOOL_CONFIRMATION_RESPONSE, | |
| correlationId: 'corr-call1', | |
| outcome: ToolConfirmationOutcome.ProceedOnce, | |
| }), | |
| ); | |
| expect(mockMessageBus.publish).toHaveBeenCalledWith( | |
| expect.objectContaining({ | |
| correlationId: 'corr-call2', | |
| outcome: ToolConfirmationOutcome.ProceedOnce, | |
| }), | |
| ); | |
| }); | |
| it('should only auto-approve edit tools when switching to AUTO_EDIT mode', async () => { | |
| const awaitingApprovalToolCalls: TrackedToolCall[] = [ | |
| createMockToolCall('replace', 'call1', 'edit'), | |
| createMockToolCall('write_file', 'call2', 'edit'), | |
| createMockToolCall('read_file', 'call3', 'info'), | |
| ]; | |
| const { result } = await renderTestHook(awaitingApprovalToolCalls); | |
| await act(async () => { | |
| await result.current.handleApprovalModeChange(ApprovalMode.AUTO_EDIT); | |
| }); | |
| // Only replace and write_file should be auto-approved | |
| expect(mockMessageBus.publish).toHaveBeenCalledTimes(2); | |
| expect(mockMessageBus.publish).toHaveBeenCalledWith( | |
| expect.objectContaining({ correlationId: 'corr-call1' }), | |
| ); | |
| expect(mockMessageBus.publish).toHaveBeenCalledWith( | |
| expect.objectContaining({ correlationId: 'corr-call2' }), | |
| ); | |
| expect(mockMessageBus.publish).not.toHaveBeenCalledWith( | |
| expect.objectContaining({ correlationId: 'corr-call3' }), | |
| ); | |
| }); | |
| it('should auto-approve shell commands with redirection when switching to AUTO_EDIT mode', async () => { | |
| const shellCall = createMockToolCall( | |
| SHELL_TOOL_NAME, | |
| 'call-shell', | |
| 'info', | |
| ); | |
| shellCall.request.args = { command: 'ls > files.txt' }; | |
| const { result } = await renderTestHook([shellCall]); | |
| await act(async () => { | |
| await result.current.handleApprovalModeChange(ApprovalMode.AUTO_EDIT); | |
| }); | |
| // Shell command with redirection should be auto-approved | |
| expect(mockMessageBus.publish).toHaveBeenCalledWith( | |
| expect.objectContaining({ correlationId: 'corr-call-shell' }), | |
| ); | |
| }); | |
| it('should NOT auto-approve shell commands without redirection when switching to AUTO_EDIT mode', async () => { | |
| const shellCall = createMockToolCall( | |
| SHELL_TOOL_NAME, | |
| 'call-shell', | |
| 'info', | |
| ); | |
| shellCall.request.args = { command: 'ls -la' }; | |
| const { result } = await renderTestHook([shellCall]); | |
| await act(async () => { | |
| await result.current.handleApprovalModeChange(ApprovalMode.AUTO_EDIT); | |
| }); | |
| // Regular shell command should NOT be auto-approved | |
| expect(mockMessageBus.publish).not.toHaveBeenCalled(); | |
| }); | |
| it('should not auto-approve any tools when switching to REQUIRE_CONFIRMATION mode', async () => { | |
| const awaitingApprovalToolCalls: TrackedToolCall[] = [ | |
| createMockToolCall('replace', 'call1', 'edit'), | |
| ]; | |
| const { result } = await renderTestHook(awaitingApprovalToolCalls); | |
| await act(async () => { | |
| await result.current.handleApprovalModeChange(ApprovalMode.DEFAULT); | |
| }); | |
| // No tools should be auto-approved | |
| expect(mockMessageBus.publish).not.toHaveBeenCalled(); | |
| }); | |
| it('should handle errors gracefully when auto-approving tool calls', async () => { | |
| const debuggerSpy = vi | |
| .spyOn(debugLogger, 'warn') | |
| .mockImplementation(() => {}); | |
| mockMessageBus.publish.mockRejectedValueOnce(new Error('Bus error')); | |
| const awaitingApprovalToolCalls: TrackedToolCall[] = [ | |
| createMockToolCall('replace', 'call1', 'edit'), | |
| createMockToolCall('write_file', 'call2', 'edit'), | |
| ]; | |
| const { result } = await renderTestHook(awaitingApprovalToolCalls); | |
| await act(async () => { | |
| await result.current.handleApprovalModeChange(ApprovalMode.YOLO); | |
| }); | |
| // Both should be attempted despite first error | |
| expect(mockMessageBus.publish).toHaveBeenCalledTimes(2); | |
| expect(debuggerSpy).toHaveBeenCalledWith( | |
| 'Failed to auto-approve tool call call1:', | |
| expect.any(Error), | |
| ); | |
| debuggerSpy.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: CoreToolCallStatus.AwaitingApproval, | |
| responseSubmittedToGemini: false, | |
| // No confirmationDetails | |
| tool: { | |
| name: 'replace', | |
| displayName: 'replace', | |
| description: 'Replace text', | |
| build: vi.fn(), | |
| } as unknown as AnyDeclarativeTool, | |
| invocation: { | |
| getDescription: () => 'Mock description', | |
| } as unknown as AnyToolInvocation, | |
| correlationId: 'corr-1', | |
| } as unknown as TrackedWaitingToolCall, | |
| ]; | |
| const { result } = await 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 mixedStatusToolCalls: TrackedToolCall[] = [ | |
| createMockToolCall( | |
| 'replace', | |
| 'call1', | |
| 'edit', | |
| CoreToolCallStatus.AwaitingApproval, | |
| mockOnConfirmAwaiting, | |
| ), | |
| { | |
| request: { | |
| callId: 'call2', | |
| name: 'write_file', | |
| args: { path: '/test/file.txt', content: 'content' }, | |
| isClientInitiated: false, | |
| prompt_id: 'prompt-id-1', | |
| }, | |
| status: CoreToolCallStatus.Executing, | |
| responseSubmittedToGemini: false, | |
| tool: { | |
| name: 'write_file', | |
| displayName: 'write_file', | |
| description: 'Write file', | |
| build: vi.fn(), | |
| } as unknown as AnyDeclarativeTool, | |
| invocation: { | |
| getDescription: () => 'Mock description', | |
| } as unknown as AnyToolInvocation, | |
| startTime: Date.now(), | |
| liveOutput: 'Writing...', | |
| correlationId: 'corr-call2', | |
| } as TrackedExecutingToolCall, | |
| ]; | |
| const { result } = await renderTestHook(mixedStatusToolCalls); | |
| await act(async () => { | |
| await result.current.handleApprovalModeChange(ApprovalMode.YOLO); | |
| }); | |
| // Only the awaiting_approval tool should be processed. | |
| expect(mockMessageBus.publish).toHaveBeenCalledTimes(1); | |
| expect(mockMessageBus.publish).toHaveBeenCalledWith( | |
| expect.objectContaining({ correlationId: 'corr-call1' }), | |
| ); | |
| expect(mockMessageBus.publish).not.toHaveBeenCalledWith( | |
| expect.objectContaining({ correlationId: 'corr-call2' }), | |
| ); | |
| }); | |
| it('should inject a notification message when manually exiting Plan Mode', async () => { | |
| // Setup mockConfig to return PLAN mode initially | |
| (mockConfig.getApprovalMode as Mock).mockReturnValue(ApprovalMode.PLAN); | |
| // Render the hook, which will initialize the previousApprovalModeRef with PLAN | |
| const { result, client } = await renderTestHook([]); | |
| // Update mockConfig to return DEFAULT mode (new mode) | |
| (mockConfig.getApprovalMode as Mock).mockReturnValue( | |
| ApprovalMode.DEFAULT, | |
| ); | |
| await act(async () => { | |
| // Trigger manual exit from Plan Mode | |
| await result.current.handleApprovalModeChange(ApprovalMode.DEFAULT); | |
| }); | |
| // Verify that addHistory was called with the notification message | |
| expect(client.addHistory).toHaveBeenCalledWith({ | |
| role: 'user', | |
| parts: [ | |
| { | |
| text: getPlanModeExitMessage(ApprovalMode.DEFAULT, true), | |
| }, | |
| ], | |
| }); | |
| }); | |
| }); | |
| 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 } = await renderHookWithProviders(() => | |
| useGeminiStream( | |
| new MockedGeminiClientClass(mockConfig), | |
| [], | |
| mockAddItem, | |
| mockConfig, | |
| mockLoadedSettings, | |
| mockOnDebugMessage, | |
| mockHandleSlashCommand, | |
| false, | |
| () => 'vscode' as EditorType, | |
| () => {}, | |
| () => Promise.resolve(), | |
| false, | |
| () => {}, | |
| () => {}, | |
| () => {}, | |
| 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), | |
| ); | |
| }); | |
| }); | |
| describe('ContextWindowWillOverflow event', () => { | |
| beforeEach(() => { | |
| vi.mocked(tokenLimit).mockReturnValue(100); | |
| }); | |
| it.each([ | |
| { | |
| name: 'without suggestion when remaining tokens are > 75% of limit', | |
| requestTokens: 20, | |
| remainingTokens: 80, | |
| expectedMessage: | |
| 'Sending this message (20 tokens) might exceed the context window limit (80 tokens left).', | |
| }, | |
| { | |
| name: 'with suggestion when remaining tokens are < 75% of limit', | |
| requestTokens: 30, | |
| remainingTokens: 70, | |
| expectedMessage: | |
| 'Sending this message (30 tokens) might exceed the context window limit (70 tokens left). Please try reducing the size of your message or use the `/compress` command to compress the chat history.', | |
| }, | |
| ])( | |
| 'should add message $name', | |
| async ({ requestTokens, remainingTokens, expectedMessage }) => { | |
| mockSendMessageStream.mockReturnValue( | |
| (async function* () { | |
| yield { | |
| type: ServerGeminiEventType.ContextWindowWillOverflow, | |
| value: { | |
| estimatedRequestTokenCount: requestTokens, | |
| remainingTokenCount: remainingTokens, | |
| }, | |
| }; | |
| })(), | |
| ); | |
| const { result } = await renderHookWithDefaults(); | |
| await act(async () => { | |
| await result.current.submitQuery('Test overflow'); | |
| }); | |
| await waitFor(() => { | |
| expect(mockAddItem).toHaveBeenCalledWith({ | |
| type: 'info', | |
| text: expectedMessage, | |
| }); | |
| }); | |
| }, | |
| ); | |
| }); | |
| it('should call onCancelSubmit when ContextWindowWillOverflow event is received', async () => { | |
| const onCancelSubmitSpy = vi.fn(); | |
| // Setup mock to return a stream with ContextWindowWillOverflow event | |
| mockSendMessageStream.mockReturnValue( | |
| (async function* () { | |
| yield { | |
| type: ServerGeminiEventType.ContextWindowWillOverflow, | |
| value: { | |
| estimatedRequestTokenCount: 100, | |
| remainingTokenCount: 50, | |
| }, | |
| }; | |
| })(), | |
| ); | |
| const { result } = await renderHookWithProviders(() => | |
| useGeminiStream( | |
| new MockedGeminiClientClass(mockConfig), | |
| [], | |
| mockAddItem, | |
| mockConfig, | |
| mockLoadedSettings, | |
| mockOnDebugMessage, | |
| mockHandleSlashCommand, | |
| false, | |
| () => 'vscode' as EditorType, | |
| () => {}, | |
| () => Promise.resolve(), | |
| false, | |
| () => {}, | |
| onCancelSubmitSpy, | |
| () => {}, | |
| 80, | |
| 24, | |
| ), | |
| ); | |
| // Submit a query | |
| await act(async () => { | |
| await result.current.submitQuery('Test overflow'); | |
| }); | |
| // Check that onCancelSubmit was called | |
| await waitFor(() => { | |
| expect(onCancelSubmitSpy).toHaveBeenCalledWith(true); | |
| }); | |
| }); | |
| it('should add informational messages when ChatCompressed event is received', async () => { | |
| vi.mocked(tokenLimit).mockReturnValue(10000); | |
| // Setup mock to return a stream with ChatCompressed event | |
| mockSendMessageStream.mockReturnValue( | |
| (async function* () { | |
| yield { | |
| type: ServerGeminiEventType.ChatCompressed, | |
| value: { | |
| originalTokenCount: 1000, | |
| newTokenCount: 500, | |
| compressionStatus: 'compressed', | |
| }, | |
| }; | |
| })(), | |
| ); | |
| const { result } = await renderHookWithDefaults(); | |
| // Submit a query | |
| await act(async () => { | |
| await result.current.submitQuery('Test compression'); | |
| }); | |
| // Check that the succinct info message was added | |
| await waitFor(() => { | |
| expect(mockAddItem).toHaveBeenCalledWith( | |
| expect.objectContaining({ | |
| type: MessageType.INFO, | |
| text: 'Context compressed from 10% to 5%.', | |
| secondaryText: 'Change threshold in /settings.', | |
| color: theme.status.warning, | |
| }), | |
| expect.any(Number), | |
| ); | |
| }); | |
| }); | |
| it.each([ | |
| { | |
| reason: 'STOP', | |
| shouldAddMessage: false, | |
| }, | |
| { | |
| reason: 'FINISH_REASON_UNSPECIFIED', | |
| shouldAddMessage: false, | |
| }, | |
| { | |
| reason: 'SAFETY', | |
| message: '⚠️ Response stopped due to safety reasons.', | |
| }, | |
| { | |
| reason: 'RECITATION', | |
| message: '⚠️ Response stopped due to recitation policy.', | |
| }, | |
| { | |
| reason: 'LANGUAGE', | |
| message: '⚠️ Response stopped due to unsupported language.', | |
| }, | |
| { | |
| reason: 'BLOCKLIST', | |
| message: '⚠️ Response stopped due to forbidden terms.', | |
| }, | |
| { | |
| reason: 'PROHIBITED_CONTENT', | |
| message: '⚠️ Response stopped due to prohibited content.', | |
| }, | |
| { | |
| reason: 'SPII', | |
| message: | |
| '⚠️ Response stopped due to sensitive personally identifiable information.', | |
| }, | |
| { | |
| reason: 'OTHER', | |
| message: '⚠️ Response stopped for other reasons.', | |
| }, | |
| { | |
| reason: 'MALFORMED_FUNCTION_CALL', | |
| message: '⚠️ Response stopped due to malformed function call.', | |
| }, | |
| { | |
| reason: 'IMAGE_SAFETY', | |
| message: '⚠️ Response stopped due to image safety violations.', | |
| }, | |
| { | |
| reason: 'UNEXPECTED_TOOL_CALL', | |
| message: '⚠️ Response stopped due to unexpected tool call.', | |
| }, | |
| ])( | |
| 'should handle $reason finish reason correctly', | |
| async ({ reason, shouldAddMessage = true, message }) => { | |
| mockSendMessageStream.mockReturnValue( | |
| (async function* () { | |
| yield { | |
| type: ServerGeminiEventType.Content, | |
| value: `Response for ${reason}`, | |
| }; | |
| yield { | |
| type: ServerGeminiEventType.Finished, | |
| value: { reason, usageMetadata: undefined }, | |
| }; | |
| })(), | |
| ); | |
| const { result } = await renderHookWithDefaults(); | |
| await act(async () => { | |
| await result.current.submitQuery(`Test ${reason}`); | |
| }); | |
| if (shouldAddMessage) { | |
| await waitFor(() => { | |
| expect(mockAddItem).toHaveBeenCalledWith( | |
| { | |
| type: 'info', | |
| text: message, | |
| }, | |
| expect.any(Number), | |
| ); | |
| }); | |
| } else { | |
| // Verify state returns to idle without any info messages | |
| await waitFor(() => { | |
| expect(result.current.streamingState).toBe(StreamingState.Idle); | |
| }); | |
| const infoMessages = mockAddItem.mock.calls.filter( | |
| (call) => call[0].type === 'info', | |
| ); | |
| expect(infoMessages).toHaveLength(0); | |
| } | |
| }, | |
| ); | |
| }); | |
| it('should flush pending text rationale before scheduling tool calls to ensure correct history order', async () => { | |
| const addItemOrder: string[] = []; | |
| let capturedOnComplete: (tools: CompletedToolCall[]) => Promise<void>; | |
| const mockScheduleToolCalls = vi.fn(async (requests) => { | |
| addItemOrder.push('scheduleToolCalls_START'); | |
| // Simulate tools completing and triggering onComplete immediately. | |
| // This mimics the behavior that caused the regression where tool results | |
| // were added to history during the await scheduleToolCalls(...) block. | |
| const tools = requests.map((r: ToolCallRequestInfo) => ({ | |
| request: r, | |
| status: CoreToolCallStatus.Success, | |
| tool: { displayName: r.name, name: r.name }, | |
| invocation: { getDescription: () => 'desc' }, | |
| response: { responseParts: [], resultDisplay: 'done' }, | |
| startTime: Date.now(), | |
| endTime: Date.now(), | |
| })); | |
| // Wait a tick for refs to be set up | |
| await new Promise((resolve) => setTimeout(resolve, 0)); | |
| await capturedOnComplete(tools); | |
| addItemOrder.push('scheduleToolCalls_END'); | |
| }); | |
| mockAddItem.mockImplementation((item: HistoryItemWithoutId) => { | |
| addItemOrder.push(`addItem:${item.type}`); | |
| }); | |
| // We need to capture the onComplete callback from useToolScheduler | |
| mockUseToolScheduler.mockImplementation((onComplete) => { | |
| capturedOnComplete = onComplete; | |
| return [ | |
| [], // toolCalls | |
| mockScheduleToolCalls, | |
| vi.fn(), // markToolsAsSubmitted | |
| vi.fn(), // setToolCallsForDisplay | |
| vi.fn(), // cancelAllToolCalls | |
| 0, // lastToolOutputTime | |
| ]; | |
| }); | |
| const { result } = await renderHookWithProviders(() => | |
| useGeminiStream( | |
| new MockedGeminiClientClass(mockConfig), | |
| [], | |
| mockAddItem, | |
| mockConfig, | |
| mockLoadedSettings, | |
| vi.fn(), | |
| vi.fn(), | |
| false, | |
| () => 'vscode' as EditorType, | |
| vi.fn(), | |
| vi.fn(), | |
| false, | |
| vi.fn(), | |
| vi.fn(), | |
| vi.fn(), | |
| 80, | |
| 24, | |
| ), | |
| ); | |
| const mockStream = (async function* () { | |
| yield { | |
| type: ServerGeminiEventType.Content, | |
| value: 'Rationale rationale.', | |
| }; | |
| yield { | |
| type: ServerGeminiEventType.ToolCallRequest, | |
| value: { callId: '1', name: 'test_tool', args: {} }, | |
| }; | |
| })(); | |
| mockSendMessageStream.mockReturnValue(mockStream); | |
| await act(async () => { | |
| await result.current.submitQuery('test input'); | |
| }); | |
| // Expectation: addItem:gemini (rationale) MUST happen before scheduleToolCalls_START | |
| const rationaleIndex = addItemOrder.indexOf('addItem:gemini'); | |
| const scheduleIndex = addItemOrder.indexOf('scheduleToolCalls_START'); | |
| const toolGroupIndex = addItemOrder.indexOf('addItem:tool_group'); | |
| expect(rationaleIndex).toBeGreaterThan(-1); | |
| expect(scheduleIndex).toBeGreaterThan(-1); | |
| expect(toolGroupIndex).toBeGreaterThan(-1); | |
| // This is the core fix validation: Rationale comes before tools are even scheduled (awaited) | |
| expect(rationaleIndex).toBeLessThan(scheduleIndex); | |
| expect(rationaleIndex).toBeLessThan(toolGroupIndex); | |
| // Ensure all state updates from recursive submitQuery are settled | |
| await waitFor(() => { | |
| expect(result.current.streamingState).toBe(StreamingState.Idle); | |
| }); | |
| }); | |
| 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 } = await renderHookWithProviders(() => | |
| useGeminiStream( | |
| mockConfig.getGeminiClient(), | |
| [], | |
| mockAddItem, | |
| mockConfig, | |
| mockLoadedSettings, | |
| mockOnDebugMessage, | |
| mockHandleSlashCommand, | |
| false, // shellModeActive | |
| vi.fn(), // getPreferredEditor | |
| vi.fn(), // onAuthError | |
| vi.fn(), // performMemoryRefresh | |
| false, // modelSwitched | |
| vi.fn(), // setModelSwitched | |
| vi.fn(), // onCancelSubmit | |
| 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, | |
| rawQuery, | |
| ); | |
| }); | |
| it('should display user query, then tool execution, then model response', async () => { | |
| const userQuery = 'read this @file(test.txt)'; | |
| const toolExecutionMessage = 'Reading file: test.txt'; | |
| const modelResponseContent = 'The content of test.txt is: Hello World!'; | |
| // Mock handleAtCommand to simulate a tool call and add a tool_group message | |
| handleAtCommandSpy.mockImplementation( | |
| async ({ addItem: atCommandAddItem, messageId }) => { | |
| atCommandAddItem( | |
| { | |
| type: 'tool_group', | |
| tools: [ | |
| { | |
| callId: 'client-read-123', | |
| name: 'read_file', | |
| description: toolExecutionMessage, | |
| status: CoreToolCallStatus.Success, | |
| resultDisplay: toolExecutionMessage, | |
| confirmationDetails: undefined, | |
| }, | |
| ], | |
| }, | |
| messageId, | |
| ); | |
| return { shouldProceed: true, processedQuery: userQuery }; | |
| }, | |
| ); | |
| // Mock the Gemini stream to return a model response after the tool | |
| mockSendMessageStream.mockReturnValue( | |
| (async function* () { | |
| yield { | |
| type: ServerGeminiEventType.Content, | |
| value: modelResponseContent, | |
| }; | |
| yield { | |
| type: ServerGeminiEventType.Finished, | |
| value: { reason: 'STOP' }, | |
| }; | |
| })(), | |
| ); | |
| const { result } = await renderTestHook(); | |
| await act(async () => { | |
| await result.current.submitQuery(userQuery); | |
| }); | |
| // Assert the order of messages added to the history | |
| await waitFor(() => { | |
| expect(mockAddItem).toHaveBeenCalledTimes(3); // User prompt + tool execution + model response | |
| // 1. User's prompt | |
| expect(mockAddItem).toHaveBeenNthCalledWith( | |
| 1, | |
| expect.objectContaining({ | |
| type: MessageType.USER, | |
| text: userQuery, | |
| }), | |
| expect.any(Number), | |
| ); | |
| // 2. Tool execution message | |
| expect(mockAddItem).toHaveBeenNthCalledWith( | |
| 2, | |
| expect.objectContaining({ | |
| type: 'tool_group', | |
| tools: expect.arrayContaining([ | |
| expect.objectContaining({ | |
| name: 'read_file', | |
| status: CoreToolCallStatus.Success, | |
| }), | |
| ]), | |
| }), | |
| expect.any(Number), | |
| ); | |
| // 3. Model's response | |
| expect(mockAddItem).toHaveBeenNthCalledWith( | |
| 3, | |
| expect.objectContaining({ | |
| type: 'gemini', | |
| text: modelResponseContent, | |
| }), | |
| expect.any(Number), | |
| ); | |
| }); | |
| }); | |
| describe('Thought Reset', () => { | |
| it('should keep full thinking entries in history when mode is full', async () => { | |
| const fullThinkingSettings: LoadedSettings = { | |
| // eslint-disable-next-line @typescript-eslint/no-misused-spread | |
| ...mockLoadedSettings, | |
| merged: { | |
| ...mockLoadedSettings.merged, | |
| ui: { inlineThinkingMode: 'full' }, | |
| }, | |
| } as unknown as LoadedSettings; | |
| mockSendMessageStream.mockReturnValue( | |
| (async function* () { | |
| yield { | |
| type: ServerGeminiEventType.Thought, | |
| value: { | |
| subject: 'Full thought', | |
| description: 'Detailed thinking', | |
| }, | |
| }; | |
| yield { | |
| type: ServerGeminiEventType.Content, | |
| value: 'Response', | |
| }; | |
| })(), | |
| ); | |
| const { result } = await renderHookWithProviders(() => | |
| useGeminiStream( | |
| new MockedGeminiClientClass(mockConfig), | |
| [], | |
| mockAddItem, | |
| mockConfig, | |
| fullThinkingSettings, | |
| mockOnDebugMessage, | |
| mockHandleSlashCommand, | |
| false, | |
| () => 'vscode' as EditorType, | |
| () => {}, | |
| () => Promise.resolve(), | |
| false, | |
| () => {}, | |
| () => {}, | |
| () => {}, | |
| 80, | |
| 24, | |
| ), | |
| ); | |
| await act(async () => { | |
| await result.current.submitQuery('Test query'); | |
| }); | |
| expect(mockAddItem).toHaveBeenCalledWith( | |
| expect.objectContaining({ | |
| type: 'thinking', | |
| thought: expect.objectContaining({ subject: 'Full thought' }), | |
| }), | |
| ); | |
| }); | |
| it('keeps thought transient and clears it on first non-thought event', async () => { | |
| mockSendMessageStream.mockReturnValue( | |
| (async function* () { | |
| yield { | |
| type: ServerGeminiEventType.Thought, | |
| value: { | |
| subject: 'Assessing intent', | |
| description: 'Inspecting context', | |
| }, | |
| }; | |
| yield { | |
| type: ServerGeminiEventType.Content, | |
| value: 'Model response content', | |
| }; | |
| yield { | |
| type: ServerGeminiEventType.Finished, | |
| value: { reason: 'STOP', usageMetadata: undefined }, | |
| }; | |
| })(), | |
| ); | |
| const { result } = await renderTestHook(); | |
| await act(async () => { | |
| await result.current.submitQuery('Test query'); | |
| }); | |
| await waitFor(() => { | |
| expect(mockAddItem).toHaveBeenCalledWith( | |
| expect.objectContaining({ | |
| type: 'gemini', | |
| text: 'Model response content', | |
| }), | |
| expect.any(Number), | |
| ); | |
| }); | |
| expect(result.current.thought).toBeNull(); | |
| expect(mockAddItem).not.toHaveBeenCalledWith( | |
| expect.objectContaining({ type: 'thinking' }), | |
| expect.any(Number), | |
| ); | |
| }); | |
| 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 } = await renderHookWithProviders(() => | |
| useGeminiStream( | |
| new MockedGeminiClientClass(mockConfig), | |
| [], | |
| mockAddItem, | |
| mockConfig, | |
| mockLoadedSettings, | |
| mockOnDebugMessage, | |
| mockHandleSlashCommand, | |
| false, | |
| () => 'vscode' as EditorType, | |
| () => {}, | |
| () => Promise.resolve(), | |
| false, | |
| () => {}, | |
| () => {}, | |
| () => {}, | |
| 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 memoize pendingHistoryItems', async () => { | |
| mockUseToolScheduler.mockReturnValue([ | |
| [], | |
| mockScheduleToolCalls, | |
| mockMarkToolsAsSubmitted, | |
| vi.fn(), | |
| mockCancelAllToolCalls, | |
| 0, | |
| ]); | |
| const { result, rerender } = await renderHookWithProviders(() => | |
| useGeminiStream( | |
| mockConfig.getGeminiClient(), | |
| [], | |
| mockAddItem, | |
| mockConfig, | |
| mockLoadedSettings, | |
| mockOnDebugMessage, | |
| mockHandleSlashCommand, | |
| false, | |
| () => 'vscode' as EditorType, | |
| () => {}, | |
| () => Promise.resolve(), | |
| false, | |
| () => {}, | |
| () => {}, | |
| () => {}, | |
| 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: CoreToolCallStatus.Executing, | |
| tool: { | |
| name: 'tool1', | |
| displayName: 'tool1', | |
| description: 'desc1', | |
| build: vi.fn(), | |
| }, | |
| invocation: { | |
| getDescription: () => 'Mock description', | |
| }, | |
| } as unknown as TrackedExecutingToolCall, | |
| ]; | |
| mockUseToolScheduler.mockReturnValue([ | |
| newToolCalls, | |
| mockScheduleToolCalls, | |
| mockMarkToolsAsSubmitted, | |
| vi.fn(), | |
| mockCancelAllToolCalls, | |
| 0, | |
| ]); | |
| 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 } = await renderHookWithProviders(() => | |
| useGeminiStream( | |
| new MockedGeminiClientClass(mockConfig), | |
| [], | |
| mockAddItem, | |
| mockConfig, | |
| mockLoadedSettings, | |
| mockOnDebugMessage, | |
| mockHandleSlashCommand, | |
| false, | |
| () => 'vscode' as EditorType, | |
| () => {}, | |
| () => Promise.resolve(), | |
| false, | |
| () => {}, | |
| () => {}, | |
| () => {}, | |
| 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 } = await renderHookWithProviders(() => | |
| useGeminiStream( | |
| new MockedGeminiClientClass(mockConfig), | |
| [], | |
| mockAddItem, | |
| mockConfig, | |
| mockLoadedSettings, | |
| mockOnDebugMessage, | |
| mockHandleSlashCommand, | |
| false, | |
| () => 'vscode' as EditorType, | |
| () => {}, | |
| () => Promise.resolve(), | |
| false, | |
| () => {}, | |
| () => {}, | |
| () => {}, | |
| 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: CoreToolCallStatus.Error, | |
| }), | |
| expect.any(Number), | |
| ); | |
| }); | |
| // Verify parseAndFormatApiError was called | |
| expect(mockParseAndFormatApiError).toHaveBeenCalledWith( | |
| { message: 'Test error' }, | |
| expect.any(String), | |
| undefined, | |
| 'gemini-2.5-pro', | |
| 'gemini-2.5-flash', | |
| ); | |
| }); | |
| it('should update lastOutputTime on Gemini thought and content events', async () => { | |
| vi.useFakeTimers(); | |
| const startTime = 1000000; | |
| vi.setSystemTime(startTime); | |
| // Mock a stream that yields a thought then content | |
| mockSendMessageStream.mockReturnValue( | |
| (async function* () { | |
| yield { | |
| type: ServerGeminiEventType.Thought, | |
| value: { subject: 'Thinking...', description: '' }, | |
| }; | |
| // Advance time for the next event | |
| vi.advanceTimersByTime(1000); | |
| yield { | |
| type: ServerGeminiEventType.Content, | |
| value: 'Hello', | |
| }; | |
| })(), | |
| ); | |
| const { result } = await renderHookWithProviders(() => | |
| useGeminiStream( | |
| new MockedGeminiClientClass(mockConfig), | |
| [], | |
| mockAddItem, | |
| mockConfig, | |
| mockLoadedSettings, | |
| mockOnDebugMessage, | |
| mockHandleSlashCommand, | |
| false, | |
| () => 'vscode' as EditorType, | |
| () => {}, | |
| () => Promise.resolve(), | |
| false, | |
| () => {}, | |
| () => {}, | |
| () => {}, | |
| 80, | |
| 24, | |
| ), | |
| ); | |
| // Reset fake timers to startTime because the asynchronous render lifecycle | |
| // (via waitUntilReady) advances the mock clock while waiting for initial | |
| // components to settle. | |
| vi.setSystemTime(startTime); | |
| // Submit query | |
| await act(async () => { | |
| await result.current.submitQuery('Test query'); | |
| }); | |
| // Verify lastOutputTime was updated | |
| // It should be the time of the last event (startTime + 1000) | |
| expect(result.current.lastOutputTime).toBe(startTime + 1000); | |
| vi.useRealTimers(); | |
| }); | |
| }); | |
| 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 } = await 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); | |
| // Mock for the initial request | |
| mockSendMessageStream.mockReturnValueOnce( | |
| (async function* () { | |
| yield { | |
| type: ServerGeminiEventType.LoopDetected, | |
| }; | |
| })(), | |
| ); | |
| // Mock for the retry request | |
| mockSendMessageStream.mockReturnValueOnce( | |
| (async function* () { | |
| yield { | |
| type: ServerGeminiEventType.Content, | |
| value: 'Retry successful', | |
| }; | |
| yield { | |
| type: ServerGeminiEventType.Finished, | |
| value: { reason: 'STOP' }, | |
| }; | |
| })(), | |
| ); | |
| const { result } = await 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 "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. Retrying request...', | |
| }); | |
| // Verify that the request was retried | |
| await waitFor(() => { | |
| expect(mockSendMessageStream).toHaveBeenCalledTimes(2); | |
| expect(mockSendMessageStream).toHaveBeenNthCalledWith( | |
| 2, | |
| 'test query', | |
| expect.any(AbortSignal), | |
| expect.any(String), | |
| undefined, | |
| 'test query', | |
| ); | |
| }); | |
| }); | |
| 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 } = await 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.', | |
| }); | |
| // Verify that the request was NOT retried | |
| expect(mockSendMessageStream).toHaveBeenCalledTimes(1); | |
| }); | |
| it('should handle multiple loop detection events properly', async () => { | |
| const { result } = await 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.', | |
| }); | |
| // Second loop detection - set up fresh mock for second call | |
| mockSendMessageStream.mockReturnValueOnce( | |
| (async function* () { | |
| yield { | |
| type: ServerGeminiEventType.LoopDetected, | |
| }; | |
| })(), | |
| ); | |
| // Mock for the retry request | |
| mockSendMessageStream.mockReturnValueOnce( | |
| (async function* () { | |
| yield { | |
| type: ServerGeminiEventType.Content, | |
| value: 'Retry successful', | |
| }; | |
| yield { | |
| type: ServerGeminiEventType.Finished, | |
| value: { reason: 'STOP' }, | |
| }; | |
| })(), | |
| ); | |
| // 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. Retrying request...', | |
| }); | |
| // Verify that the request was retried | |
| await waitFor(() => { | |
| expect(mockSendMessageStream).toHaveBeenCalledTimes(3); // 1st query, 2nd query, retry of 2nd query | |
| expect(mockSendMessageStream).toHaveBeenNthCalledWith( | |
| 3, | |
| 'second query', | |
| expect.any(AbortSignal), | |
| expect.any(String), | |
| undefined, | |
| 'second query', | |
| ); | |
| }); | |
| }); | |
| 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 } = await 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(); | |
| }); | |
| }); | |
| describe('Race Condition Prevention', () => { | |
| it('should reject concurrent submitQuery when already responding', async () => { | |
| // Stream that stays open (simulates "still responding") | |
| mockSendMessageStream.mockReturnValue( | |
| (async function* () { | |
| yield { | |
| type: ServerGeminiEventType.Content, | |
| value: 'First response', | |
| }; | |
| // Keep the stream open | |
| await new Promise(() => {}); | |
| })(), | |
| ); | |
| const { result } = await renderTestHook(); | |
| // Start first query without awaiting (fire-and-forget, like existing tests) | |
| await act(async () => { | |
| // eslint-disable-next-line @typescript-eslint/no-floating-promises | |
| result.current.submitQuery('first query'); | |
| }); | |
| // Wait for the stream to start responding | |
| await waitFor(() => { | |
| expect(result.current.streamingState).toBe(StreamingState.Responding); | |
| }); | |
| // Try a second query while first is still responding | |
| await act(async () => { | |
| // eslint-disable-next-line @typescript-eslint/no-floating-promises | |
| result.current.submitQuery('second query'); | |
| }); | |
| // Should have only called sendMessageStream once (second was rejected) | |
| expect(mockSendMessageStream).toHaveBeenCalledTimes(1); | |
| }); | |
| it('should allow continuation queries via loop detection retry', async () => { | |
| const mockLoopDetectionService = { | |
| disableForSession: vi.fn(), | |
| }; | |
| const mockClient = { | |
| ...new MockedGeminiClientClass(mockConfig), | |
| getLoopDetectionService: () => mockLoopDetectionService, | |
| }; | |
| mockConfig.getGeminiClient = vi.fn().mockReturnValue(mockClient); | |
| // First call triggers loop detection | |
| mockSendMessageStream.mockReturnValueOnce( | |
| (async function* () { | |
| yield { | |
| type: ServerGeminiEventType.LoopDetected, | |
| }; | |
| })(), | |
| ); | |
| // Retry call succeeds | |
| mockSendMessageStream.mockReturnValueOnce( | |
| (async function* () { | |
| yield { | |
| type: ServerGeminiEventType.Content, | |
| value: 'Retry success', | |
| }; | |
| yield { | |
| type: ServerGeminiEventType.Finished, | |
| value: { reason: 'STOP' }, | |
| }; | |
| })(), | |
| ); | |
| const { result } = await renderTestHook(); | |
| await act(async () => { | |
| await result.current.submitQuery('test query'); | |
| }); | |
| await waitFor(() => { | |
| expect( | |
| result.current.loopDetectionConfirmationRequest, | |
| ).not.toBeNull(); | |
| }); | |
| // User selects "disable" which triggers a continuation query | |
| await act(async () => { | |
| result.current.loopDetectionConfirmationRequest?.onComplete({ | |
| userSelection: 'disable', | |
| }); | |
| }); | |
| // Verify disableForSession was called | |
| expect( | |
| mockLoopDetectionService.disableForSession, | |
| ).toHaveBeenCalledTimes(1); | |
| // Continuation query should have gone through (2 total calls) | |
| await waitFor(() => { | |
| expect(mockSendMessageStream).toHaveBeenCalledTimes(2); | |
| expect(mockSendMessageStream).toHaveBeenNthCalledWith( | |
| 2, | |
| 'test query', | |
| expect.any(AbortSignal), | |
| expect.any(String), | |
| undefined, | |
| 'test query', | |
| ); | |
| }); | |
| }); | |
| }); | |
| }); | |
| describe('Agent Execution Events', () => { | |
| it('should handle AgentExecutionStopped event with systemMessage', async () => { | |
| mockSendMessageStream.mockReturnValue( | |
| (async function* () { | |
| yield { | |
| type: ServerGeminiEventType.AgentExecutionStopped, | |
| value: { | |
| reason: 'hook-reason', | |
| systemMessage: 'Custom stop message', | |
| }, | |
| }; | |
| })(), | |
| ); | |
| const { result } = await renderTestHook(); | |
| await act(async () => { | |
| await result.current.submitQuery('test stop'); | |
| }); | |
| await waitFor(() => { | |
| expect(mockAddItem).toHaveBeenCalledWith( | |
| { | |
| type: MessageType.INFO, | |
| text: 'Agent execution stopped: Custom stop message', | |
| }, | |
| expect.any(Number), | |
| ); | |
| expect(result.current.streamingState).toBe(StreamingState.Idle); | |
| }); | |
| }); | |
| it('should handle AgentExecutionStopped event by falling back to reason when systemMessage is missing', async () => { | |
| mockSendMessageStream.mockReturnValue( | |
| (async function* () { | |
| yield { | |
| type: ServerGeminiEventType.AgentExecutionStopped, | |
| value: { reason: 'Stopped by hook' }, | |
| }; | |
| })(), | |
| ); | |
| const { result } = await renderTestHook(); | |
| await act(async () => { | |
| await result.current.submitQuery('test stop'); | |
| }); | |
| await waitFor(() => { | |
| expect(mockAddItem).toHaveBeenCalledWith( | |
| { | |
| type: MessageType.INFO, | |
| text: 'Agent execution stopped: Stopped by hook', | |
| }, | |
| expect.any(Number), | |
| ); | |
| expect(result.current.streamingState).toBe(StreamingState.Idle); | |
| }); | |
| }); | |
| it('should handle AgentExecutionBlocked event with systemMessage', async () => { | |
| mockSendMessageStream.mockReturnValue( | |
| (async function* () { | |
| yield { | |
| type: ServerGeminiEventType.AgentExecutionBlocked, | |
| value: { | |
| reason: 'hook-reason', | |
| systemMessage: 'Custom block message', | |
| }, | |
| }; | |
| })(), | |
| ); | |
| const { result } = await renderTestHook(); | |
| await act(async () => { | |
| await result.current.submitQuery('test block'); | |
| }); | |
| await waitFor(() => { | |
| expect(mockAddItem).toHaveBeenCalledWith( | |
| { | |
| type: MessageType.WARNING, | |
| text: 'Agent execution blocked: Custom block message', | |
| }, | |
| expect.any(Number), | |
| ); | |
| }); | |
| }); | |
| it('should handle AgentExecutionBlocked event by falling back to reason when systemMessage is missing', async () => { | |
| mockSendMessageStream.mockReturnValue( | |
| (async function* () { | |
| yield { | |
| type: ServerGeminiEventType.AgentExecutionBlocked, | |
| value: { reason: 'Blocked by hook' }, | |
| }; | |
| })(), | |
| ); | |
| const { result } = await renderTestHook(); | |
| await act(async () => { | |
| await result.current.submitQuery('test block'); | |
| }); | |
| await waitFor(() => { | |
| expect(mockAddItem).toHaveBeenCalledWith( | |
| { | |
| type: MessageType.WARNING, | |
| text: 'Agent execution blocked: Blocked by hook', | |
| }, | |
| expect.any(Number), | |
| ); | |
| }); | |
| }); | |
| }); | |
| describe('Stream Splitting', () => { | |
| it('should not add empty history item when splitting message results in empty or whitespace-only beforeText', async () => { | |
| // Mock split point to always be 0, causing beforeText to be empty | |
| vi.mocked(findLastSafeSplitPoint).mockReturnValue(0); | |
| mockSendMessageStream.mockReturnValue( | |
| (async function* () { | |
| yield { type: ServerGeminiEventType.Content, value: 'test content' }; | |
| })(), | |
| ); | |
| const { result } = await renderTestHook(); | |
| await act(async () => { | |
| await result.current.submitQuery('user query'); | |
| }); | |
| await waitFor(() => { | |
| // We expect the stream to be processed. | |
| // Since beforeText is empty (0 split), addItem should NOT be called for it. | |
| // addItem IS called for the user query "user query". | |
| }); | |
| // Check addItem calls. | |
| // It should be called for user query and for the content. | |
| expect(mockAddItem).toHaveBeenCalledTimes(2); | |
| expect(mockAddItem).toHaveBeenCalledWith( | |
| expect.objectContaining({ type: 'user', text: 'user query' }), | |
| expect.any(Number), | |
| ); | |
| expect(mockAddItem).toHaveBeenLastCalledWith( | |
| expect.objectContaining({ | |
| type: 'gemini_content', | |
| text: 'test content', | |
| }), | |
| expect.any(Number), | |
| ); | |
| // Verify that pendingHistoryItem is empty after (afterText). | |
| expect(result.current.pendingHistoryItems.length).toEqual(0); | |
| // Reset mock | |
| vi.mocked(findLastSafeSplitPoint).mockReset(); | |
| vi.mocked(findLastSafeSplitPoint).mockImplementation( | |
| (s: string) => s.length, | |
| ); | |
| }); | |
| it('should add whitespace-only history item when splitting message', async () => { | |
| // Input: " content" | |
| // Split at 3 -> before: " ", after: "content" | |
| vi.mocked(findLastSafeSplitPoint).mockReturnValue(3); | |
| mockSendMessageStream.mockReturnValue( | |
| (async function* () { | |
| yield { type: ServerGeminiEventType.Content, value: ' content' }; | |
| })(), | |
| ); | |
| const { result } = await renderTestHook(); | |
| await act(async () => { | |
| await result.current.submitQuery('user query'); | |
| }); | |
| await waitFor(() => {}); | |
| expect(mockAddItem).toHaveBeenCalledTimes(3); | |
| expect(mockAddItem).toHaveBeenCalledWith( | |
| expect.objectContaining({ type: 'user', text: 'user query' }), | |
| expect.any(Number), | |
| ); | |
| expect(mockAddItem).toHaveBeenLastCalledWith( | |
| expect.objectContaining({ | |
| type: 'gemini_content', | |
| text: 'content', | |
| }), | |
| expect.any(Number), | |
| ); | |
| expect(result.current.pendingHistoryItems.length).toEqual(0); | |
| }); | |
| }); | |
| it('should trace UserPrompt telemetry on submitQuery', async () => { | |
| const { result } = await renderTestHook(); | |
| mockSendMessageStream.mockReturnValue( | |
| (async function* () { | |
| yield { type: ServerGeminiEventType.Content, value: 'Response' }; | |
| })(), | |
| ); | |
| await act(async () => { | |
| await result.current.submitQuery('telemetry test query'); | |
| }); | |
| const userPromptCall = mockRunInDevTraceSpan.mock.calls.find( | |
| (call) => | |
| call[0].operation === GeminiCliOperation.UserPrompt || | |
| call[0].operation === 'UserPrompt', | |
| ); | |
| expect(userPromptCall).toBeDefined(); | |
| const spanMetadata = {} as SpanMetadata; | |
| await act(async () => { | |
| await userPromptCall; | |
| }); | |
| expect(spanMetadata.input).toBe('telemetry test query'); | |
| }); | |
| }); | |