2025-05-29 22:30:18 +00:00
/**
* @license
* Copyright 2025 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
2025-06-02 01:50:28 -07:00
/* eslint-disable @typescript-eslint/no-explicit-any */
2026-03-12 02:51:40 +05:30
import {
describe ,
it ,
expect ,
vi ,
beforeEach ,
type Mock ,
type MockInstance ,
} from 'vitest' ;
2025-10-28 10:32:15 -07:00
import { act } from 'react' ;
2026-01-27 16:06:24 -08:00
import { renderHookWithProviders } from '../../test-utils/render.js' ;
2025-10-30 11:50:26 -07:00
import { waitFor } from '../../test-utils/async.js' ;
2025-08-22 14:12:05 -07:00
import { useGeminiStream } from './useGeminiStream.js' ;
2025-08-12 14:05:49 -07:00
import { useKeypress } from './useKeypress.js' ;
2025-08-20 15:51:31 -04:00
import * as atCommandProcessor from './atCommandProcessor.js' ;
2026-03-12 02:51:40 +05:30
import {
useToolScheduler ,
type TrackedToolCall ,
type TrackedCompletedToolCall ,
type TrackedExecutingToolCall ,
type TrackedCancelledToolCall ,
type TrackedWaitingToolCall ,
2026-01-21 00:18:42 -05:00
} from './useToolScheduler.js' ;
2025-08-26 00:04:53 +02:00
import type {
2025-07-22 06:57:11 +09:00
Config ,
EditorType ,
2025-08-06 10:50:02 -07:00
AnyToolInvocation ,
2026-02-26 18:26:16 -08:00
SpanMetadata ,
2025-07-22 06:57:11 +09:00
} from '@google/gemini-cli-core' ;
2025-07-07 16:45:44 -04:00
import {
2026-02-13 11:27:20 -05:00
CoreToolCallStatus ,
2025-08-26 00:04:53 +02:00
ApprovalMode ,
AuthType ,
GeminiEventType as ServerGeminiEventType ,
ToolErrorType ,
2025-09-14 20:20:21 -07:00
ToolConfirmationOutcome ,
2026-02-13 11:14:35 +09:00
MessageBusType ,
2025-10-09 10:22:26 -07:00
tokenLimit ,
2025-10-28 19:05:48 +00:00
debugLogger ,
2026-01-13 23:03:19 -05:00
coreEvents ,
CoreEvent ,
2026-01-15 15:33:16 -05:00
MCPDiscoveryState ,
2026-02-26 18:26:16 -08:00
GeminiCliOperation ,
2026-02-24 14:31:41 -05:00
getPlanModeExitMessage ,
2025-08-26 00:04:53 +02:00
} from '@google/gemini-cli-core' ;
import type { Part , PartListUnion } from '@google/genai' ;
import type { UseHistoryManagerReturn } from './useHistoryManager.js' ;
2025-10-27 09:59:08 -07:00
import type { SlashCommandProcessorResult } from '../types.js' ;
2026-02-13 17:20:14 -05:00
import { MessageType , StreamingState } from '../types.js' ;
2026-02-17 18:41:43 -08:00
2025-08-26 00:04:53 +02:00
import type { LoadedSettings } from '../../config/settings.js' ;
2026-02-18 12:53:06 -08:00
import { findLastSafeSplitPoint } from '../utils/markdownUtilities.js' ;
2026-03-03 01:22:29 -08:00
import { theme } from '../semantic-colors.js' ;
2025-05-29 22:30:18 +00:00
2025-06-02 01:50:28 -07:00
// --- MOCKS ---
const mockSendMessageStream = vi
. fn ( )
. mockReturnValue ( ( async function * ( ) { } ) ( ) ) ;
const mockStartChat = vi . fn ( ) ;
2026-02-13 11:14:35 +09:00
const mockMessageBus = {
publish : vi.fn ( ) ,
subscribe : vi.fn ( ) ,
unsubscribe : vi.fn ( ) ,
} ;
2025-06-02 01:50:28 -07:00
2025-06-02 22:30:52 -07:00
const MockedGeminiClientClass = vi . hoisted ( ( ) = >
vi . fn ( ) . mockImplementation ( function ( this : any , _config : any ) {
2025-06-02 01:50:28 -07:00
// _config
this . startChat = mockStartChat ;
this . sendMessageStream = mockSendMessageStream ;
2025-06-08 11:14:45 -07:00
this . addHistory = vi . fn ( ) ;
2026-02-18 14:05:50 -08:00
this . generateContent = vi . fn ( ) . mockResolvedValue ( {
candidates : [
{ content : { parts : [ { text : 'Got it. Focusing on tests only.' } ] } } ,
] ,
} ) ;
2026-01-27 16:06:24 -08:00
this . getCurrentSequenceModel = vi . fn ( ) . mockReturnValue ( 'test-model' ) ;
2025-10-06 13:34:00 -06:00
this . getChat = vi . fn ( ) . mockReturnValue ( {
recordCompletedToolCalls : vi.fn ( ) ,
} ) ;
2025-09-02 23:29:07 -06:00
this . getChatRecordingService = vi . fn ( ) . mockReturnValue ( {
recordThought : vi.fn ( ) ,
initialize : vi.fn ( ) ,
recordMessage : vi.fn ( ) ,
recordMessageTokens : vi.fn ( ) ,
recordToolCalls : vi.fn ( ) ,
getConversationFile : vi.fn ( ) ,
} ) ;
2026-01-30 09:53:09 -08:00
this . getCurrentSequenceModel = vi
. fn ( )
. mockReturnValue ( 'gemini-2.0-flash-exp' ) ;
2025-06-02 22:30:52 -07:00
} ) ,
) ;
2025-06-22 09:26:48 -05:00
const MockedUserPromptEvent = vi . hoisted ( ( ) = >
vi . fn ( ) . mockImplementation ( ( ) = > { } ) ,
) ;
2025-08-13 17:57:11 +00:00
const mockParseAndFormatApiError = vi . hoisted ( ( ) = > vi . fn ( ) ) ;
2026-03-12 00:03:54 -04:00
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' )
) ;
} ,
) ;
2025-06-22 09:26:48 -05:00
2026-01-27 13:17:40 -08:00
const MockValidationRequiredError = vi . hoisted (
( ) = >
class extends Error {
userHandled = false ;
} ,
) ;
2026-02-26 18:26:16 -08:00
const mockRunInDevTraceSpan = vi . hoisted ( ( ) = >
vi . fn ( async ( opts , fn ) = > {
const metadata : SpanMetadata = {
name : opts.operation ,
attributes : opts.attributes || { } ,
} ;
return await fn ( {
metadata ,
endSpan : vi.fn ( ) ,
} ) ;
} ) ,
) ;
2025-06-25 05:41:11 -07:00
vi . mock ( '@google/gemini-cli-core' , async ( importOriginal ) = > {
2025-06-02 22:30:52 -07:00
const actualCoreModule = ( await importOriginal ( ) ) as any ;
2025-06-02 01:50:28 -07:00
return {
2025-06-11 15:33:09 -04:00
. . . actualCoreModule ,
2026-03-12 00:03:54 -04:00
isBackgroundExecutionData : mockIsBackgroundExecutionData ,
2025-06-11 15:33:09 -04:00
GitService : vi.fn ( ) ,
GeminiClient : MockedGeminiClientClass ,
2025-06-22 09:26:48 -05:00
UserPromptEvent : MockedUserPromptEvent ,
2026-01-27 13:17:40 -08:00
ValidationRequiredError : MockValidationRequiredError ,
2025-08-13 17:57:11 +00:00
parseAndFormatApiError : mockParseAndFormatApiError ,
2025-10-09 10:22:26 -07:00
tokenLimit : vi.fn ( ) . mockReturnValue ( 100 ) , // Mock tokenLimit
2026-01-27 16:06:24 -08:00
recordToolCallInteractions : vi.fn ( ) . mockResolvedValue ( undefined ) ,
getCodeAssistServer : vi.fn ( ) . mockReturnValue ( undefined ) ,
2026-02-26 18:26:16 -08:00
runInDevTraceSpan : mockRunInDevTraceSpan ,
2025-06-02 01:50:28 -07:00
} ;
} ) ;
2026-01-21 00:18:42 -05:00
const mockUseToolScheduler = useToolScheduler as Mock ;
vi . mock ( './useToolScheduler.js' , async ( importOriginal ) = > {
2025-06-02 01:50:28 -07:00
const actualSchedulerModule = ( await importOriginal ( ) ) as any ;
2025-05-29 22:30:18 +00:00
return {
2025-06-02 01:50:28 -07:00
. . . ( actualSchedulerModule || { } ) ,
2026-01-21 00:18:42 -05:00
useToolScheduler : vi.fn ( ) ,
2025-05-29 22:30:18 +00:00
} ;
} ) ;
2025-08-12 14:05:49 -07:00
vi . mock ( './useKeypress.js' , ( ) = > ( {
useKeypress : vi.fn ( ) ,
} ) ) ;
2025-06-02 01:50:28 -07:00
vi . mock ( './shellCommandProcessor.js' , ( ) = > ( {
useShellCommandProcessor : vi.fn ( ) . mockReturnValue ( {
handleShellCommand : vi.fn ( ) ,
2026-01-13 06:19:53 -08:00
activeShellPtyId : null ,
lastShellOutputTime : 0 ,
2025-06-02 01:50:28 -07:00
} ) ,
} ) ) ;
2025-08-20 15:51:31 -04:00
vi . mock ( './atCommandProcessor.js' ) ;
2025-06-02 01:50:28 -07:00
vi . mock ( '../utils/markdownUtilities.js' , ( ) = > ( {
findLastSafeSplitPoint : vi.fn ( ( s : string ) = > s . length ) ,
} ) ) ;
vi . mock ( './useStateAndRef.js' , ( ) = > ( {
useStateAndRef : vi.fn ( ( initial ) = > {
let val = initial ;
const ref = { current : val } ;
const setVal = vi . fn ( ( updater ) = > {
if ( typeof updater === 'function' ) {
val = updater ( val ) ;
} else {
val = updater ;
}
ref . current = val ;
} ) ;
2025-09-17 15:37:13 -07:00
return [ val , ref , setVal ] ;
2025-06-02 01:50:28 -07:00
} ) ,
} ) ) ;
vi . mock ( './useLogger.js' , ( ) = > ( {
useLogger : vi.fn ( ) . mockReturnValue ( {
logMessage : vi.fn ( ) . mockResolvedValue ( undefined ) ,
} ) ,
} ) ) ;
2025-07-10 00:19:30 +05:30
const mockStartNewPrompt = vi . fn ( ) ;
2025-06-09 20:25:37 -04:00
const mockAddUsage = vi . fn ( ) ;
2026-02-17 14:59:33 -08:00
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 ) ,
} ) ) ,
} ;
} ) ;
2025-06-09 20:25:37 -04:00
2025-06-02 01:50:28 -07:00
vi . mock ( './slashCommandProcessor.js' , ( ) = > ( {
handleSlashCommand : vi.fn ( ) . mockReturnValue ( false ) ,
} ) ) ;
2025-11-11 07:50:11 -08:00
vi . mock ( './useAlternateBuffer.js' , ( ) = > ( {
useAlternateBuffer : vi.fn ( ( ) = > false ) ,
} ) ) ;
2025-06-02 01:50:28 -07:00
// --- END MOCKS ---
// --- Tests for useGeminiStream Hook ---
describe ( 'useGeminiStream' , ( ) = > {
2026-01-27 16:06:24 -08:00
let mockAddItem = vi . fn ( ) ;
let mockOnDebugMessage = vi . fn ( ) ;
let mockHandleSlashCommand = vi . fn ( ) . mockResolvedValue ( false ) ;
2025-06-02 01:50:28 -07:00
let mockScheduleToolCalls : Mock ;
let mockCancelAllToolCalls : Mock ;
let mockMarkToolsAsSubmitted : Mock ;
2025-08-20 15:51:31 -04:00
let handleAtCommandSpy : MockInstance ;
2025-06-02 01:50:28 -07:00
2026-01-27 16:06:24 -08:00
const emptyHistory : any [ ] = [ ] ;
let capturedOnComplete : any = 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 ;
} ) ;
2025-06-02 22:30:52 -07:00
2026-01-27 16:06:24 -08:00
const mockMcpClientManager = {
getDiscoveryState : vi.fn ( ) . mockReturnValue ( MCPDiscoveryState . COMPLETED ) ,
getMcpServerCount : vi.fn ( ) . mockReturnValue ( 0 ) ,
} ;
2026-01-15 15:33:16 -05:00
2026-01-27 16:06:24 -08:00
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 ( ) ,
2026-03-06 19:14:44 -08:00
resetBillingTurnState : vi.fn ( ) ,
2026-01-27 16:06:24 -08:00
getQuotaErrorOccurred : vi.fn ( ( ) = > false ) ,
getModel : vi.fn ( ( ) = > 'gemini-2.5-pro' ) ,
getContentGeneratorConfig : vi.fn ( ( ) = > ( {
2025-07-11 22:17:46 +05:30
model : 'test-model' ,
apiKey : 'test-key' ,
vertexai : false ,
authType : AuthType.USE_GEMINI ,
2026-01-27 16:06:24 -08:00
} ) ) ,
getContentGenerator : vi.fn ( ) ,
isInteractive : ( ) = > false ,
getExperiments : ( ) = > { } ,
getMaxSessionTurns : vi.fn ( ( ) = > 100 ) ,
isJitContextEnabled : vi.fn ( ( ) = > false ) ,
getGlobalMemory : vi.fn ( ( ) = > '' ) ,
getUserMemory : vi.fn ( ( ) = > '' ) ,
2026-02-13 11:14:35 +09:00
getMessageBus : vi.fn ( ( ) = > mockMessageBus ) ,
2026-02-18 14:05:50 -08:00
getBaseLlmClient : vi.fn ( ( ) = > ( {
generateContent : vi.fn ( ) . mockResolvedValue ( {
candidates : [
{ content : { parts : [ { text : 'Got it. Focusing on tests only.' } ] } } ,
] ,
} ) ,
} ) ) ,
2026-01-27 16:06:24 -08:00
getIdeMode : vi.fn ( ( ) = > false ) ,
getEnableHooks : vi.fn ( ( ) = > false ) ,
} as unknown as Config ;
2025-07-11 22:17:46 +05:30
2026-01-27 16:06:24 -08:00
beforeEach ( ( ) = > {
vi . clearAllMocks ( ) ; // Clear mocks before each test
mockAddItem = vi . fn ( ) ;
2025-06-02 01:50:28 -07:00
mockOnDebugMessage = vi . fn ( ) ;
2025-06-11 15:33:09 -04:00
mockHandleSlashCommand = vi . fn ( ) . mockResolvedValue ( false ) ;
2025-06-02 01:50:28 -07:00
// Mock return value for useReactToolScheduler
mockScheduleToolCalls = vi . fn ( ) ;
mockCancelAllToolCalls = vi . fn ( ) ;
mockMarkToolsAsSubmitted = vi . fn ( ) ;
2026-01-27 16:06:24 -08:00
// Reset properties of mockConfig if needed
( mockConfig . getCheckpointingEnabled as Mock ) . mockReturnValue ( false ) ;
( mockConfig . getApprovalMode as Mock ) . mockReturnValue ( ApprovalMode . DEFAULT ) ;
2025-06-02 01:50:28 -07:00
// Default mock for useReactToolScheduler to prevent toolCalls being undefined initially
2026-01-21 00:18:42 -05:00
mockUseToolScheduler . mockReturnValue ( [
2025-06-02 01:50:28 -07:00
[ ] , // Default to empty array for toolCalls
mockScheduleToolCalls ,
mockMarkToolsAsSubmitted ,
2025-10-27 09:59:08 -07:00
vi . fn ( ) , // setToolCallsForDisplay
mockCancelAllToolCalls ,
2026-01-13 06:19:53 -08:00
0 , // lastToolOutputTime
2025-06-02 01:50:28 -07:00
] ) ;
// 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 * ( ) { } ) ( ) ) ;
2025-08-20 15:51:31 -04:00
handleAtCommandSpy = vi . spyOn ( atCommandProcessor , 'handleAtCommand' ) ;
2026-01-15 15:33:16 -05:00
vi . spyOn ( coreEvents , 'emitFeedback' ) ;
2025-06-02 01:50:28 -07:00
} ) ;
2025-06-12 02:21:54 +01:00
const mockLoadedSettings : LoadedSettings = {
2026-02-27 14:15:10 -05:00
merged : {
preferredEditor : 'vscode' ,
ui : { errorVerbosity : 'full' } ,
} ,
2025-06-12 02:21:54 +01:00
user : { path : '/user/settings.json' , settings : { } } ,
workspace : { path : '/workspace/.gemini/settings.json' , settings : { } } ,
errors : [ ] ,
forScope : vi.fn ( ) ,
setValue : vi.fn ( ) ,
} as unknown as LoadedSettings ;
2025-06-08 11:14:45 -07:00
const renderTestHook = (
initialToolCalls : TrackedToolCall [ ] = [ ] ,
geminiClient? : any ,
2026-02-27 14:15:10 -05:00
loadedSettings : LoadedSettings = mockLoadedSettings ,
2025-06-08 11:14:45 -07:00
) = > {
const client = geminiClient || mockConfig . getGeminiClient ( ) ;
2026-01-27 16:06:24 -08:00
let lastToolCalls = initialToolCalls ;
2025-06-08 11:14:45 -07:00
2025-10-27 09:59:08 -07:00
const initialProps = {
client ,
2026-01-27 16:06:24 -08:00
history : emptyHistory ,
2025-10-27 09:59:08 -07:00
addItem : mockAddItem as unknown as UseHistoryManagerReturn [ 'addItem' ] ,
config : mockConfig ,
onDebugMessage : mockOnDebugMessage ,
handleSlashCommand : mockHandleSlashCommand as unknown as (
cmd : PartListUnion ,
) = > Promise < SlashCommandProcessorResult | false > ,
shellModeActive : false ,
2026-02-27 14:15:10 -05:00
loadedSettings ,
2025-10-27 09:59:08 -07:00
toolCalls : initialToolCalls ,
} ;
2026-01-27 16:06:24 -08:00
mockUseToolScheduler . mockImplementation ( ( onComplete ) = > {
capturedOnComplete = onComplete ;
return [
lastToolCalls ,
mockScheduleToolCalls ,
mockMarkToolsAsSubmitted ,
( updater : any ) = > {
lastToolCalls =
typeof updater === 'function' ? updater ( lastToolCalls ) : updater ;
rerender ( { . . . initialProps , toolCalls : lastToolCalls } ) ;
} ,
( . . . args : any [ ] ) = > {
2025-10-27 09:59:08 -07:00
mockCancelAllToolCalls ( . . . args ) ;
2026-01-27 16:06:24 -08:00
lastToolCalls = lastToolCalls . map ( ( tc ) = > {
2025-10-27 09:59:08 -07:00
if (
2026-02-13 11:27:20 -05:00
tc . status === CoreToolCallStatus . AwaitingApproval ||
tc . status === CoreToolCallStatus . Executing ||
tc . status === CoreToolCallStatus . Scheduled ||
tc . status === CoreToolCallStatus . Validating
2025-10-27 09:59:08 -07:00
) {
return {
. . . tc ,
2026-02-13 11:27:20 -05:00
status : CoreToolCallStatus.Cancelled ,
2025-10-27 09:59:08 -07:00
response : {
callId : tc.request.callId ,
responseParts : [ ] ,
resultDisplay : 'Request cancelled.' ,
} ,
2026-01-27 16:06:24 -08:00
responseSubmittedToGemini : true ,
2025-10-27 09:59:08 -07:00
} as any as TrackedCancelledToolCall ;
}
return tc ;
} ) ;
2026-01-27 16:06:24 -08:00
rerender ( { . . . initialProps , toolCalls : lastToolCalls } ) ;
} ,
0 ,
] ;
} ) ;
2025-10-27 09:59:08 -07:00
2026-01-27 16:06:24 -08:00
const { result , rerender } = renderHookWithProviders (
( props : typeof initialProps ) = >
useGeminiStream (
2025-06-08 15:42:49 -07:00
props . client ,
2025-06-11 15:33:09 -04:00
props . history ,
2025-06-08 15:42:49 -07:00
props . addItem ,
props . config ,
2025-08-28 16:42:54 -07:00
props . loadedSettings ,
2025-06-08 15:42:49 -07:00
props . onDebugMessage ,
props . handleSlashCommand ,
props . shellModeActive ,
2026-01-27 16:06:24 -08:00
mockGetPreferredEditor ,
mockOnAuthError ,
mockPerformMemoryRefresh ,
2025-07-09 13:55:56 -04:00
false ,
2026-01-27 16:06:24 -08:00
mockSetModelSwitchedFromQuotaError ,
mockOnCancelSubmit ,
mockSetShellInputFocused ,
2025-09-11 13:27:27 -07:00
80 ,
24 ,
2026-01-27 16:06:24 -08:00
) ,
2025-06-08 15:42:49 -07:00
{
2025-10-27 09:59:08 -07:00
initialProps ,
2025-06-08 15:42:49 -07:00
} ,
2025-06-02 01:50:28 -07:00
) ;
return {
result ,
rerender ,
mockMarkToolsAsSubmitted ,
mockSendMessageStream ,
2025-06-08 11:14:45 -07:00
client ,
2025-06-02 01:50:28 -07:00
} ;
} ;
2025-11-03 23:40:57 +05:30
// Helper to create mock tool calls - reduces boilerplate
const createMockToolCall = (
toolName : string ,
callId : string ,
confirmationType : 'edit' | 'info' ,
2026-02-13 11:27:20 -05:00
status : TrackedToolCall [ 'status' ] = CoreToolCallStatus . AwaitingApproval ,
mockOnConfirm : Mock = vi . fn ( ) ,
2025-11-03 23:40:57 +05:30
) : TrackedWaitingToolCall = > ( {
request : {
callId ,
name : toolName ,
args : { } ,
isClientInitiated : false ,
prompt_id : 'prompt-id-1' ,
} ,
2026-02-13 11:27:20 -05:00
status : status as CoreToolCallStatus . AwaitingApproval ,
2025-11-03 23:40:57 +05:30
responseSubmittedToGemini : false ,
confirmationDetails :
confirmationType === 'edit'
? {
type : 'edit' ,
title : 'Confirm Edit' ,
fileName : 'file.txt' ,
filePath : '/test/file.txt' ,
fileDiff : 'fake diff' ,
originalContent : 'old' ,
newContent : 'new' ,
2026-02-13 11:27:20 -05:00
onConfirm : mockOnConfirm ,
2025-11-03 23:40:57 +05:30
}
: {
type : 'info' ,
title : ` ${ toolName } confirmation ` ,
prompt : ` Execute ${ toolName } ? ` ,
2026-02-13 11:27:20 -05:00
onConfirm : mockOnConfirm ,
2025-11-03 23:40:57 +05:30
} ,
tool : {
name : toolName ,
displayName : toolName ,
description : ` ${ toolName } description ` ,
build : vi.fn ( ) ,
} as any ,
invocation : {
getDescription : ( ) = > 'Mock description' ,
} as unknown as AnyToolInvocation ,
2026-02-13 11:14:35 +09:00
correlationId : ` corr- ${ callId } ` ,
2025-11-03 23:40:57 +05:30
} ) ;
// Helper to render hook with default parameters - reduces boilerplate
const renderHookWithDefaults = (
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 ;
2026-01-27 16:06:24 -08:00
return renderHookWithProviders ( ( ) = >
2025-11-03 23:40:57 +05:30
useGeminiStream (
new MockedGeminiClientClass ( mockConfig ) ,
[ ] ,
mockAddItem ,
mockConfig ,
mockLoadedSettings ,
mockOnDebugMessage ,
mockHandleSlashCommand ,
shellModeActive ,
( ) = > 'vscode' as EditorType ,
onAuthError ,
performMemoryRefresh ,
modelSwitched ,
setModelSwitched ,
onCancelSubmit ,
setShellInputFocused ,
80 ,
24 ,
) ,
) ;
} ;
2025-06-02 01:50:28 -07:00
it ( 'should not submit tool responses if not all tool calls are completed' , ( ) = > {
const toolCalls : TrackedToolCall [ ] = [
{
2025-06-22 01:35:36 -04:00
request : {
callId : 'call1' ,
name : 'tool1' ,
args : { } ,
isClientInitiated : false ,
2025-07-10 00:19:30 +05:30
prompt_id : 'prompt-id-1' ,
2025-06-22 01:35:36 -04:00
} ,
2026-02-13 11:27:20 -05:00
status : CoreToolCallStatus.Success ,
2025-06-02 01:50:28 -07:00
responseSubmittedToGemini : false ,
response : {
callId : 'call1' ,
responseParts : [ { text : 'tool 1 response' } ] ,
error : undefined ,
2025-08-20 15:51:31 -04:00
errorType : undefined , // FIX: Added missing property
2025-06-02 01:50:28 -07:00
resultDisplay : 'Tool 1 success display' ,
} ,
tool : {
name : 'tool1' ,
2025-08-06 10:50:02 -07:00
displayName : 'tool1' ,
2025-06-02 01:50:28 -07:00
description : 'desc1' ,
2025-08-06 10:50:02 -07:00
build : vi.fn ( ) ,
2025-06-02 01:50:28 -07:00
} as any ,
2025-08-06 10:50:02 -07:00
invocation : {
getDescription : ( ) = > ` Mock description ` ,
} as unknown as AnyToolInvocation ,
2025-06-02 01:50:28 -07:00
startTime : Date.now ( ) ,
endTime : Date.now ( ) ,
} as TrackedCompletedToolCall ,
{
2025-07-10 00:19:30 +05:30
request : {
callId : 'call2' ,
name : 'tool2' ,
args : { } ,
prompt_id : 'prompt-id-1' ,
} ,
2026-02-13 11:27:20 -05:00
status : CoreToolCallStatus.Executing ,
2025-06-02 01:50:28 -07:00
responseSubmittedToGemini : false ,
tool : {
name : 'tool2' ,
2025-08-06 10:50:02 -07:00
displayName : 'tool2' ,
2025-06-02 01:50:28 -07:00
description : 'desc2' ,
2025-08-06 10:50:02 -07:00
build : vi.fn ( ) ,
2025-06-02 01:50:28 -07:00
} as any ,
2025-08-06 10:50:02 -07:00
invocation : {
getDescription : ( ) = > ` Mock description ` ,
} as unknown as AnyToolInvocation ,
2025-06-02 01:50:28 -07:00
startTime : Date.now ( ) ,
liveOutput : '...' ,
} as TrackedExecutingToolCall ,
] ;
const { mockMarkToolsAsSubmitted , mockSendMessageStream } =
renderTestHook ( toolCalls ) ;
// Effect for submitting tool responses depends on toolCalls and isResponding
// isResponding is initially false, so the effect should run.
expect ( mockMarkToolsAsSubmitted ) . not . toHaveBeenCalled ( ) ;
expect ( mockSendMessageStream ) . not . toHaveBeenCalled ( ) ; // submitQuery uses this
} ) ;
2026-03-12 00:03:54 -04:00
it ( 'should expose activePtyId for non-shell executing tools that report an execution ID' , ( ) = > {
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 } = renderTestHook ( [ remoteExecutingTool ] ) ;
expect ( result . current . activePtyId ) . toBe ( 4242 ) ;
} ) ;
2025-06-02 01:50:28 -07:00
it ( 'should submit tool responses when all tool calls are completed and ready' , async ( ) = > {
2025-08-22 14:12:05 -07:00
const toolCall1ResponseParts : Part [ ] = [ { text : 'tool 1 final response' } ] ;
const toolCall2ResponseParts : Part [ ] = [ { text : 'tool 2 final response' } ] ;
2025-06-22 01:35:36 -04:00
const completedToolCalls : TrackedToolCall [ ] = [
2025-06-02 01:50:28 -07:00
{
2025-06-22 01:35:36 -04:00
request : {
2025-06-02 01:50:28 -07:00
callId : 'call1' ,
name : 'tool1' ,
2025-06-22 01:35:36 -04:00
args : { } ,
isClientInitiated : false ,
2025-07-10 00:19:30 +05:30
prompt_id : 'prompt-id-2' ,
2025-06-22 01:35:36 -04:00
} ,
2026-02-13 11:27:20 -05:00
status : CoreToolCallStatus.Success ,
2025-06-22 01:35:36 -04:00
responseSubmittedToGemini : false ,
2025-08-20 15:51:31 -04:00
response : {
callId : 'call1' ,
responseParts : toolCall1ResponseParts ,
errorType : undefined , // FIX: Added missing property
} ,
2025-08-06 10:50:02 -07:00
tool : {
displayName : 'MockTool' ,
} ,
invocation : {
getDescription : ( ) = > ` Mock description ` ,
} as unknown as AnyToolInvocation ,
2025-06-02 01:50:28 -07:00
} as TrackedCompletedToolCall ,
{
2025-06-22 01:35:36 -04:00
request : {
2025-06-02 01:50:28 -07:00
callId : 'call2' ,
name : 'tool2' ,
2025-06-22 01:35:36 -04:00
args : { } ,
isClientInitiated : false ,
2025-07-10 00:19:30 +05:30
prompt_id : 'prompt-id-2' ,
2025-06-22 01:35:36 -04:00
} ,
2026-02-13 11:27:20 -05:00
status : CoreToolCallStatus.Error ,
2025-06-22 01:35:36 -04:00
responseSubmittedToGemini : false ,
2025-08-20 15:51:31 -04:00
response : {
callId : 'call2' ,
responseParts : toolCall2ResponseParts ,
errorType : ToolErrorType.UNHANDLED_EXCEPTION , // FIX: Added missing property
} ,
2025-06-22 01:35:36 -04:00
} as TrackedCompletedToolCall , // Treat error as a form of completion for submission
2025-06-02 01:50:28 -07:00
] ;
2025-06-27 16:39:54 -07:00
// Capture the onComplete callback
let capturedOnComplete :
| ( ( completedTools : TrackedToolCall [ ] ) = > Promise < void > )
| null = null ;
2026-01-21 00:18:42 -05:00
mockUseToolScheduler . mockImplementation ( ( onComplete ) = > {
2025-06-27 16:39:54 -07:00
capturedOnComplete = onComplete ;
2026-01-27 16:06:24 -08:00
return [
[ ] ,
mockScheduleToolCalls ,
mockMarkToolsAsSubmitted ,
vi . fn ( ) ,
mockCancelAllToolCalls ,
0 ,
] ;
2025-06-27 16:39:54 -07:00
} ) ;
2026-01-27 16:06:24 -08:00
renderHookWithProviders ( ( ) = >
2025-06-22 01:35:36 -04:00
useGeminiStream (
new MockedGeminiClientClass ( mockConfig ) ,
[ ] ,
mockAddItem ,
mockConfig ,
2025-08-28 16:42:54 -07:00
mockLoadedSettings ,
2025-06-22 01:35:36 -04:00
mockOnDebugMessage ,
mockHandleSlashCommand ,
false ,
( ) = > 'vscode' as EditorType ,
( ) = > { } ,
( ) = > Promise . resolve ( ) ,
2025-07-09 13:55:56 -04:00
false ,
( ) = > { } ,
2025-08-06 15:19:10 -04:00
( ) = > { } ,
( ) = > { } ,
2025-09-11 13:27:27 -07:00
80 ,
24 ,
2025-06-22 01:35:36 -04:00
) ,
) ;
2025-06-02 01:50:28 -07:00
2025-06-27 16:39:54 -07:00
// Trigger the onComplete callback with completed tools
await act ( async ( ) = > {
if ( capturedOnComplete ) {
2026-01-27 16:06:24 -08:00
// Wait a tick for refs to be set up
await new Promise ( ( resolve ) = > setTimeout ( resolve , 0 ) ) ;
2025-06-27 16:39:54 -07:00
await capturedOnComplete ( completedToolCalls ) ;
}
2025-06-08 15:42:49 -07:00
} ) ;
2025-06-02 01:50:28 -07:00
2025-10-30 11:50:26 -07:00
await waitFor ( ( ) = > {
2025-06-22 01:35:36 -04:00
expect ( mockMarkToolsAsSubmitted ) . toHaveBeenCalledTimes ( 1 ) ;
expect ( mockSendMessageStream ) . toHaveBeenCalledTimes ( 1 ) ;
2025-06-02 01:50:28 -07:00
} ) ;
2025-08-22 14:12:05 -07:00
const expectedMergedResponse = [
. . . toolCall1ResponseParts ,
. . . toolCall2ResponseParts ,
] ;
2025-06-22 01:35:36 -04:00
expect ( mockSendMessageStream ) . toHaveBeenCalledWith (
2025-06-02 01:50:28 -07:00
expectedMergedResponse ,
2025-06-03 02:10:54 +00:00
expect . any ( AbortSignal ) ,
2025-07-10 00:19:30 +05:30
'prompt-id-2' ,
2026-01-30 10:09:27 -08:00
undefined ,
false ,
expectedMergedResponse ,
2025-06-02 01:50:28 -07:00
) ;
} ) ;
2025-06-08 11:14:45 -07:00
2026-02-18 14:05:50 -08:00
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 ,
] ;
} ) ;
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.' ,
) ;
2026-02-26 18:26:16 -08:00
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 , endSpan : vi.fn ( ) } ) ;
} ) ;
expect ( metadata ) . toMatchObject ( {
input : sentParts ,
} ) ;
2026-02-18 14:05:50 -08:00
} ) ;
2025-06-08 11:14:45 -07:00
it ( 'should handle all tool calls being cancelled' , async ( ) = > {
2025-06-22 01:35:36 -04:00
const cancelledToolCalls : TrackedToolCall [ ] = [
2025-06-08 11:14:45 -07:00
{
2025-06-22 01:35:36 -04:00
request : {
2025-06-08 11:14:45 -07:00
callId : '1' ,
2025-06-22 01:35:36 -04:00
name : 'testTool' ,
args : { } ,
isClientInitiated : false ,
2025-07-10 00:19:30 +05:30
prompt_id : 'prompt-id-3' ,
2025-06-08 11:14:45 -07:00
} ,
2026-02-13 11:27:20 -05:00
status : CoreToolCallStatus.Cancelled ,
2025-08-20 15:51:31 -04:00
response : {
callId : '1' ,
2026-02-13 11:27:20 -05:00
responseParts : [ { text : CoreToolCallStatus.Cancelled } ] ,
2025-08-20 15:51:31 -04:00
errorType : undefined , // FIX: Added missing property
} ,
2025-06-08 11:14:45 -07:00
responseSubmittedToGemini : false ,
2025-08-06 10:50:02 -07:00
tool : {
displayName : 'mock tool' ,
} ,
invocation : {
getDescription : ( ) = > ` Mock description ` ,
} as unknown as AnyToolInvocation ,
2025-06-22 01:35:36 -04:00
} as TrackedCancelledToolCall ,
2025-06-08 11:14:45 -07:00
] ;
const client = new MockedGeminiClientClass ( mockConfig ) ;
2025-06-22 01:35:36 -04:00
2025-06-27 16:39:54 -07:00
// Capture the onComplete callback
let capturedOnComplete :
| ( ( completedTools : TrackedToolCall [ ] ) = > Promise < void > )
| null = null ;
2026-01-21 00:18:42 -05:00
mockUseToolScheduler . mockImplementation ( ( onComplete ) = > {
2025-06-27 16:39:54 -07:00
capturedOnComplete = onComplete ;
2026-01-27 16:06:24 -08:00
return [
[ ] ,
mockScheduleToolCalls ,
mockMarkToolsAsSubmitted ,
vi . fn ( ) ,
mockCancelAllToolCalls ,
0 ,
] ;
2025-06-27 16:39:54 -07:00
} ) ;
2026-01-27 16:06:24 -08:00
renderHookWithProviders ( ( ) = >
2025-06-22 01:35:36 -04:00
useGeminiStream (
client ,
[ ] ,
mockAddItem ,
mockConfig ,
2025-08-28 16:42:54 -07:00
mockLoadedSettings ,
2025-06-22 01:35:36 -04:00
mockOnDebugMessage ,
mockHandleSlashCommand ,
false ,
( ) = > 'vscode' as EditorType ,
( ) = > { } ,
( ) = > Promise . resolve ( ) ,
2025-07-09 13:55:56 -04:00
false ,
( ) = > { } ,
2025-08-06 15:19:10 -04:00
( ) = > { } ,
( ) = > { } ,
2025-09-11 13:27:27 -07:00
80 ,
24 ,
2025-06-22 01:35:36 -04:00
) ,
2025-06-08 11:14:45 -07:00
) ;
2025-06-27 16:39:54 -07:00
// Trigger the onComplete callback with cancelled tools
await act ( async ( ) = > {
if ( capturedOnComplete ) {
2026-01-27 16:06:24 -08:00
// Wait a tick for refs to be set up
await new Promise ( ( resolve ) = > setTimeout ( resolve , 0 ) ) ;
2025-06-27 16:39:54 -07:00
await capturedOnComplete ( cancelledToolCalls ) ;
}
2025-06-08 11:14:45 -07:00
} ) ;
2025-10-30 11:50:26 -07:00
await waitFor ( ( ) = > {
2025-06-22 01:35:36 -04:00
expect ( mockMarkToolsAsSubmitted ) . toHaveBeenCalledWith ( [ '1' ] ) ;
2025-06-08 11:14:45 -07:00
expect ( client . addHistory ) . toHaveBeenCalledWith ( {
role : 'user' ,
2026-02-13 11:27:20 -05:00
parts : [ { text : CoreToolCallStatus.Cancelled } ] ,
2025-06-08 11:14:45 -07:00
} ) ;
2025-06-22 01:35:36 -04:00
// Ensure we do NOT call back to the API
expect ( mockSendMessageStream ) . not . toHaveBeenCalled ( ) ;
2025-06-08 11:14:45 -07:00
} ) ;
} ) ;
2025-06-09 20:25:37 -04:00
2025-12-31 07:22:53 +08:00
it ( 'should stop agent execution immediately when a tool call returns STOP_EXECUTION error' , async ( ) = > {
const stopExecutionToolCalls : TrackedToolCall [ ] = [
{
request : {
callId : 'stop-call' ,
name : 'stopTool' ,
args : { } ,
isClientInitiated : false ,
prompt_id : 'prompt-id-stop' ,
} ,
2026-02-13 11:27:20 -05:00
status : CoreToolCallStatus.Error ,
2025-12-31 07:22:53 +08:00
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 ) ;
2026-01-27 16:06:24 -08:00
const { result } = renderTestHook ( [ ] , client ) ;
2025-12-31 07:22:53 +08:00
// Trigger the onComplete callback with STOP_EXECUTION tool
await act ( async ( ) = > {
if ( capturedOnComplete ) {
2026-01-27 16:06:24 -08:00
await capturedOnComplete ( stopExecutionToolCalls ) ;
2025-12-31 07:22:53 +08:00
}
} ) ;
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 ) ;
} ) ;
2026-02-27 14:15:10 -05:00
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 : TrackedToolCall [ ] = [
{
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 = {
. . . mockLoadedSettings ,
merged : {
. . . mockLoadedSettings . merged ,
ui : { errorVerbosity : 'low' } ,
} ,
} as LoadedSettings ;
const client = new MockedGeminiClientClass ( mockConfig ) ;
const { result } = 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 ) ;
2026-03-04 13:20:08 -08:00
// The failure hint should NOT be present if the suppressed error note was shown
expect ( failureHintIndex ) . toBe ( - 1 ) ;
2026-02-27 14:15:10 -05:00
expect ( noteIndex ) . toBeLessThan ( stopIndex ) ;
2025-12-31 07:22:53 +08:00
} ) ;
2025-07-05 13:56:39 -07:00
it ( 'should group multiple cancelled tool call responses into a single history entry' , async ( ) = > {
const cancelledToolCall1 : TrackedCancelledToolCall = {
request : {
callId : 'cancel-1' ,
name : 'toolA' ,
args : { } ,
isClientInitiated : false ,
2025-07-10 00:19:30 +05:30
prompt_id : 'prompt-id-7' ,
2025-07-05 13:56:39 -07:00
} ,
tool : {
name : 'toolA' ,
2025-08-06 10:50:02 -07:00
displayName : 'toolA' ,
2025-07-05 13:56:39 -07:00
description : 'descA' ,
2025-08-06 10:50:02 -07:00
build : vi.fn ( ) ,
2025-07-05 13:56:39 -07:00
} as any ,
2025-08-06 10:50:02 -07:00
invocation : {
getDescription : ( ) = > ` Mock description ` ,
} as unknown as AnyToolInvocation ,
2026-02-13 11:27:20 -05:00
status : CoreToolCallStatus.Cancelled ,
2025-07-05 13:56:39 -07:00
response : {
callId : 'cancel-1' ,
responseParts : [
{ functionResponse : { name : 'toolA' , id : 'cancel-1' } } ,
] ,
resultDisplay : undefined ,
error : undefined ,
2025-08-20 15:51:31 -04:00
errorType : undefined , // FIX: Added missing property
2025-07-05 13:56:39 -07:00
} ,
responseSubmittedToGemini : false ,
} ;
const cancelledToolCall2 : TrackedCancelledToolCall = {
request : {
callId : 'cancel-2' ,
name : 'toolB' ,
args : { } ,
isClientInitiated : false ,
2025-07-10 00:19:30 +05:30
prompt_id : 'prompt-id-8' ,
2025-07-05 13:56:39 -07:00
} ,
tool : {
name : 'toolB' ,
2025-08-06 10:50:02 -07:00
displayName : 'toolB' ,
2025-07-05 13:56:39 -07:00
description : 'descB' ,
2025-08-06 10:50:02 -07:00
build : vi.fn ( ) ,
2025-07-05 13:56:39 -07:00
} as any ,
2025-08-06 10:50:02 -07:00
invocation : {
getDescription : ( ) = > ` Mock description ` ,
} as unknown as AnyToolInvocation ,
2026-02-13 11:27:20 -05:00
status : CoreToolCallStatus.Cancelled ,
2025-07-05 13:56:39 -07:00
response : {
callId : 'cancel-2' ,
responseParts : [
{ functionResponse : { name : 'toolB' , id : 'cancel-2' } } ,
] ,
resultDisplay : undefined ,
error : undefined ,
2025-08-20 15:51:31 -04:00
errorType : undefined , // FIX: Added missing property
2025-07-05 13:56:39 -07:00
} ,
responseSubmittedToGemini : false ,
} ;
const allCancelledTools = [ cancelledToolCall1 , cancelledToolCall2 ] ;
const client = new MockedGeminiClientClass ( mockConfig ) ;
let capturedOnComplete :
| ( ( completedTools : TrackedToolCall [ ] ) = > Promise < void > )
| null = null ;
2026-01-21 00:18:42 -05:00
mockUseToolScheduler . mockImplementation ( ( onComplete ) = > {
2025-07-05 13:56:39 -07:00
capturedOnComplete = onComplete ;
2026-01-27 16:06:24 -08:00
return [
[ ] ,
mockScheduleToolCalls ,
mockMarkToolsAsSubmitted ,
vi . fn ( ) ,
mockCancelAllToolCalls ,
0 ,
] ;
2025-07-05 13:56:39 -07:00
} ) ;
2026-01-27 16:06:24 -08:00
renderHookWithProviders ( ( ) = >
2025-07-05 13:56:39 -07:00
useGeminiStream (
client ,
[ ] ,
mockAddItem ,
mockConfig ,
2025-08-28 16:42:54 -07:00
mockLoadedSettings ,
2025-07-05 13:56:39 -07:00
mockOnDebugMessage ,
mockHandleSlashCommand ,
false ,
( ) = > 'vscode' as EditorType ,
( ) = > { } ,
( ) = > Promise . resolve ( ) ,
2025-07-09 13:55:56 -04:00
false ,
( ) = > { } ,
2025-08-06 15:19:10 -04:00
( ) = > { } ,
( ) = > { } ,
2025-09-11 13:27:27 -07:00
80 ,
24 ,
2025-07-05 13:56:39 -07:00
) ,
) ;
// Trigger the onComplete callback with multiple cancelled tools
await act ( async ( ) = > {
if ( capturedOnComplete ) {
2026-01-27 16:06:24 -08:00
// Wait a tick for refs to be set up
await new Promise ( ( resolve ) = > setTimeout ( resolve , 0 ) ) ;
2025-07-05 13:56:39 -07:00
await capturedOnComplete ( allCancelledTools ) ;
}
} ) ;
2025-10-30 11:50:26 -07:00
await waitFor ( ( ) = > {
2025-07-05 13:56:39 -07:00
// The tools should be marked as submitted locally
expect ( mockMarkToolsAsSubmitted ) . toHaveBeenCalledWith ( [
'cancel-1' ,
'cancel-2' ,
] ) ;
// Crucially, addHistory should be called only ONCE
expect ( client . addHistory ) . toHaveBeenCalledTimes ( 1 ) ;
// And that single call should contain BOTH function responses
expect ( client . addHistory ) . toHaveBeenCalledWith ( {
role : 'user' ,
parts : [
2025-12-12 17:43:43 -08:00
. . . cancelledToolCall1 . response . responseParts ,
. . . cancelledToolCall2 . response . responseParts ,
2025-07-05 13:56:39 -07:00
] ,
} ) ;
// No message should be sent back to the API for a turn with only cancellations
expect ( mockSendMessageStream ) . not . toHaveBeenCalled ( ) ;
} ) ;
} ) ;
2025-06-19 18:25:23 -07:00
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 [ ] = [
{
2025-06-27 16:39:54 -07:00
request : {
callId : 'call1' ,
name : 'tool1' ,
args : { } ,
isClientInitiated : false ,
2025-07-10 00:19:30 +05:30
prompt_id : 'prompt-id-4' ,
2025-06-27 16:39:54 -07:00
} ,
2026-02-13 11:27:20 -05:00
status : CoreToolCallStatus.Executing ,
2025-06-19 18:25:23 -07:00
responseSubmittedToGemini : false ,
tool : {
name : 'tool1' ,
2025-08-06 10:50:02 -07:00
displayName : 'tool1' ,
2025-06-19 18:25:23 -07:00
description : 'desc' ,
2025-08-06 10:50:02 -07:00
build : vi.fn ( ) ,
2025-06-19 18:25:23 -07:00
} as any ,
2025-08-06 10:50:02 -07:00
invocation : {
getDescription : ( ) = > ` Mock description ` ,
} as unknown as AnyToolInvocation ,
2025-06-19 18:25:23 -07:00
startTime : Date.now ( ) ,
} as TrackedExecutingToolCall ,
] ;
const completedToolCalls : TrackedToolCall [ ] = [
{
. . . ( initialToolCalls [ 0 ] as TrackedExecutingToolCall ) ,
2026-02-13 11:27:20 -05:00
status : CoreToolCallStatus.Success ,
2025-06-19 18:25:23 -07:00
response : {
callId : 'call1' ,
responseParts : toolCallResponseParts ,
error : undefined ,
2025-08-20 15:51:31 -04:00
errorType : undefined , // FIX: Added missing property
2025-06-19 18:25:23 -07:00
resultDisplay : 'Tool 1 success display' ,
} ,
endTime : Date.now ( ) ,
} as TrackedCompletedToolCall ,
] ;
2025-06-27 16:39:54 -07:00
// Capture the onComplete callback
let capturedOnComplete :
| ( ( completedTools : TrackedToolCall [ ] ) = > Promise < void > )
| null = null ;
let currentToolCalls = initialToolCalls ;
2026-01-21 00:18:42 -05:00
mockUseToolScheduler . mockImplementation ( ( onComplete ) = > {
2025-06-27 16:39:54 -07:00
capturedOnComplete = onComplete ;
return [
currentToolCalls ,
mockScheduleToolCalls ,
mockMarkToolsAsSubmitted ,
2025-10-27 09:59:08 -07:00
vi . fn ( ) , // setToolCallsForDisplay
2026-01-27 16:06:24 -08:00
mockCancelAllToolCalls ,
0 ,
2025-06-27 16:39:54 -07:00
] ;
} ) ;
2026-01-27 16:06:24 -08:00
const { result , rerender } = renderHookWithProviders ( ( ) = >
2025-06-27 16:39:54 -07:00
useGeminiStream (
new MockedGeminiClientClass ( mockConfig ) ,
[ ] ,
mockAddItem ,
mockConfig ,
2025-08-28 16:42:54 -07:00
mockLoadedSettings ,
2025-06-27 16:39:54 -07:00
mockOnDebugMessage ,
mockHandleSlashCommand ,
false ,
( ) = > 'vscode' as EditorType ,
( ) = > { } ,
( ) = > Promise . resolve ( ) ,
2025-07-09 13:55:56 -04:00
false ,
( ) = > { } ,
2025-08-06 15:19:10 -04:00
( ) = > { } ,
( ) = > { } ,
2025-09-11 13:27:27 -07:00
80 ,
24 ,
2025-06-27 16:39:54 -07:00
) ,
) ;
2025-06-19 18:25:23 -07:00
// 1. Initial state should be Responding because a tool is executing.
expect ( result . current . streamingState ) . toBe ( StreamingState . Responding ) ;
2025-06-27 16:39:54 -07:00
// 2. Update the tool calls to completed state and rerender
currentToolCalls = completedToolCalls ;
2026-01-21 00:18:42 -05:00
mockUseToolScheduler . mockImplementation ( ( onComplete ) = > {
2025-06-27 16:39:54 -07:00
capturedOnComplete = onComplete ;
return [
completedToolCalls ,
mockScheduleToolCalls ,
mockMarkToolsAsSubmitted ,
2025-10-27 09:59:08 -07:00
vi . fn ( ) , // setToolCallsForDisplay
2026-01-27 16:06:24 -08:00
mockCancelAllToolCalls ,
0 ,
2025-06-27 16:39:54 -07:00
] ;
} ) ;
2025-06-19 18:25:23 -07:00
act ( ( ) = > {
2025-06-27 16:39:54 -07:00
rerender ( ) ;
2025-06-19 18:25:23 -07:00
} ) ;
// 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 ) ;
2025-06-27 16:39:54 -07:00
// 4. Trigger the onComplete callback to simulate tool completion
await act ( async ( ) = > {
if ( capturedOnComplete ) {
2026-01-27 16:06:24 -08:00
// Wait a tick for refs to be set up
await new Promise ( ( resolve ) = > setTimeout ( resolve , 0 ) ) ;
2025-06-27 16:39:54 -07:00
await capturedOnComplete ( completedToolCalls ) ;
}
} ) ;
// 5. Wait for submitQuery to be called
2025-10-30 11:50:26 -07:00
await waitFor ( ( ) = > {
2025-06-19 18:25:23 -07:00
expect ( mockSendMessageStream ) . toHaveBeenCalledWith (
toolCallResponseParts ,
expect . any ( AbortSignal ) ,
2025-07-10 00:19:30 +05:30
'prompt-id-4' ,
2026-01-30 10:09:27 -08:00
undefined ,
false ,
toolCallResponseParts ,
2025-06-19 18:25:23 -07:00
) ;
} ) ;
2025-06-27 16:39:54 -07:00
// 6. After submission, the state should remain Responding until the stream completes.
2025-06-19 18:25:23 -07:00
expect ( result . current . streamingState ) . toBe ( StreamingState . Responding ) ;
} ) ;
2025-06-20 23:01:44 -04:00
describe ( 'User Cancellation' , ( ) = > {
2025-08-12 14:05:49 -07:00
let keypressCallback : ( key : any ) = > void ;
const mockUseKeypress = useKeypress as Mock ;
2025-06-20 23:01:44 -04:00
beforeEach ( ( ) = > {
2025-08-12 14:05:49 -07:00
// Capture the callback passed to useKeypress
mockUseKeypress . mockImplementation ( ( callback , options ) = > {
if ( options . isActive ) {
keypressCallback = callback ;
} else {
keypressCallback = ( ) = > { } ;
}
2025-06-20 23:01:44 -04:00
} ) ;
} ) ;
const simulateEscapeKeyPress = ( ) = > {
act ( ( ) = > {
2025-08-12 14:05:49 -07:00
keypressCallback ( { name : 'escape' } ) ;
2025-06-20 23:01:44 -04:00
} ) ;
} ;
it ( 'should cancel an in-progress stream when escape is pressed' , async ( ) = > {
const mockStream = ( async function * ( ) {
yield { type : 'content' , value : 'Part 1' } ;
// Keep the stream open
await new Promise ( ( ) = > { } ) ;
} ) ( ) ;
mockSendMessageStream . mockReturnValue ( mockStream ) ;
const { result } = renderTestHook ( ) ;
// Start a query
await act ( async ( ) = > {
2025-12-05 16:12:49 -08:00
// eslint-disable-next-line @typescript-eslint/no-floating-promises
2025-06-20 23:01:44 -04:00
result . current . submitQuery ( 'test query' ) ;
} ) ;
// Wait for the first part of the response
2025-10-30 11:50:26 -07:00
await waitFor ( ( ) = > {
2025-06-20 23:01:44 -04:00
expect ( result . current . streamingState ) . toBe ( StreamingState . Responding ) ;
} ) ;
// Simulate escape key press
simulateEscapeKeyPress ( ) ;
// Verify cancellation message is added
2025-10-30 11:50:26 -07:00
await waitFor ( ( ) = > {
2026-01-13 14:15:04 -05:00
expect ( mockAddItem ) . toHaveBeenCalledWith ( {
type : MessageType . INFO ,
text : 'Request cancelled.' ,
} ) ;
2025-06-20 23:01:44 -04:00
} ) ;
// Verify state is reset
expect ( result . current . streamingState ) . toBe ( StreamingState . Idle ) ;
} ) ;
2025-08-06 15:19:10 -04:00
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 ) ;
2026-01-27 16:06:24 -08:00
const { result } = renderHookWithProviders ( ( ) = >
2025-08-06 15:19:10 -04:00
useGeminiStream (
mockConfig . getGeminiClient ( ) ,
[ ] ,
mockAddItem ,
mockConfig ,
2025-08-28 16:42:54 -07:00
mockLoadedSettings ,
2025-08-06 15:19:10 -04:00
mockOnDebugMessage ,
mockHandleSlashCommand ,
false ,
( ) = > 'vscode' as EditorType ,
( ) = > { } ,
( ) = > Promise . resolve ( ) ,
false ,
( ) = > { } ,
cancelSubmitSpy ,
2025-09-11 13:27:27 -07:00
( ) = > { } ,
80 ,
24 ,
2025-08-06 15:19:10 -04:00
) ,
) ;
// Start a query
await act ( async ( ) = > {
2025-12-05 16:12:49 -08:00
// eslint-disable-next-line @typescript-eslint/no-floating-promises
2025-08-06 15:19:10 -04:00
result . current . submitQuery ( 'test query' ) ;
} ) ;
simulateEscapeKeyPress ( ) ;
2025-11-19 11:11:36 +08:00
expect ( cancelSubmitSpy ) . toHaveBeenCalledWith ( false ) ;
2025-08-06 15:19:10 -04:00
} ) ;
2025-09-11 13:27:27 -07:00
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 ) ;
2026-01-27 16:06:24 -08:00
const { result } = renderHookWithProviders ( ( ) = >
2025-09-11 13:27:27 -07:00
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 ( ) = > {
2025-12-05 16:12:49 -08:00
// eslint-disable-next-line @typescript-eslint/no-floating-promises
2025-09-11 13:27:27 -07:00
result . current . submitQuery ( 'test query' ) ;
} ) ;
simulateEscapeKeyPress ( ) ;
expect ( setShellInputFocusedSpy ) . toHaveBeenCalledWith ( false ) ;
} ) ;
2025-06-20 23:01:44 -04:00
it ( 'should not do anything if escape is pressed when not responding' , ( ) = > {
const { result } = renderTestHook ( ) ;
expect ( result . current . streamingState ) . toBe ( StreamingState . Idle ) ;
// Simulate escape key press
simulateEscapeKeyPress ( ) ;
// No change should happen, no cancellation message
expect ( mockAddItem ) . not . toHaveBeenCalledWith (
expect . objectContaining ( {
text : 'Request cancelled.' ,
} ) ,
) ;
} ) ;
it ( 'should prevent further processing after cancellation' , async ( ) = > {
let continueStream : ( ) = > void ;
const streamPromise = new Promise < void > ( ( resolve ) = > {
continueStream = resolve ;
} ) ;
const mockStream = ( async function * ( ) {
yield { type : 'content' , value : 'Initial' } ;
await streamPromise ; // Wait until we manually continue
yield { type : 'content' , value : ' Canceled' } ;
} ) ( ) ;
mockSendMessageStream . mockReturnValue ( mockStream ) ;
const { result } = renderTestHook ( ) ;
await act ( async ( ) = > {
2025-12-05 16:12:49 -08:00
// eslint-disable-next-line @typescript-eslint/no-floating-promises
2025-06-20 23:01:44 -04:00
result . current . submitQuery ( 'long running query' ) ;
} ) ;
2025-10-30 11:50:26 -07:00
await waitFor ( ( ) = > {
2025-06-20 23:01:44 -04:00
expect ( result . current . streamingState ) . toBe ( StreamingState . Responding ) ;
} ) ;
// Cancel the request
simulateEscapeKeyPress ( ) ;
// Allow the stream to continue
2025-10-30 11:50:26 -07:00
await act ( async ( ) = > {
2025-06-20 23:01:44 -04:00
continueStream ( ) ;
2025-10-30 11:50:26 -07:00
// Wait a bit to see if the second part is processed
await new Promise ( ( resolve ) = > setTimeout ( resolve , 50 ) ) ;
2025-06-20 23:01:44 -04:00
} ) ;
// 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 ) ;
} ) ;
2025-10-27 09:59:08 -07:00
it ( 'should cancel if a tool call is in progress' , async ( ) = > {
2025-06-20 23:01:44 -04:00
const toolCalls : TrackedToolCall [ ] = [
{
request : { callId : 'call1' , name : 'tool1' , args : { } } ,
2026-02-13 11:27:20 -05:00
status : CoreToolCallStatus.Executing ,
2025-06-20 23:01:44 -04:00
responseSubmittedToGemini : false ,
tool : {
name : 'tool1' ,
description : 'desc1' ,
2025-08-06 10:50:02 -07:00
build : vi.fn ( ) . mockImplementation ( ( _ ) = > ( {
getDescription : ( ) = > ` Mock description ` ,
} ) ) ,
2025-06-20 23:01:44 -04:00
} as any ,
2025-08-06 10:50:02 -07:00
invocation : {
getDescription : ( ) = > ` Mock description ` ,
} ,
2025-06-20 23:01:44 -04:00
startTime : Date.now ( ) ,
liveOutput : '...' ,
} as TrackedExecutingToolCall ,
] ;
const { result } = renderTestHook ( toolCalls ) ;
// State is `Responding` because a tool is running
expect ( result . current . streamingState ) . toBe ( StreamingState . Responding ) ;
// Try to cancel
simulateEscapeKeyPress ( ) ;
2025-10-27 09:59:08 -07:00
// The cancel function should be called
expect ( mockCancelAllToolCalls ) . toHaveBeenCalled ( ) ;
} ) ;
it ( 'should cancel a request when a tool is awaiting confirmation' , async ( ) = > {
const mockOnConfirm = vi . fn ( ) . mockResolvedValue ( undefined ) ;
const toolCalls : TrackedToolCall [ ] = [
{
request : {
callId : 'confirm-call' ,
name : 'some_tool' ,
args : { } ,
isClientInitiated : false ,
prompt_id : 'prompt-id-1' ,
} ,
2026-02-13 11:27:20 -05:00
status : CoreToolCallStatus.AwaitingApproval ,
2025-10-27 09:59:08 -07:00
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 } = 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
2025-10-30 11:50:26 -07:00
await waitFor ( ( ) = > {
2025-10-27 09:59:08 -07:00
expect ( mockAddItem ) . toHaveBeenCalledWith (
expect . objectContaining ( {
text : 'Request cancelled.' ,
} ) ,
) ;
} ) ;
// The final state should be idle
expect ( result . current . streamingState ) . toBe ( StreamingState . Idle ) ;
2025-06-20 23:01:44 -04:00
} ) ;
} ) ;
2025-06-22 01:35:36 -04:00
2026-01-13 23:03:19 -05:00
describe ( 'Retry Handling' , ( ) = > {
it ( 'should update retryStatus when CoreEvent.RetryAttempt is emitted' , async ( ) = > {
const { result } = 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 ) . toEqual ( retryPayload ) ;
} ) ;
it ( 'should reset retryStatus when isResponding becomes false' , async ( ) = > {
const { result } = 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 ( ) ;
} ) ;
} ) ;
2025-07-07 16:45:44 -04:00
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' ,
2025-06-22 01:35:36 -04:00
toolName : 'save_memory' ,
toolArgs : { fact : 'test fact' } ,
} ;
mockHandleSlashCommand . mockResolvedValue ( clientToolRequest ) ;
2025-07-07 16:45:44 -04:00
const { result } = renderTestHook ( ) ;
2025-06-22 01:35:36 -04:00
2025-07-07 16:45:44 -04:00
await act ( async ( ) = > {
await result . current . submitQuery ( '/memory add "test fact"' ) ;
} ) ;
2025-06-22 01:35:36 -04:00
2025-10-30 11:50:26 -07:00
await waitFor ( ( ) = > {
2025-07-07 16:45:44 -04:00
expect ( mockScheduleToolCalls ) . toHaveBeenCalledWith (
[
expect . objectContaining ( {
name : 'save_memory' ,
args : { fact : 'test fact' } ,
isClientInitiated : true ,
} ) ,
] ,
expect . any ( AbortSignal ) ,
) ;
expect ( mockSendMessageStream ) . not . toHaveBeenCalled ( ) ;
2025-06-27 16:39:54 -07:00
} ) ;
2025-07-07 16:45:44 -04:00
} ) ;
2025-06-27 16:39:54 -07:00
2025-07-07 16:45:44 -04:00
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 ) ;
2025-06-22 01:35:36 -04:00
2025-07-07 16:45:44 -04:00
const { result } = renderTestHook ( ) ;
2025-06-22 01:35:36 -04:00
2025-06-27 16:39:54 -07:00
await act ( async ( ) = > {
2025-07-07 16:45:44 -04:00
await result . current . submitQuery ( '/help' ) ;
2025-06-22 01:35:36 -04:00
} ) ;
2025-10-30 11:50:26 -07:00
await waitFor ( ( ) = > {
2025-07-07 16:45:44 -04:00
expect ( mockHandleSlashCommand ) . toHaveBeenCalledWith ( '/help' ) ;
expect ( mockScheduleToolCalls ) . not . toHaveBeenCalled ( ) ;
expect ( mockSendMessageStream ) . not . toHaveBeenCalled ( ) ; // No LLM call made
2025-06-22 01:35:36 -04:00
} ) ;
} ) ;
2025-07-22 00:34:55 -04:00
it ( 'should call Gemini with prompt content when slash command returns a `submit_prompt` action' , async ( ) = > {
const customCommandResult : SlashCommandProcessorResult = {
type : 'submit_prompt' ,
content : 'This is the actual prompt from the command file.' ,
} ;
mockHandleSlashCommand . mockResolvedValue ( customCommandResult ) ;
const { result , mockSendMessageStream : localMockSendMessageStream } =
renderTestHook ( ) ;
await act ( async ( ) = > {
await result . current . submitQuery ( '/my-custom-command' ) ;
} ) ;
2025-10-30 11:50:26 -07:00
await waitFor ( ( ) = > {
2025-07-22 00:34:55 -04:00
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 ) ,
2026-01-30 10:09:27 -08:00
undefined ,
false ,
'/my-custom-command' ,
2025-07-22 00:34:55 -04:00
) ;
expect ( mockScheduleToolCalls ) . not . toHaveBeenCalled ( ) ;
} ) ;
} ) ;
it ( 'should correctly handle a submit_prompt action with empty content' , async ( ) = > {
const emptyPromptResult : SlashCommandProcessorResult = {
type : 'submit_prompt' ,
content : '' ,
} ;
mockHandleSlashCommand . mockResolvedValue ( emptyPromptResult ) ;
const { result , mockSendMessageStream : localMockSendMessageStream } =
renderTestHook ( ) ;
await act ( async ( ) = > {
await result . current . submitQuery ( '/emptycmd' ) ;
} ) ;
2025-10-30 11:50:26 -07:00
await waitFor ( ( ) = > {
2025-07-22 00:34:55 -04:00
expect ( mockHandleSlashCommand ) . toHaveBeenCalledWith ( '/emptycmd' ) ;
expect ( localMockSendMessageStream ) . toHaveBeenCalledWith (
'' ,
expect . any ( AbortSignal ) ,
expect . any ( String ) ,
2026-01-30 10:09:27 -08:00
undefined ,
false ,
'/emptycmd' ,
2025-07-22 00:34:55 -04:00
) ;
} ) ;
} ) ;
2025-08-26 11:51:27 +08:00
it ( 'should not call handleSlashCommand for line comments' , async ( ) = > {
const { result , mockSendMessageStream : localMockSendMessageStream } =
renderTestHook ( ) ;
await act ( async ( ) = > {
await result . current . submitQuery ( '// This is a line comment' ) ;
} ) ;
2025-10-30 11:50:26 -07:00
await waitFor ( ( ) = > {
2025-08-26 11:51:27 +08:00
expect ( mockHandleSlashCommand ) . not . toHaveBeenCalled ( ) ;
expect ( localMockSendMessageStream ) . toHaveBeenCalledWith (
'// This is a line comment' ,
expect . any ( AbortSignal ) ,
expect . any ( String ) ,
2026-01-30 10:09:27 -08:00
undefined ,
false ,
'// This is a line comment' ,
2025-08-26 11:51:27 +08:00
) ;
} ) ;
} ) ;
it ( 'should not call handleSlashCommand for block comments' , async ( ) = > {
const { result , mockSendMessageStream : localMockSendMessageStream } =
renderTestHook ( ) ;
await act ( async ( ) = > {
await result . current . submitQuery ( '/* This is a block comment */' ) ;
} ) ;
2025-10-30 11:50:26 -07:00
await waitFor ( ( ) = > {
2025-08-26 11:51:27 +08:00
expect ( mockHandleSlashCommand ) . not . toHaveBeenCalled ( ) ;
expect ( localMockSendMessageStream ) . toHaveBeenCalledWith (
'/* This is a block comment */' ,
expect . any ( AbortSignal ) ,
expect . any ( String ) ,
2026-01-30 10:09:27 -08:00
undefined ,
false ,
'/* This is a block comment */' ,
2025-08-26 11:51:27 +08:00
) ;
} ) ;
} ) ;
2025-10-17 23:00:27 +05:30
it ( 'should not call handleSlashCommand is shell mode is active' , async ( ) = > {
2026-01-27 16:06:24 -08:00
const { result } = renderHookWithProviders ( ( ) = >
2025-10-17 23:00:27 +05:30
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' ) ;
} ) ;
2025-10-30 11:50:26 -07:00
await waitFor ( ( ) = > {
2025-10-17 23:00:27 +05:30
expect ( mockHandleSlashCommand ) . not . toHaveBeenCalled ( ) ;
} ) ;
} ) ;
2025-06-22 01:35:36 -04:00
} ) ;
describe ( 'Memory Refresh on save_memory' , ( ) = > {
it ( 'should call performMemoryRefresh when a save_memory tool call completes successfully' , async ( ) = > {
const mockPerformMemoryRefresh = vi . fn ( ) ;
const completedToolCall : TrackedCompletedToolCall = {
request : {
callId : 'save-mem-call-1' ,
name : 'save_memory' ,
args : { fact : 'test' } ,
isClientInitiated : true ,
2025-07-10 00:19:30 +05:30
prompt_id : 'prompt-id-6' ,
2025-06-22 01:35:36 -04:00
} ,
2026-02-13 11:27:20 -05:00
status : CoreToolCallStatus.Success ,
2025-06-22 01:35:36 -04:00
responseSubmittedToGemini : false ,
response : {
callId : 'save-mem-call-1' ,
responseParts : [ { text : 'Memory saved' } ] ,
resultDisplay : 'Success: Memory saved' ,
error : undefined ,
2025-08-20 15:51:31 -04:00
errorType : undefined , // FIX: Added missing property
2025-06-22 01:35:36 -04:00
} ,
tool : {
name : 'save_memory' ,
2025-08-06 10:50:02 -07:00
displayName : 'save_memory' ,
2025-06-22 01:35:36 -04:00
description : 'Saves memory' ,
2025-08-06 10:50:02 -07:00
build : vi.fn ( ) ,
2025-06-22 01:35:36 -04:00
} as any ,
2025-08-06 10:50:02 -07:00
invocation : {
getDescription : ( ) = > ` Mock description ` ,
} as unknown as AnyToolInvocation ,
2025-06-22 01:35:36 -04:00
} ;
2025-06-27 16:39:54 -07:00
// Capture the onComplete callback
let capturedOnComplete :
| ( ( completedTools : TrackedToolCall [ ] ) = > Promise < void > )
| null = null ;
2025-06-22 01:35:36 -04:00
2026-01-21 00:18:42 -05:00
mockUseToolScheduler . mockImplementation ( ( onComplete ) = > {
2025-06-27 16:39:54 -07:00
capturedOnComplete = onComplete ;
2026-01-27 16:06:24 -08:00
return [
[ ] ,
mockScheduleToolCalls ,
mockMarkToolsAsSubmitted ,
vi . fn ( ) ,
mockCancelAllToolCalls ,
0 ,
] ;
2025-06-27 16:39:54 -07:00
} ) ;
2026-01-27 16:06:24 -08:00
renderHookWithProviders ( ( ) = >
2025-06-22 01:35:36 -04:00
useGeminiStream (
new MockedGeminiClientClass ( mockConfig ) ,
[ ] ,
mockAddItem ,
mockConfig ,
2025-08-28 16:42:54 -07:00
mockLoadedSettings ,
2025-06-22 01:35:36 -04:00
mockOnDebugMessage ,
mockHandleSlashCommand ,
false ,
( ) = > 'vscode' as EditorType ,
( ) = > { } ,
mockPerformMemoryRefresh ,
2025-07-09 13:55:56 -04:00
false ,
( ) = > { } ,
2025-08-06 15:19:10 -04:00
( ) = > { } ,
( ) = > { } ,
2025-09-11 13:27:27 -07:00
80 ,
24 ,
2025-06-22 01:35:36 -04:00
) ,
) ;
2025-06-27 16:39:54 -07:00
// Trigger the onComplete callback with the completed save_memory tool
await act ( async ( ) = > {
if ( capturedOnComplete ) {
2026-01-27 16:06:24 -08:00
// Wait a tick for refs to be set up
await new Promise ( ( resolve ) = > setTimeout ( resolve , 0 ) ) ;
2025-06-27 16:39:54 -07:00
await capturedOnComplete ( [ completedToolCall ] ) ;
}
2025-06-22 01:35:36 -04:00
} ) ;
2025-10-30 11:50:26 -07:00
await waitFor ( ( ) = > {
2025-06-22 01:35:36 -04:00
expect ( mockPerformMemoryRefresh ) . toHaveBeenCalledTimes ( 1 ) ;
} ) ;
} ) ;
} ) ;
2025-06-23 23:43:00 -04:00
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' ) ;
2025-06-30 17:11:54 -07:00
const mockAuthType = AuthType . LOGIN_WITH_GOOGLE ;
2025-06-23 23:43:00 -04:00
mockParseAndFormatApiError . mockClear ( ) ;
mockSendMessageStream . mockReturnValue (
( async function * ( ) {
yield { type : 'content' , value : '' } ;
throw mockError ;
} ) ( ) ,
) ;
const testConfig = {
. . . mockConfig ,
2026-01-27 13:17:40 -08:00
getContentGenerator : vi.fn ( ) ,
2025-06-23 23:43:00 -04:00
getContentGeneratorConfig : vi.fn ( ( ) = > ( {
authType : mockAuthType ,
} ) ) ,
2025-07-09 10:18:15 -04:00
getModel : vi.fn ( ( ) = > 'gemini-2.5-pro' ) ,
2025-06-23 23:43:00 -04:00
} as unknown as Config ;
2026-01-27 16:06:24 -08:00
const { result } = renderHookWithProviders ( ( ) = >
2025-06-23 23:43:00 -04:00
useGeminiStream (
new MockedGeminiClientClass ( testConfig ) ,
[ ] ,
mockAddItem ,
testConfig ,
2025-08-28 16:42:54 -07:00
mockLoadedSettings ,
2025-06-23 23:43:00 -04:00
mockOnDebugMessage ,
mockHandleSlashCommand ,
false ,
( ) = > 'vscode' as EditorType ,
( ) = > { } ,
( ) = > Promise . resolve ( ) ,
2025-07-09 13:55:56 -04:00
false ,
( ) = > { } ,
2025-08-06 15:19:10 -04:00
( ) = > { } ,
( ) = > { } ,
2025-09-11 13:27:27 -07:00
80 ,
24 ,
2025-06-23 23:43:00 -04:00
) ,
) ;
// 2. Action
await act ( async ( ) = > {
await result . current . submitQuery ( 'test query' ) ;
} ) ;
// 3. Assertion
2025-10-30 11:50:26 -07:00
await waitFor ( ( ) = > {
2025-06-23 23:43:00 -04:00
expect ( mockParseAndFormatApiError ) . toHaveBeenCalledWith (
'Rate limit exceeded' ,
mockAuthType ,
2025-07-09 10:18:15 -04:00
undefined ,
'gemini-2.5-pro' ,
'gemini-2.5-flash' ,
2025-06-23 23:43:00 -04:00
) ;
} ) ;
} ) ;
} ) ;
2025-07-22 06:57:11 +09:00
2025-09-14 20:20:21 -07:00
describe ( 'handleApprovalModeChange' , ( ) = > {
it ( 'should auto-approve all pending tool calls when switching to YOLO mode' , async ( ) = > {
const awaitingApprovalToolCalls : TrackedToolCall [ ] = [
2026-02-13 11:14:35 +09:00
createMockToolCall ( 'replace' , 'call1' , 'edit' ) ,
createMockToolCall ( 'read_file' , 'call2' , 'info' ) ,
2025-09-14 20:20:21 -07:00
] ;
const { result } = renderTestHook ( awaitingApprovalToolCalls ) ;
await act ( async ( ) = > {
await result . current . handleApprovalModeChange ( ApprovalMode . YOLO ) ;
} ) ;
// Both tool calls should be auto-approved
2026-02-13 11:14:35 +09:00
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 ,
} ) ,
2025-09-14 20:20:21 -07:00
) ;
} ) ;
it ( 'should only auto-approve edit tools when switching to AUTO_EDIT mode' , async ( ) = > {
const awaitingApprovalToolCalls : TrackedToolCall [ ] = [
2026-02-13 11:14:35 +09:00
createMockToolCall ( 'replace' , 'call1' , 'edit' ) ,
createMockToolCall ( 'write_file' , 'call2' , 'edit' ) ,
createMockToolCall ( 'read_file' , 'call3' , 'info' ) ,
2025-09-14 20:20:21 -07:00
] ;
const { result } = renderTestHook ( awaitingApprovalToolCalls ) ;
await act ( async ( ) = > {
await result . current . handleApprovalModeChange ( ApprovalMode . AUTO_EDIT ) ;
} ) ;
// Only replace and write_file should be auto-approved
2026-02-13 11:14:35 +09:00
expect ( mockMessageBus . publish ) . toHaveBeenCalledTimes ( 2 ) ;
expect ( mockMessageBus . publish ) . toHaveBeenCalledWith (
expect . objectContaining ( { correlationId : 'corr-call1' } ) ,
2025-09-14 20:20:21 -07:00
) ;
2026-02-13 11:14:35 +09:00
expect ( mockMessageBus . publish ) . toHaveBeenCalledWith (
expect . objectContaining ( { correlationId : 'corr-call2' } ) ,
) ;
expect ( mockMessageBus . publish ) . not . toHaveBeenCalledWith (
expect . objectContaining ( { correlationId : 'corr-call3' } ) ,
2025-09-14 20:20:21 -07:00
) ;
} ) ;
it ( 'should not auto-approve any tools when switching to REQUIRE_CONFIRMATION mode' , async ( ) = > {
const awaitingApprovalToolCalls : TrackedToolCall [ ] = [
2026-02-13 11:14:35 +09:00
createMockToolCall ( 'replace' , 'call1' , 'edit' ) ,
2025-09-14 20:20:21 -07:00
] ;
const { result } = renderTestHook ( awaitingApprovalToolCalls ) ;
await act ( async ( ) = > {
2025-10-19 17:16:16 -07:00
await result . current . handleApprovalModeChange ( ApprovalMode . DEFAULT ) ;
2025-09-14 20:20:21 -07:00
} ) ;
// No tools should be auto-approved
2026-02-13 11:14:35 +09:00
expect ( mockMessageBus . publish ) . not . toHaveBeenCalled ( ) ;
2025-09-14 20:20:21 -07:00
} ) ;
it ( 'should handle errors gracefully when auto-approving tool calls' , async ( ) = > {
2025-10-28 19:05:48 +00:00
const debuggerSpy = vi
. spyOn ( debugLogger , 'warn' )
2025-09-14 20:20:21 -07:00
. mockImplementation ( ( ) = > { } ) ;
2026-02-13 11:14:35 +09:00
mockMessageBus . publish . mockRejectedValueOnce ( new Error ( 'Bus error' ) ) ;
2025-09-14 20:20:21 -07:00
const awaitingApprovalToolCalls : TrackedToolCall [ ] = [
2026-02-13 11:14:35 +09:00
createMockToolCall ( 'replace' , 'call1' , 'edit' ) ,
createMockToolCall ( 'write_file' , 'call2' , 'edit' ) ,
2025-09-14 20:20:21 -07:00
] ;
const { result } = renderTestHook ( awaitingApprovalToolCalls ) ;
await act ( async ( ) = > {
await result . current . handleApprovalModeChange ( ApprovalMode . YOLO ) ;
} ) ;
2026-02-13 11:14:35 +09:00
// Both should be attempted despite first error
expect ( mockMessageBus . publish ) . toHaveBeenCalledTimes ( 2 ) ;
2025-10-28 19:05:48 +00:00
expect ( debuggerSpy ) . toHaveBeenCalledWith (
2026-02-13 11:14:35 +09:00
'Failed to auto-approve tool call call1:' ,
2025-09-14 20:20:21 -07:00
expect . any ( Error ) ,
) ;
2025-10-28 19:05:48 +00:00
debuggerSpy . mockRestore ( ) ;
2025-09-14 20:20:21 -07:00
} ) ;
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' ,
} ,
2026-02-13 11:27:20 -05:00
status : CoreToolCallStatus.AwaitingApproval ,
2025-09-14 20:20:21 -07:00
responseSubmittedToGemini : false ,
// No confirmationDetails
tool : {
name : 'replace' ,
displayName : 'replace' ,
description : 'Replace text' ,
build : vi.fn ( ) ,
} as any ,
invocation : {
getDescription : ( ) = > 'Mock description' ,
} as unknown as AnyToolInvocation ,
2026-02-13 11:14:35 +09:00
correlationId : 'corr-1' ,
2025-10-19 17:16:16 -07:00
} as unknown as TrackedWaitingToolCall ,
2025-09-14 20:20:21 -07:00
] ;
const { result } = renderTestHook ( awaitingApprovalToolCalls ) ;
// Should not throw an error
await act ( async ( ) = > {
await result . current . handleApprovalModeChange ( ApprovalMode . YOLO ) ;
} ) ;
} ) ;
it ( 'should only process tool calls with awaiting_approval status' , async ( ) = > {
2026-02-13 11:27:20 -05:00
const mockOnConfirmAwaiting = vi . fn ( ) . mockResolvedValue ( undefined ) ;
2025-09-14 20:20:21 -07:00
const mixedStatusToolCalls : TrackedToolCall [ ] = [
2026-02-13 11:27:20 -05:00
createMockToolCall (
'replace' ,
'call1' ,
'edit' ,
CoreToolCallStatus . AwaitingApproval ,
mockOnConfirmAwaiting ,
) ,
2025-09-14 20:20:21 -07:00
{
request : {
callId : 'call2' ,
name : 'write_file' ,
args : { path : '/test/file.txt' , content : 'content' } ,
isClientInitiated : false ,
prompt_id : 'prompt-id-1' ,
} ,
2026-02-13 11:27:20 -05:00
status : CoreToolCallStatus.Executing ,
2025-09-14 20:20:21 -07:00
responseSubmittedToGemini : false ,
tool : {
name : 'write_file' ,
displayName : 'write_file' ,
description : 'Write file' ,
build : vi.fn ( ) ,
} as any ,
invocation : {
getDescription : ( ) = > 'Mock description' ,
} as unknown as AnyToolInvocation ,
startTime : Date.now ( ) ,
liveOutput : 'Writing...' ,
2026-02-13 11:14:35 +09:00
correlationId : 'corr-call2' ,
2025-09-14 20:20:21 -07:00
} as TrackedExecutingToolCall ,
] ;
const { result } = renderTestHook ( mixedStatusToolCalls ) ;
await act ( async ( ) = > {
await result . current . handleApprovalModeChange ( ApprovalMode . YOLO ) ;
} ) ;
2026-02-13 11:14:35 +09:00
// 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' } ) ,
) ;
2025-09-14 20:20:21 -07:00
} ) ;
2026-02-24 14:31:41 -05:00
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 } = 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 ) ,
} ,
] ,
} ) ;
} ) ;
2025-09-14 20:20:21 -07:00
} ) ;
2025-07-22 06:57:11 +09:00
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...' ,
} ;
2025-09-02 23:29:07 -06:00
yield {
type : ServerGeminiEventType . Finished ,
value : { reason : 'MAX_TOKENS' , usageMetadata : undefined } ,
} ;
2025-07-22 06:57:11 +09:00
} ) ( ) ,
) ;
2026-01-27 16:06:24 -08:00
const { result } = renderHookWithProviders ( ( ) = >
2025-07-22 06:57:11 +09:00
useGeminiStream (
new MockedGeminiClientClass ( mockConfig ) ,
[ ] ,
mockAddItem ,
mockConfig ,
2025-08-28 16:42:54 -07:00
mockLoadedSettings ,
2025-07-22 06:57:11 +09:00
mockOnDebugMessage ,
mockHandleSlashCommand ,
false ,
( ) = > 'vscode' as EditorType ,
( ) = > { } ,
( ) = > Promise . resolve ( ) ,
false ,
( ) = > { } ,
2025-08-06 15:19:10 -04:00
( ) = > { } ,
( ) = > { } ,
2025-09-11 13:27:27 -07:00
80 ,
24 ,
2025-07-22 06:57:11 +09:00
) ,
) ;
// Submit a query
await act ( async ( ) = > {
await result . current . submitQuery ( 'Generate long text' ) ;
} ) ;
// Check that the info message was added
2025-10-30 11:50:26 -07:00
await waitFor ( ( ) = > {
2025-07-22 06:57:11 +09:00
expect ( mockAddItem ) . toHaveBeenCalledWith (
{
type : 'info' ,
text : '⚠️ Response truncated due to token limits.' ,
} ,
expect . any ( Number ) ,
) ;
} ) ;
} ) ;
2025-10-09 10:22:26 -07:00
describe ( 'ContextWindowWillOverflow event' , ( ) = > {
beforeEach ( ( ) = > {
vi . mocked ( tokenLimit ) . mockReturnValue ( 100 ) ;
} ) ;
2025-10-08 15:20:44 -07:00
2025-11-03 23:40:57 +05:30
it . each ( [
{
name : 'without suggestion when remaining tokens are > 75% of limit' ,
requestTokens : 20 ,
remainingTokens : 80 ,
expectedMessage :
2026-03-03 01:22:29 -08:00
'Sending this message (20 tokens) might exceed the context window limit (80 tokens left).' ,
2025-11-03 23:40:57 +05:30
} ,
{
name : 'with suggestion when remaining tokens are < 75% of limit' ,
requestTokens : 30 ,
remainingTokens : 70 ,
expectedMessage :
2026-03-03 01:22:29 -08:00
'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.' ,
2025-11-03 23:40:57 +05:30
} ,
] ) (
'should add message $name' ,
async ( { requestTokens , remainingTokens , expectedMessage } ) = > {
mockSendMessageStream . mockReturnValue (
( async function * ( ) {
yield {
type : ServerGeminiEventType . ContextWindowWillOverflow ,
value : {
estimatedRequestTokenCount : requestTokens ,
remainingTokenCount : remainingTokens ,
} ,
} ;
} ) ( ) ,
2025-10-09 10:22:26 -07:00
) ;
2025-10-08 15:20:44 -07:00
2025-11-03 23:40:57 +05:30
const { result } = renderHookWithDefaults ( ) ;
2025-10-09 10:22:26 -07:00
2025-11-03 23:40:57 +05:30
await act ( async ( ) = > {
await result . current . submitQuery ( 'Test overflow' ) ;
} ) ;
2025-10-09 10:22:26 -07:00
2025-11-03 23:40:57 +05:30
await waitFor ( ( ) = > {
2026-01-13 14:15:04 -05:00
expect ( mockAddItem ) . toHaveBeenCalledWith ( {
type : 'info' ,
text : expectedMessage ,
} ) ;
2025-11-03 23:40:57 +05:30
} ) ;
} ,
) ;
2025-10-08 15:20:44 -07:00
} ) ;
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 ,
} ,
} ;
} ) ( ) ,
) ;
2026-01-27 16:06:24 -08:00
const { result } = renderHookWithProviders ( ( ) = >
2025-10-08 15:20:44 -07:00
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
2025-10-30 11:50:26 -07:00
await waitFor ( ( ) = > {
2025-11-19 11:11:36 +08:00
expect ( onCancelSubmitSpy ) . toHaveBeenCalledWith ( true ) ;
2025-10-08 15:20:44 -07:00
} ) ;
} ) ;
2026-03-03 01:22:29 -08:00
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 } = 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 ) ,
) ;
} ) ;
} ) ;
2025-11-03 23:40:57 +05:30
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 } ) = > {
2025-07-22 06:57:11 +09:00
mockSendMessageStream . mockReturnValue (
( async function * ( ) {
yield {
type : ServerGeminiEventType . Content ,
value : ` Response for ${ reason } ` ,
} ;
2025-09-02 23:29:07 -06:00
yield {
type : ServerGeminiEventType . Finished ,
value : { reason , usageMetadata : undefined } ,
} ;
2025-07-22 06:57:11 +09:00
} ) ( ) ,
) ;
2025-11-03 23:40:57 +05:30
const { result } = renderHookWithDefaults ( ) ;
2025-07-22 06:57:11 +09:00
await act ( async ( ) = > {
await result . current . submitQuery ( ` Test ${ reason } ` ) ;
} ) ;
2025-11-03 23:40:57 +05:30
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' ,
2025-07-22 06:57:11 +09:00
) ;
2025-11-03 23:40:57 +05:30
expect ( infoMessages ) . toHaveLength ( 0 ) ;
}
} ,
) ;
2025-07-22 06:57:11 +09:00
} ) ;
2025-07-28 23:27:33 +05:30
2026-01-19 17:22:15 -08:00
it ( 'should flush pending text rationale before scheduling tool calls to ensure correct history order' , async ( ) = > {
const addItemOrder : string [ ] = [ ] ;
let capturedOnComplete : any ;
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 : any ) = > ( {
request : r ,
2026-02-13 11:27:20 -05:00
status : CoreToolCallStatus.Success ,
2026-01-19 17:22:15 -08:00
tool : { displayName : r.name , name : r.name } ,
invocation : { getDescription : ( ) = > 'desc' } ,
response : { responseParts : [ ] , resultDisplay : 'done' } ,
startTime : Date.now ( ) ,
endTime : Date.now ( ) ,
} ) ) ;
2026-01-27 16:06:24 -08:00
// Wait a tick for refs to be set up
await new Promise ( ( resolve ) = > setTimeout ( resolve , 0 ) ) ;
2026-01-19 17:22:15 -08:00
await capturedOnComplete ( tools ) ;
addItemOrder . push ( 'scheduleToolCalls_END' ) ;
} ) ;
mockAddItem . mockImplementation ( ( item : any ) = > {
addItemOrder . push ( ` addItem: ${ item . type } ` ) ;
} ) ;
2026-01-21 00:18:42 -05:00
// We need to capture the onComplete callback from useToolScheduler
mockUseToolScheduler . mockImplementation ( ( onComplete ) = > {
2026-01-19 17:22:15 -08:00
capturedOnComplete = onComplete ;
return [
[ ] , // toolCalls
mockScheduleToolCalls ,
vi . fn ( ) , // markToolsAsSubmitted
vi . fn ( ) , // setToolCallsForDisplay
vi . fn ( ) , // cancelAllToolCalls
0 , // lastToolOutputTime
] ;
} ) ;
2026-01-27 16:06:24 -08:00
const { result } = renderHookWithProviders ( ( ) = >
2026-01-19 17:22:15 -08:00
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 ) ;
2026-02-26 18:26:16 -08:00
// Ensure all state updates from recursive submitQuery are settled
await waitFor ( ( ) = > {
expect ( result . current . streamingState ) . toBe ( StreamingState . Idle ) ;
} ) ;
2026-01-19 17:22:15 -08:00
} ) ;
2025-08-20 15:51:31 -04:00
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 ,
} ) ;
2026-01-27 16:06:24 -08:00
const { result } = renderHookWithProviders ( ( ) = >
2025-08-20 15:51:31 -04:00
useGeminiStream (
2025-12-12 17:43:43 -08:00
mockConfig . getGeminiClient ( ) ,
2025-08-20 15:51:31 -04:00
[ ] ,
mockAddItem ,
mockConfig ,
2025-08-28 16:42:54 -07:00
mockLoadedSettings ,
2025-08-20 15:51:31 -04:00
mockOnDebugMessage ,
mockHandleSlashCommand ,
false , // shellModeActive
vi . fn ( ) , // getPreferredEditor
vi . fn ( ) , // onAuthError
vi . fn ( ) , // performMemoryRefresh
false , // modelSwitched
vi . fn ( ) , // setModelSwitched
vi . fn ( ) , // onCancelSubmit
2025-09-11 13:27:27 -07:00
vi . fn ( ) , // setShellInputFocused
80 , // terminalWidth
24 , // terminalHeight
2025-08-20 15:51:31 -04:00
) ,
) ;
await act ( async ( ) = > {
await result . current . submitQuery ( rawQuery ) ;
} ) ;
expect ( handleAtCommandSpy ) . toHaveBeenCalledWith (
expect . objectContaining ( {
query : rawQuery ,
} ) ,
) ;
expect ( mockAddItem ) . toHaveBeenCalledWith (
{
type : MessageType . USER ,
text : rawQuery ,
} ,
userMessageTimestamp ,
) ;
2025-08-22 15:43:53 -04:00
// FIX: The expectation now matches the actual call signature.
2025-08-20 15:51:31 -04:00
expect ( mockSendMessageStream ) . toHaveBeenCalledWith (
2025-08-22 15:43:53 -04:00
processedQueryParts , // Argument 1: The parts array directly
expect . any ( AbortSignal ) , // Argument 2: An AbortSignal
expect . any ( String ) , // Argument 3: The prompt_id string
2026-01-30 10:09:27 -08:00
undefined ,
false ,
rawQuery ,
2025-08-22 15:43:53 -04:00
) ;
} ) ;
2026-01-23 15:50:45 +00:00
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 ,
2026-02-13 17:20:14 -05:00
status : CoreToolCallStatus.Success ,
2026-01-23 15:50:45 +00:00
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 } = 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' ,
2026-02-13 17:20:14 -05:00
status : CoreToolCallStatus.Success ,
2026-01-23 15:50:45 +00:00
} ) ,
] ) ,
} ) ,
expect . any ( Number ) ,
) ;
// 3. Model's response
expect ( mockAddItem ) . toHaveBeenNthCalledWith (
3 ,
expect . objectContaining ( {
type : 'gemini' ,
text : modelResponseContent ,
} ) ,
expect . any ( Number ) ,
) ;
} ) ;
} ) ;
2025-08-22 15:43:53 -04:00
describe ( 'Thought Reset' , ( ) = > {
2026-02-09 19:24:41 -08:00
it ( 'should keep full thinking entries in history when mode is full' , async ( ) = > {
const fullThinkingSettings : LoadedSettings = {
. . . 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 } = 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 } = 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 ) ,
) ;
} ) ;
2025-08-22 15:43:53 -04:00
it ( 'should reset thought to null when starting a new prompt' , async ( ) = > {
2025-08-25 16:21:47 +02:00
// First, simulate a response with a thought
2025-08-22 15:43:53 -04:00
mockSendMessageStream . mockReturnValue (
( async function * ( ) {
yield {
type : ServerGeminiEventType . Thought ,
value : {
subject : 'Previous thought' ,
description : 'Old description' ,
} ,
} ;
yield {
type : ServerGeminiEventType . Content ,
value : 'Some response content' ,
} ;
2025-09-02 23:29:07 -06:00
yield {
type : ServerGeminiEventType . Finished ,
value : { reason : 'STOP' , usageMetadata : undefined } ,
} ;
2025-08-22 15:43:53 -04:00
} ) ( ) ,
) ;
2026-01-27 16:06:24 -08:00
const { result } = renderHookWithProviders ( ( ) = >
2025-08-22 15:43:53 -04:00
useGeminiStream (
new MockedGeminiClientClass ( mockConfig ) ,
[ ] ,
mockAddItem ,
mockConfig ,
2025-08-28 16:42:54 -07:00
mockLoadedSettings ,
2025-08-22 15:43:53 -04:00
mockOnDebugMessage ,
mockHandleSlashCommand ,
false ,
( ) = > 'vscode' as EditorType ,
( ) = > { } ,
( ) = > Promise . resolve ( ) ,
false ,
( ) = > { } ,
( ) = > { } ,
( ) = > { } ,
2025-09-11 13:27:27 -07:00
80 ,
24 ,
2025-08-22 15:43:53 -04:00
) ,
) ;
2025-08-25 16:21:47 +02:00
// Submit first query to set a thought
2025-08-22 15:43:53 -04:00
await act ( async ( ) = > {
await result . current . submitQuery ( 'First query' ) ;
} ) ;
2025-08-25 16:21:47 +02:00
// Wait for the first response to complete
2025-10-30 11:50:26 -07:00
await waitFor ( ( ) = > {
2025-08-22 15:43:53 -04:00
expect ( mockAddItem ) . toHaveBeenCalledWith (
expect . objectContaining ( {
type : 'gemini' ,
text : 'Some response content' ,
} ) ,
expect . any ( Number ) ,
) ;
} ) ;
2025-08-25 16:21:47 +02:00
// Now simulate a new response without a thought
2025-08-22 15:43:53 -04:00
mockSendMessageStream . mockReturnValue (
( async function * ( ) {
yield {
type : ServerGeminiEventType . Content ,
value : 'New response content' ,
} ;
2025-09-02 23:29:07 -06:00
yield {
type : ServerGeminiEventType . Finished ,
value : { reason : 'STOP' , usageMetadata : undefined } ,
} ;
2025-08-22 15:43:53 -04:00
} ) ( ) ,
) ;
2025-08-25 16:21:47 +02:00
// Submit second query - thought should be reset
2025-08-22 15:43:53 -04:00
await act ( async ( ) = > {
await result . current . submitQuery ( 'Second query' ) ;
} ) ;
2025-08-25 16:21:47 +02:00
// 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
2025-10-30 11:50:26 -07:00
await waitFor ( ( ) = > {
2025-08-22 15:43:53 -04:00
expect ( mockAddItem ) . toHaveBeenCalledWith (
expect . objectContaining ( {
type : 'gemini' ,
text : 'New response content' ,
} ) ,
expect . any ( Number ) ,
) ;
} ) ;
} ) ;
2025-09-17 15:37:13 -07:00
it ( 'should memoize pendingHistoryItems' , ( ) = > {
2026-01-21 00:18:42 -05:00
mockUseToolScheduler . mockReturnValue ( [
2025-09-17 15:37:13 -07:00
[ ] ,
mockScheduleToolCalls ,
mockMarkToolsAsSubmitted ,
2026-01-27 16:06:24 -08:00
vi . fn ( ) ,
mockCancelAllToolCalls ,
0 ,
2025-09-17 15:37:13 -07:00
] ) ;
2026-01-27 16:06:24 -08:00
const { result , rerender } = renderHookWithProviders ( ( ) = >
2025-09-17 15:37:13 -07:00
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 : { } } ,
2026-02-13 11:27:20 -05:00
status : CoreToolCallStatus.Executing ,
2025-09-17 15:37:13 -07:00
tool : {
name : 'tool1' ,
displayName : 'tool1' ,
description : 'desc1' ,
build : vi.fn ( ) ,
} ,
invocation : {
getDescription : ( ) = > 'Mock description' ,
} ,
} as unknown as TrackedExecutingToolCall ,
] ;
2026-01-21 00:18:42 -05:00
mockUseToolScheduler . mockReturnValue ( [
2025-09-17 15:37:13 -07:00
newToolCalls ,
mockScheduleToolCalls ,
mockMarkToolsAsSubmitted ,
2026-01-27 16:06:24 -08:00
vi . fn ( ) ,
mockCancelAllToolCalls ,
0 ,
2025-09-17 15:37:13 -07:00
] ) ;
rerender ( ) ;
const thirdResult = result . current . pendingHistoryItems ;
expect ( thirdResult ) . not . toStrictEqual ( secondResult ) ;
} ) ;
2025-08-22 15:43:53 -04:00
it ( 'should reset thought to null when user cancels' , async ( ) = > {
2025-08-25 16:21:47 +02:00
// Mock a stream that yields a thought then gets cancelled
2025-08-22 15:43:53 -04:00
mockSendMessageStream . mockReturnValue (
( async function * ( ) {
yield {
type : ServerGeminiEventType . Thought ,
value : { subject : 'Some thought' , description : 'Description' } ,
} ;
yield { type : ServerGeminiEventType . UserCancelled } ;
} ) ( ) ,
) ;
2026-01-27 16:06:24 -08:00
const { result } = renderHookWithProviders ( ( ) = >
2025-08-22 15:43:53 -04:00
useGeminiStream (
new MockedGeminiClientClass ( mockConfig ) ,
[ ] ,
mockAddItem ,
mockConfig ,
2025-08-28 16:42:54 -07:00
mockLoadedSettings ,
2025-08-22 15:43:53 -04:00
mockOnDebugMessage ,
mockHandleSlashCommand ,
false ,
( ) = > 'vscode' as EditorType ,
( ) = > { } ,
( ) = > Promise . resolve ( ) ,
false ,
( ) = > { } ,
( ) = > { } ,
( ) = > { } ,
2025-09-11 13:27:27 -07:00
80 ,
24 ,
2025-08-22 15:43:53 -04:00
) ,
) ;
2025-08-25 16:21:47 +02:00
// Submit query
2025-08-22 15:43:53 -04:00
await act ( async ( ) = > {
await result . current . submitQuery ( 'Test query' ) ;
} ) ;
2025-08-25 16:21:47 +02:00
// Verify cancellation message was added
2025-10-30 11:50:26 -07:00
await waitFor ( ( ) = > {
2025-08-22 15:43:53 -04:00
expect ( mockAddItem ) . toHaveBeenCalledWith (
expect . objectContaining ( {
type : 'info' ,
text : 'User cancelled the request.' ,
} ) ,
expect . any ( Number ) ,
) ;
} ) ;
2025-08-25 16:21:47 +02:00
// Verify state is reset to idle
2025-08-22 15:43:53 -04:00
expect ( result . current . streamingState ) . toBe ( StreamingState . Idle ) ;
} ) ;
it ( 'should reset thought to null when there is an error' , async ( ) = > {
2025-08-25 16:21:47 +02:00
// Mock a stream that yields a thought then encounters an error
2025-08-22 15:43:53 -04:00
mockSendMessageStream . mockReturnValue (
( async function * ( ) {
yield {
type : ServerGeminiEventType . Thought ,
value : { subject : 'Some thought' , description : 'Description' } ,
} ;
yield {
type : ServerGeminiEventType . Error ,
value : { error : { message : 'Test error' } } ,
} ;
} ) ( ) ,
) ;
2026-01-27 16:06:24 -08:00
const { result } = renderHookWithProviders ( ( ) = >
2025-08-22 15:43:53 -04:00
useGeminiStream (
new MockedGeminiClientClass ( mockConfig ) ,
[ ] ,
mockAddItem ,
mockConfig ,
2025-08-28 16:42:54 -07:00
mockLoadedSettings ,
2025-08-22 15:43:53 -04:00
mockOnDebugMessage ,
mockHandleSlashCommand ,
false ,
( ) = > 'vscode' as EditorType ,
( ) = > { } ,
( ) = > Promise . resolve ( ) ,
false ,
( ) = > { } ,
( ) = > { } ,
( ) = > { } ,
2025-09-11 13:27:27 -07:00
80 ,
24 ,
2025-08-22 15:43:53 -04:00
) ,
) ;
2025-08-25 16:21:47 +02:00
// Submit query
2025-08-22 15:43:53 -04:00
await act ( async ( ) = > {
await result . current . submitQuery ( 'Test query' ) ;
} ) ;
2025-08-25 16:21:47 +02:00
// Verify error message was added
2025-10-30 11:50:26 -07:00
await waitFor ( ( ) = > {
2025-08-22 15:43:53 -04:00
expect ( mockAddItem ) . toHaveBeenCalledWith (
expect . objectContaining ( {
2026-02-13 11:27:20 -05:00
type : CoreToolCallStatus . Error ,
2025-08-22 15:43:53 -04:00
} ) ,
expect . any ( Number ) ,
) ;
} ) ;
2025-08-25 16:21:47 +02:00
// Verify parseAndFormatApiError was called
2025-08-22 15:43:53 -04:00
expect ( mockParseAndFormatApiError ) . toHaveBeenCalledWith (
{ message : 'Test error' } ,
expect . any ( String ) ,
undefined ,
'gemini-2.5-pro' ,
'gemini-2.5-flash' ,
) ;
} ) ;
2026-01-13 06:19:53 -08:00
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' ,
} ;
} ) ( ) ,
) ;
2026-01-27 16:06:24 -08:00
const { result } = renderHookWithProviders ( ( ) = >
2026-01-13 06:19:53 -08:00
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 lastOutputTime was updated
// It should be the time of the last event (startTime + 1000)
expect ( result . current . lastOutputTime ) . toBe ( startTime + 1000 ) ;
vi . useRealTimers ( ) ;
} ) ;
2025-08-22 15:43:53 -04:00
} ) ;
2025-09-10 22:20:13 -07:00
describe ( 'Loop Detection Confirmation' , ( ) = > {
beforeEach ( ( ) = > {
// Add mock for getLoopDetectionService to the config
const mockLoopDetectionService = {
disableForSession : vi.fn ( ) ,
} ;
mockConfig . getGeminiClient = vi . fn ( ) . mockReturnValue ( {
. . . new MockedGeminiClientClass ( mockConfig ) ,
getLoopDetectionService : ( ) = > mockLoopDetectionService ,
} ) ;
} ) ;
it ( 'should set loopDetectionConfirmationRequest when LoopDetected event is received' , async ( ) = > {
mockSendMessageStream . mockReturnValue (
( async function * ( ) {
yield {
type : ServerGeminiEventType . Content ,
value : 'Some content' ,
} ;
yield {
type : ServerGeminiEventType . LoopDetected ,
} ;
} ) ( ) ,
) ;
const { result } = renderTestHook ( ) ;
await act ( async ( ) = > {
await result . current . submitQuery ( 'test query' ) ;
} ) ;
2025-10-30 11:50:26 -07:00
await waitFor ( ( ) = > {
2025-09-10 22:20:13 -07:00
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 ) ;
2025-10-21 13:27:57 -07:00
// Mock for the initial request
mockSendMessageStream . mockReturnValueOnce (
2025-09-10 22:20:13 -07:00
( async function * ( ) {
yield {
type : ServerGeminiEventType . LoopDetected ,
} ;
} ) ( ) ,
) ;
2025-10-21 13:27:57 -07:00
// Mock for the retry request
mockSendMessageStream . mockReturnValueOnce (
( async function * ( ) {
yield {
type : ServerGeminiEventType . Content ,
value : 'Retry successful' ,
} ;
yield {
type : ServerGeminiEventType . Finished ,
value : { reason : 'STOP' } ,
} ;
} ) ( ) ,
) ;
2025-09-10 22:20:13 -07:00
const { result } = renderTestHook ( ) ;
await act ( async ( ) = > {
await result . current . submitQuery ( 'test query' ) ;
} ) ;
// Wait for confirmation request to be set
2025-10-30 11:50:26 -07:00
await waitFor ( ( ) = > {
2025-09-10 22:20:13 -07:00
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
2026-01-13 14:15:04 -05:00
expect ( mockAddItem ) . toHaveBeenCalledWith ( {
type : 'info' ,
text : 'Loop detection has been disabled for this session. Retrying request...' ,
} ) ;
2025-10-21 13:27:57 -07:00
// Verify that the request was retried
2025-10-30 11:50:26 -07:00
await waitFor ( ( ) = > {
2025-10-21 13:27:57 -07:00
expect ( mockSendMessageStream ) . toHaveBeenCalledTimes ( 2 ) ;
expect ( mockSendMessageStream ) . toHaveBeenNthCalledWith (
2 ,
'test query' ,
expect . any ( AbortSignal ) ,
expect . any ( String ) ,
2026-01-30 10:09:27 -08:00
undefined ,
false ,
'test query' ,
2025-10-21 13:27:57 -07:00
) ;
} ) ;
2025-09-10 22:20:13 -07:00
} ) ;
it ( 'should keep loop detection enabled and show message when user selects "keep"' , async ( ) = > {
const mockLoopDetectionService = {
disableForSession : vi.fn ( ) ,
} ;
const mockClient = {
. . . new MockedGeminiClientClass ( mockConfig ) ,
getLoopDetectionService : ( ) = > mockLoopDetectionService ,
} ;
mockConfig . getGeminiClient = vi . fn ( ) . mockReturnValue ( mockClient ) ;
mockSendMessageStream . mockReturnValue (
( async function * ( ) {
yield {
type : ServerGeminiEventType . LoopDetected ,
} ;
} ) ( ) ,
) ;
const { result } = renderTestHook ( ) ;
await act ( async ( ) = > {
await result . current . submitQuery ( 'test query' ) ;
} ) ;
// Wait for confirmation request to be set
2025-10-30 11:50:26 -07:00
await waitFor ( ( ) = > {
2025-09-10 22:20:13 -07:00
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
2026-01-13 14:15:04 -05:00
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.' ,
} ) ;
2025-10-21 13:27:57 -07:00
// Verify that the request was NOT retried
expect ( mockSendMessageStream ) . toHaveBeenCalledTimes ( 1 ) ;
2025-09-10 22:20:13 -07:00
} ) ;
it ( 'should handle multiple loop detection events properly' , async ( ) = > {
const { result } = renderTestHook ( ) ;
// First loop detection - set up fresh mock for first call
mockSendMessageStream . mockReturnValueOnce (
( async function * ( ) {
yield {
type : ServerGeminiEventType . LoopDetected ,
} ;
} ) ( ) ,
) ;
// First loop detection
await act ( async ( ) = > {
await result . current . submitQuery ( 'first query' ) ;
} ) ;
2025-10-30 11:50:26 -07:00
await waitFor ( ( ) = > {
2025-09-10 22:20:13 -07:00
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
2026-01-13 14:15:04 -05:00
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.' ,
} ) ;
2025-09-10 22:20:13 -07:00
// Second loop detection - set up fresh mock for second call
mockSendMessageStream . mockReturnValueOnce (
( async function * ( ) {
yield {
type : ServerGeminiEventType . LoopDetected ,
} ;
} ) ( ) ,
) ;
2025-10-21 13:27:57 -07:00
// Mock for the retry request
mockSendMessageStream . mockReturnValueOnce (
( async function * ( ) {
yield {
type : ServerGeminiEventType . Content ,
value : 'Retry successful' ,
} ;
yield {
type : ServerGeminiEventType . Finished ,
value : { reason : 'STOP' } ,
} ;
} ) ( ) ,
) ;
2025-09-10 22:20:13 -07:00
// Second loop detection
await act ( async ( ) = > {
await result . current . submitQuery ( 'second query' ) ;
} ) ;
2025-10-30 11:50:26 -07:00
await waitFor ( ( ) = > {
2025-09-10 22:20:13 -07:00
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
2026-01-13 14:15:04 -05:00
expect ( mockAddItem ) . toHaveBeenCalledWith ( {
type : 'info' ,
text : 'Loop detection has been disabled for this session. Retrying request...' ,
} ) ;
2025-10-21 13:27:57 -07:00
// Verify that the request was retried
2025-10-30 11:50:26 -07:00
await waitFor ( ( ) = > {
2025-10-21 13:27:57 -07:00
expect ( mockSendMessageStream ) . toHaveBeenCalledTimes ( 3 ) ; // 1st query, 2nd query, retry of 2nd query
expect ( mockSendMessageStream ) . toHaveBeenNthCalledWith (
3 ,
'second query' ,
expect . any ( AbortSignal ) ,
expect . any ( String ) ,
2026-01-30 10:09:27 -08:00
undefined ,
false ,
'second query' ,
2025-10-21 13:27:57 -07:00
) ;
} ) ;
2025-09-10 22:20:13 -07:00
} ) ;
it ( 'should process LoopDetected event after moving pending history to history' , async ( ) = > {
mockSendMessageStream . mockReturnValue (
( async function * ( ) {
yield {
type : ServerGeminiEventType . Content ,
value : 'Some response content' ,
} ;
yield {
type : ServerGeminiEventType . LoopDetected ,
} ;
} ) ( ) ,
) ;
const { result } = renderTestHook ( ) ;
await act ( async ( ) = > {
await result . current . submitQuery ( 'test query' ) ;
} ) ;
// Verify that the content was added to history before the loop detection dialog
2025-10-30 11:50:26 -07:00
await waitFor ( ( ) = > {
2025-09-10 22:20:13 -07:00
expect ( mockAddItem ) . toHaveBeenCalledWith (
expect . objectContaining ( {
type : 'gemini' ,
text : 'Some response content' ,
} ) ,
expect . any ( Number ) ,
) ;
} ) ;
// Then verify loop detection confirmation request was set
2025-10-30 11:50:26 -07:00
await waitFor ( ( ) = > {
2025-09-10 22:20:13 -07:00
expect ( result . current . loopDetectionConfirmationRequest ) . not . toBeNull ( ) ;
} ) ;
} ) ;
2026-03-10 15:41:16 -03:00
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 } = 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 } = 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 ,
false ,
'test query' ,
) ;
} ) ;
} ) ;
} ) ;
2025-09-10 22:20:13 -07:00
} ) ;
2026-01-04 18:58:34 -08:00
describe ( 'Agent Execution Events' , ( ) = > {
2026-01-09 15:47:14 -05:00
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 } = 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 ( ) = > {
2026-01-04 18:58:34 -08:00
mockSendMessageStream . mockReturnValue (
( async function * ( ) {
yield {
type : ServerGeminiEventType . AgentExecutionStopped ,
value : { reason : 'Stopped by hook' } ,
} ;
} ) ( ) ,
) ;
const { result } = 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 ) ;
} ) ;
} ) ;
2026-01-09 15:47:14 -05:00
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 } = 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 ( ) = > {
2026-01-04 18:58:34 -08:00
mockSendMessageStream . mockReturnValue (
( async function * ( ) {
yield {
type : ServerGeminiEventType . AgentExecutionBlocked ,
value : { reason : 'Blocked by hook' } ,
} ;
} ) ( ) ,
) ;
const { result } = 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 ) ,
) ;
} ) ;
} ) ;
} ) ;
2026-02-18 12:53:06 -08:00
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 } = 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 } = 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 ) ;
} ) ;
} ) ;
2026-02-26 18:26:16 -08:00
it ( 'should trace UserPrompt telemetry on submitQuery' , async ( ) = > {
const { result } = 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 ! [ 1 ] ( { metadata : spanMetadata , endSpan : vi.fn ( ) } ) ;
} ) ;
expect ( spanMetadata . input ) . toBe ( 'telemetry test query' ) ;
} ) ;
2025-06-02 01:50:28 -07:00
} ) ;