mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-03-15 16:41:11 -07:00
3854 lines
114 KiB
TypeScript
3854 lines
114 KiB
TypeScript
/**
|
|
* @license
|
|
* Copyright 2025 Google LLC
|
|
* SPDX-License-Identifier: Apache-2.0
|
|
*/
|
|
|
|
/* eslint-disable @typescript-eslint/no-explicit-any */
|
|
import type { Mock, MockInstance } from 'vitest';
|
|
import { describe, it, expect, vi, beforeEach } from 'vitest';
|
|
import { act } from 'react';
|
|
import { renderHookWithProviders } from '../../test-utils/render.js';
|
|
import { waitFor } from '../../test-utils/async.js';
|
|
import { useGeminiStream } from './useGeminiStream.js';
|
|
import { useKeypress } from './useKeypress.js';
|
|
import * as atCommandProcessor from './atCommandProcessor.js';
|
|
import type {
|
|
TrackedToolCall,
|
|
TrackedCompletedToolCall,
|
|
TrackedExecutingToolCall,
|
|
TrackedCancelledToolCall,
|
|
TrackedWaitingToolCall,
|
|
} from './useToolScheduler.js';
|
|
import { useToolScheduler } from './useToolScheduler.js';
|
|
import type {
|
|
Config,
|
|
EditorType,
|
|
AnyToolInvocation,
|
|
SpanMetadata,
|
|
} from '@google/gemini-cli-core';
|
|
import {
|
|
CoreToolCallStatus,
|
|
ApprovalMode,
|
|
AuthType,
|
|
GeminiEventType as ServerGeminiEventType,
|
|
ToolErrorType,
|
|
ToolConfirmationOutcome,
|
|
MessageBusType,
|
|
tokenLimit,
|
|
debugLogger,
|
|
coreEvents,
|
|
CoreEvent,
|
|
MCPDiscoveryState,
|
|
GeminiCliOperation,
|
|
getPlanModeExitMessage,
|
|
} from '@google/gemini-cli-core';
|
|
import type { Part, PartListUnion } from '@google/genai';
|
|
import type { UseHistoryManagerReturn } from './useHistoryManager.js';
|
|
import type { SlashCommandProcessorResult } from '../types.js';
|
|
import { MessageType, StreamingState } from '../types.js';
|
|
|
|
import type { LoadedSettings } from '../../config/settings.js';
|
|
import { findLastSafeSplitPoint } from '../utils/markdownUtilities.js';
|
|
import { theme } from '../semantic-colors.js';
|
|
|
|
// --- MOCKS ---
|
|
const mockSendMessageStream = vi
|
|
.fn()
|
|
.mockReturnValue((async function* () {})());
|
|
const mockStartChat = vi.fn();
|
|
const mockMessageBus = {
|
|
publish: vi.fn(),
|
|
subscribe: vi.fn(),
|
|
unsubscribe: vi.fn(),
|
|
};
|
|
|
|
const MockedGeminiClientClass = vi.hoisted(() =>
|
|
vi.fn().mockImplementation(function (this: any, _config: any) {
|
|
// _config
|
|
this.startChat = mockStartChat;
|
|
this.sendMessageStream = mockSendMessageStream;
|
|
this.addHistory = vi.fn();
|
|
this.generateContent = vi.fn().mockResolvedValue({
|
|
candidates: [
|
|
{ content: { parts: [{ text: 'Got it. Focusing on tests only.' }] } },
|
|
],
|
|
});
|
|
this.getCurrentSequenceModel = vi.fn().mockReturnValue('test-model');
|
|
this.getChat = vi.fn().mockReturnValue({
|
|
recordCompletedToolCalls: vi.fn(),
|
|
});
|
|
this.getChatRecordingService = vi.fn().mockReturnValue({
|
|
recordThought: vi.fn(),
|
|
initialize: vi.fn(),
|
|
recordMessage: vi.fn(),
|
|
recordMessageTokens: vi.fn(),
|
|
recordToolCalls: vi.fn(),
|
|
getConversationFile: vi.fn(),
|
|
});
|
|
this.getCurrentSequenceModel = vi
|
|
.fn()
|
|
.mockReturnValue('gemini-2.0-flash-exp');
|
|
}),
|
|
);
|
|
|
|
const MockedUserPromptEvent = vi.hoisted(() =>
|
|
vi.fn().mockImplementation(() => {}),
|
|
);
|
|
const mockParseAndFormatApiError = vi.hoisted(() => vi.fn());
|
|
|
|
const MockValidationRequiredError = vi.hoisted(
|
|
() =>
|
|
class extends Error {
|
|
userHandled = false;
|
|
},
|
|
);
|
|
|
|
const mockRunInDevTraceSpan = vi.hoisted(() =>
|
|
vi.fn(async (opts, fn) => {
|
|
const metadata: SpanMetadata = {
|
|
name: opts.operation,
|
|
attributes: opts.attributes || {},
|
|
};
|
|
return await fn({
|
|
metadata,
|
|
endSpan: vi.fn(),
|
|
});
|
|
}),
|
|
);
|
|
|
|
vi.mock('@google/gemini-cli-core', async (importOriginal) => {
|
|
const actualCoreModule = (await importOriginal()) as any;
|
|
return {
|
|
...actualCoreModule,
|
|
GitService: vi.fn(),
|
|
GeminiClient: MockedGeminiClientClass,
|
|
UserPromptEvent: MockedUserPromptEvent,
|
|
ValidationRequiredError: MockValidationRequiredError,
|
|
parseAndFormatApiError: mockParseAndFormatApiError,
|
|
tokenLimit: vi.fn().mockReturnValue(100), // Mock tokenLimit
|
|
recordToolCallInteractions: vi.fn().mockResolvedValue(undefined),
|
|
getCodeAssistServer: vi.fn().mockReturnValue(undefined),
|
|
runInDevTraceSpan: mockRunInDevTraceSpan,
|
|
};
|
|
});
|
|
|
|
const mockUseToolScheduler = useToolScheduler as Mock;
|
|
vi.mock('./useToolScheduler.js', async (importOriginal) => {
|
|
const actualSchedulerModule = (await importOriginal()) as any;
|
|
return {
|
|
...(actualSchedulerModule || {}),
|
|
useToolScheduler: vi.fn(),
|
|
};
|
|
});
|
|
|
|
vi.mock('./useKeypress.js', () => ({
|
|
useKeypress: vi.fn(),
|
|
}));
|
|
|
|
vi.mock('./shellCommandProcessor.js', () => ({
|
|
useShellCommandProcessor: vi.fn().mockReturnValue({
|
|
handleShellCommand: vi.fn(),
|
|
activeShellPtyId: null,
|
|
lastShellOutputTime: 0,
|
|
}),
|
|
}));
|
|
|
|
vi.mock('./atCommandProcessor.js');
|
|
|
|
vi.mock('../utils/markdownUtilities.js', () => ({
|
|
findLastSafeSplitPoint: vi.fn((s: string) => s.length),
|
|
}));
|
|
|
|
vi.mock('./useStateAndRef.js', () => ({
|
|
useStateAndRef: vi.fn((initial) => {
|
|
let val = initial;
|
|
const ref = { current: val };
|
|
const setVal = vi.fn((updater) => {
|
|
if (typeof updater === 'function') {
|
|
val = updater(val);
|
|
} else {
|
|
val = updater;
|
|
}
|
|
ref.current = val;
|
|
});
|
|
return [val, ref, setVal];
|
|
}),
|
|
}));
|
|
|
|
vi.mock('./useLogger.js', () => ({
|
|
useLogger: vi.fn().mockReturnValue({
|
|
logMessage: vi.fn().mockResolvedValue(undefined),
|
|
}),
|
|
}));
|
|
|
|
const mockStartNewPrompt = vi.fn();
|
|
const mockAddUsage = vi.fn();
|
|
vi.mock('../contexts/SessionContext.js', async (importOriginal) => {
|
|
const actual = (await importOriginal()) as any;
|
|
return {
|
|
...actual,
|
|
useSessionStats: vi.fn(() => ({
|
|
startNewPrompt: mockStartNewPrompt,
|
|
addUsage: mockAddUsage,
|
|
getPromptCount: vi.fn(() => 5),
|
|
})),
|
|
};
|
|
});
|
|
|
|
vi.mock('./slashCommandProcessor.js', () => ({
|
|
handleSlashCommand: vi.fn().mockReturnValue(false),
|
|
}));
|
|
|
|
vi.mock('./useAlternateBuffer.js', () => ({
|
|
useAlternateBuffer: vi.fn(() => false),
|
|
}));
|
|
|
|
// --- END MOCKS ---
|
|
|
|
// --- Tests for useGeminiStream Hook ---
|
|
describe('useGeminiStream', () => {
|
|
let mockAddItem = vi.fn();
|
|
let mockOnDebugMessage = vi.fn();
|
|
let mockHandleSlashCommand = vi.fn().mockResolvedValue(false);
|
|
let mockScheduleToolCalls: Mock;
|
|
let mockCancelAllToolCalls: Mock;
|
|
let mockMarkToolsAsSubmitted: Mock;
|
|
let handleAtCommandSpy: MockInstance;
|
|
|
|
const emptyHistory: 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;
|
|
});
|
|
|
|
const mockMcpClientManager = {
|
|
getDiscoveryState: vi.fn().mockReturnValue(MCPDiscoveryState.COMPLETED),
|
|
getMcpServerCount: vi.fn().mockReturnValue(0),
|
|
};
|
|
|
|
const mockConfig: Config = {
|
|
apiKey: 'test-api-key',
|
|
model: 'gemini-pro',
|
|
sandbox: false,
|
|
targetDir: '/test/dir',
|
|
debugMode: false,
|
|
question: undefined,
|
|
coreTools: [],
|
|
toolDiscoveryCommand: undefined,
|
|
toolCallCommand: undefined,
|
|
mcpServerCommand: undefined,
|
|
mcpServers: undefined,
|
|
userAgent: 'test-agent',
|
|
userMemory: '',
|
|
geminiMdFileCount: 0,
|
|
alwaysSkipModificationConfirmation: false,
|
|
vertexai: false,
|
|
showMemoryUsage: false,
|
|
contextFileName: undefined,
|
|
storage: {
|
|
getProjectTempDir: vi.fn(() => '/test/temp'),
|
|
getProjectTempCheckpointsDir: vi.fn(() => '/test/temp/checkpoints'),
|
|
} as any,
|
|
getToolRegistry: vi.fn(
|
|
() => ({ getToolSchemaList: vi.fn(() => []) }) as any,
|
|
),
|
|
getProjectRoot: vi.fn(() => '/test/dir'),
|
|
getCheckpointingEnabled: vi.fn(() => false),
|
|
getGeminiClient: mockGetGeminiClient,
|
|
getMcpClientManager: () => mockMcpClientManager as any,
|
|
getApprovalMode: vi.fn(() => ApprovalMode.DEFAULT),
|
|
getUsageStatisticsEnabled: () => true,
|
|
getDebugMode: () => false,
|
|
addHistory: vi.fn(),
|
|
getSessionId: vi.fn(() => 'test-session-id'),
|
|
setQuotaErrorOccurred: vi.fn(),
|
|
resetBillingTurnState: vi.fn(),
|
|
getQuotaErrorOccurred: vi.fn(() => false),
|
|
getModel: vi.fn(() => 'gemini-2.5-pro'),
|
|
getContentGeneratorConfig: vi.fn(() => ({
|
|
model: 'test-model',
|
|
apiKey: 'test-key',
|
|
vertexai: false,
|
|
authType: AuthType.USE_GEMINI,
|
|
})),
|
|
getContentGenerator: vi.fn(),
|
|
isInteractive: () => false,
|
|
getExperiments: () => {},
|
|
getMaxSessionTurns: vi.fn(() => 100),
|
|
isJitContextEnabled: vi.fn(() => false),
|
|
getGlobalMemory: vi.fn(() => ''),
|
|
getUserMemory: vi.fn(() => ''),
|
|
getMessageBus: vi.fn(() => mockMessageBus),
|
|
getBaseLlmClient: vi.fn(() => ({
|
|
generateContent: vi.fn().mockResolvedValue({
|
|
candidates: [
|
|
{ content: { parts: [{ text: 'Got it. Focusing on tests only.' }] } },
|
|
],
|
|
}),
|
|
})),
|
|
getIdeMode: vi.fn(() => false),
|
|
getEnableHooks: vi.fn(() => false),
|
|
} as unknown as Config;
|
|
|
|
beforeEach(() => {
|
|
vi.clearAllMocks(); // Clear mocks before each test
|
|
mockAddItem = vi.fn();
|
|
mockOnDebugMessage = vi.fn();
|
|
mockHandleSlashCommand = vi.fn().mockResolvedValue(false);
|
|
|
|
// Mock return value for useReactToolScheduler
|
|
mockScheduleToolCalls = vi.fn();
|
|
mockCancelAllToolCalls = vi.fn();
|
|
mockMarkToolsAsSubmitted = vi.fn();
|
|
|
|
// Reset properties of mockConfig if needed
|
|
(mockConfig.getCheckpointingEnabled as Mock).mockReturnValue(false);
|
|
(mockConfig.getApprovalMode as Mock).mockReturnValue(ApprovalMode.DEFAULT);
|
|
|
|
// Default mock for useReactToolScheduler to prevent toolCalls being undefined initially
|
|
mockUseToolScheduler.mockReturnValue([
|
|
[], // Default to empty array for toolCalls
|
|
mockScheduleToolCalls,
|
|
mockMarkToolsAsSubmitted,
|
|
vi.fn(), // setToolCallsForDisplay
|
|
mockCancelAllToolCalls,
|
|
0, // lastToolOutputTime
|
|
]);
|
|
|
|
// Reset mocks for GeminiClient instance methods (startChat and sendMessageStream)
|
|
// The GeminiClient constructor itself is mocked at the module level.
|
|
mockStartChat.mockClear().mockResolvedValue({
|
|
sendMessageStream: mockSendMessageStream,
|
|
} as unknown as any); // GeminiChat -> any
|
|
mockSendMessageStream
|
|
.mockClear()
|
|
.mockReturnValue((async function* () {})());
|
|
handleAtCommandSpy = vi.spyOn(atCommandProcessor, 'handleAtCommand');
|
|
vi.spyOn(coreEvents, 'emitFeedback');
|
|
});
|
|
|
|
const mockLoadedSettings: LoadedSettings = {
|
|
merged: {
|
|
preferredEditor: 'vscode',
|
|
ui: { errorVerbosity: 'full' },
|
|
},
|
|
user: { path: '/user/settings.json', settings: {} },
|
|
workspace: { path: '/workspace/.gemini/settings.json', settings: {} },
|
|
errors: [],
|
|
forScope: vi.fn(),
|
|
setValue: vi.fn(),
|
|
} as unknown as LoadedSettings;
|
|
|
|
const renderTestHook = (
|
|
initialToolCalls: TrackedToolCall[] = [],
|
|
geminiClient?: any,
|
|
loadedSettings: LoadedSettings = mockLoadedSettings,
|
|
) => {
|
|
const client = geminiClient || mockConfig.getGeminiClient();
|
|
let lastToolCalls = initialToolCalls;
|
|
|
|
const initialProps = {
|
|
client,
|
|
history: emptyHistory,
|
|
addItem: mockAddItem as unknown as UseHistoryManagerReturn['addItem'],
|
|
config: mockConfig,
|
|
onDebugMessage: mockOnDebugMessage,
|
|
handleSlashCommand: mockHandleSlashCommand as unknown as (
|
|
cmd: PartListUnion,
|
|
) => Promise<SlashCommandProcessorResult | false>,
|
|
shellModeActive: false,
|
|
loadedSettings,
|
|
toolCalls: initialToolCalls,
|
|
};
|
|
|
|
mockUseToolScheduler.mockImplementation((onComplete) => {
|
|
capturedOnComplete = onComplete;
|
|
return [
|
|
lastToolCalls,
|
|
mockScheduleToolCalls,
|
|
mockMarkToolsAsSubmitted,
|
|
(updater: any) => {
|
|
lastToolCalls =
|
|
typeof updater === 'function' ? updater(lastToolCalls) : updater;
|
|
rerender({ ...initialProps, toolCalls: lastToolCalls });
|
|
},
|
|
(...args: any[]) => {
|
|
mockCancelAllToolCalls(...args);
|
|
lastToolCalls = lastToolCalls.map((tc) => {
|
|
if (
|
|
tc.status === CoreToolCallStatus.AwaitingApproval ||
|
|
tc.status === CoreToolCallStatus.Executing ||
|
|
tc.status === CoreToolCallStatus.Scheduled ||
|
|
tc.status === CoreToolCallStatus.Validating
|
|
) {
|
|
return {
|
|
...tc,
|
|
status: CoreToolCallStatus.Cancelled,
|
|
response: {
|
|
callId: tc.request.callId,
|
|
responseParts: [],
|
|
resultDisplay: 'Request cancelled.',
|
|
},
|
|
responseSubmittedToGemini: true,
|
|
} as any as TrackedCancelledToolCall;
|
|
}
|
|
return tc;
|
|
});
|
|
rerender({ ...initialProps, toolCalls: lastToolCalls });
|
|
},
|
|
0,
|
|
];
|
|
});
|
|
|
|
const { result, rerender } = renderHookWithProviders(
|
|
(props: typeof initialProps) =>
|
|
useGeminiStream(
|
|
props.client,
|
|
props.history,
|
|
props.addItem,
|
|
props.config,
|
|
props.loadedSettings,
|
|
props.onDebugMessage,
|
|
props.handleSlashCommand,
|
|
props.shellModeActive,
|
|
mockGetPreferredEditor,
|
|
mockOnAuthError,
|
|
mockPerformMemoryRefresh,
|
|
false,
|
|
mockSetModelSwitchedFromQuotaError,
|
|
mockOnCancelSubmit,
|
|
mockSetShellInputFocused,
|
|
80,
|
|
24,
|
|
),
|
|
{
|
|
initialProps,
|
|
},
|
|
);
|
|
return {
|
|
result,
|
|
rerender,
|
|
mockMarkToolsAsSubmitted,
|
|
mockSendMessageStream,
|
|
client,
|
|
};
|
|
};
|
|
|
|
// Helper to create mock tool calls - reduces boilerplate
|
|
const createMockToolCall = (
|
|
toolName: string,
|
|
callId: string,
|
|
confirmationType: 'edit' | 'info',
|
|
status: TrackedToolCall['status'] = CoreToolCallStatus.AwaitingApproval,
|
|
mockOnConfirm: Mock = vi.fn(),
|
|
): TrackedWaitingToolCall => ({
|
|
request: {
|
|
callId,
|
|
name: toolName,
|
|
args: {},
|
|
isClientInitiated: false,
|
|
prompt_id: 'prompt-id-1',
|
|
},
|
|
status: status as CoreToolCallStatus.AwaitingApproval,
|
|
responseSubmittedToGemini: false,
|
|
confirmationDetails:
|
|
confirmationType === 'edit'
|
|
? {
|
|
type: 'edit',
|
|
title: 'Confirm Edit',
|
|
fileName: 'file.txt',
|
|
filePath: '/test/file.txt',
|
|
fileDiff: 'fake diff',
|
|
originalContent: 'old',
|
|
newContent: 'new',
|
|
onConfirm: mockOnConfirm,
|
|
}
|
|
: {
|
|
type: 'info',
|
|
title: `${toolName} confirmation`,
|
|
prompt: `Execute ${toolName}?`,
|
|
onConfirm: mockOnConfirm,
|
|
},
|
|
tool: {
|
|
name: toolName,
|
|
displayName: toolName,
|
|
description: `${toolName} description`,
|
|
build: vi.fn(),
|
|
} as any,
|
|
invocation: {
|
|
getDescription: () => 'Mock description',
|
|
} as unknown as AnyToolInvocation,
|
|
correlationId: `corr-${callId}`,
|
|
});
|
|
|
|
// Helper to render hook with default parameters - reduces boilerplate
|
|
const renderHookWithDefaults = (
|
|
options: {
|
|
shellModeActive?: boolean;
|
|
onCancelSubmit?: () => void;
|
|
setShellInputFocused?: (focused: boolean) => void;
|
|
performMemoryRefresh?: () => Promise<void>;
|
|
onAuthError?: () => void;
|
|
setModelSwitched?: Mock;
|
|
modelSwitched?: boolean;
|
|
} = {},
|
|
) => {
|
|
const {
|
|
shellModeActive = false,
|
|
onCancelSubmit = () => {},
|
|
setShellInputFocused = () => {},
|
|
performMemoryRefresh = () => Promise.resolve(),
|
|
onAuthError = () => {},
|
|
setModelSwitched = vi.fn(),
|
|
modelSwitched = false,
|
|
} = options;
|
|
|
|
return renderHookWithProviders(() =>
|
|
useGeminiStream(
|
|
new MockedGeminiClientClass(mockConfig),
|
|
[],
|
|
mockAddItem,
|
|
mockConfig,
|
|
mockLoadedSettings,
|
|
mockOnDebugMessage,
|
|
mockHandleSlashCommand,
|
|
shellModeActive,
|
|
() => 'vscode' as EditorType,
|
|
onAuthError,
|
|
performMemoryRefresh,
|
|
modelSwitched,
|
|
setModelSwitched,
|
|
onCancelSubmit,
|
|
setShellInputFocused,
|
|
80,
|
|
24,
|
|
),
|
|
);
|
|
};
|
|
|
|
it('should not submit tool responses if not all tool calls are completed', () => {
|
|
const toolCalls: TrackedToolCall[] = [
|
|
{
|
|
request: {
|
|
callId: 'call1',
|
|
name: 'tool1',
|
|
args: {},
|
|
isClientInitiated: false,
|
|
prompt_id: 'prompt-id-1',
|
|
},
|
|
status: CoreToolCallStatus.Success,
|
|
responseSubmittedToGemini: false,
|
|
response: {
|
|
callId: 'call1',
|
|
responseParts: [{ text: 'tool 1 response' }],
|
|
error: undefined,
|
|
errorType: undefined, // FIX: Added missing property
|
|
resultDisplay: 'Tool 1 success display',
|
|
},
|
|
tool: {
|
|
name: 'tool1',
|
|
displayName: 'tool1',
|
|
description: 'desc1',
|
|
build: vi.fn(),
|
|
} as any,
|
|
invocation: {
|
|
getDescription: () => `Mock description`,
|
|
} as unknown as AnyToolInvocation,
|
|
startTime: Date.now(),
|
|
endTime: Date.now(),
|
|
} as TrackedCompletedToolCall,
|
|
{
|
|
request: {
|
|
callId: 'call2',
|
|
name: 'tool2',
|
|
args: {},
|
|
prompt_id: 'prompt-id-1',
|
|
},
|
|
status: CoreToolCallStatus.Executing,
|
|
responseSubmittedToGemini: false,
|
|
tool: {
|
|
name: 'tool2',
|
|
displayName: 'tool2',
|
|
description: 'desc2',
|
|
build: vi.fn(),
|
|
} as any,
|
|
invocation: {
|
|
getDescription: () => `Mock description`,
|
|
} as unknown as AnyToolInvocation,
|
|
startTime: Date.now(),
|
|
liveOutput: '...',
|
|
} as TrackedExecutingToolCall,
|
|
];
|
|
|
|
const { mockMarkToolsAsSubmitted, mockSendMessageStream } =
|
|
renderTestHook(toolCalls);
|
|
|
|
// Effect for submitting tool responses depends on toolCalls and isResponding
|
|
// isResponding is initially false, so the effect should run.
|
|
|
|
expect(mockMarkToolsAsSubmitted).not.toHaveBeenCalled();
|
|
expect(mockSendMessageStream).not.toHaveBeenCalled(); // submitQuery uses this
|
|
});
|
|
|
|
it('should submit tool responses when all tool calls are completed and ready', async () => {
|
|
const toolCall1ResponseParts: Part[] = [{ text: 'tool 1 final response' }];
|
|
const toolCall2ResponseParts: Part[] = [{ text: 'tool 2 final response' }];
|
|
const completedToolCalls: TrackedToolCall[] = [
|
|
{
|
|
request: {
|
|
callId: 'call1',
|
|
name: 'tool1',
|
|
args: {},
|
|
isClientInitiated: false,
|
|
prompt_id: 'prompt-id-2',
|
|
},
|
|
status: CoreToolCallStatus.Success,
|
|
responseSubmittedToGemini: false,
|
|
response: {
|
|
callId: 'call1',
|
|
responseParts: toolCall1ResponseParts,
|
|
errorType: undefined, // FIX: Added missing property
|
|
},
|
|
tool: {
|
|
displayName: 'MockTool',
|
|
},
|
|
invocation: {
|
|
getDescription: () => `Mock description`,
|
|
} as unknown as AnyToolInvocation,
|
|
} as TrackedCompletedToolCall,
|
|
{
|
|
request: {
|
|
callId: 'call2',
|
|
name: 'tool2',
|
|
args: {},
|
|
isClientInitiated: false,
|
|
prompt_id: 'prompt-id-2',
|
|
},
|
|
status: CoreToolCallStatus.Error,
|
|
responseSubmittedToGemini: false,
|
|
response: {
|
|
callId: 'call2',
|
|
responseParts: toolCall2ResponseParts,
|
|
errorType: ToolErrorType.UNHANDLED_EXCEPTION, // FIX: Added missing property
|
|
},
|
|
} as TrackedCompletedToolCall, // Treat error as a form of completion for submission
|
|
];
|
|
|
|
// Capture the onComplete callback
|
|
let capturedOnComplete:
|
|
| ((completedTools: TrackedToolCall[]) => Promise<void>)
|
|
| null = null;
|
|
|
|
mockUseToolScheduler.mockImplementation((onComplete) => {
|
|
capturedOnComplete = onComplete;
|
|
return [
|
|
[],
|
|
mockScheduleToolCalls,
|
|
mockMarkToolsAsSubmitted,
|
|
vi.fn(),
|
|
mockCancelAllToolCalls,
|
|
0,
|
|
];
|
|
});
|
|
|
|
renderHookWithProviders(() =>
|
|
useGeminiStream(
|
|
new MockedGeminiClientClass(mockConfig),
|
|
[],
|
|
mockAddItem,
|
|
mockConfig,
|
|
mockLoadedSettings,
|
|
mockOnDebugMessage,
|
|
mockHandleSlashCommand,
|
|
false,
|
|
() => 'vscode' as EditorType,
|
|
() => {},
|
|
() => Promise.resolve(),
|
|
false,
|
|
() => {},
|
|
() => {},
|
|
() => {},
|
|
80,
|
|
24,
|
|
),
|
|
);
|
|
|
|
// Trigger the onComplete callback with completed tools
|
|
await act(async () => {
|
|
if (capturedOnComplete) {
|
|
// Wait a tick for refs to be set up
|
|
await new Promise((resolve) => setTimeout(resolve, 0));
|
|
await capturedOnComplete(completedToolCalls);
|
|
}
|
|
});
|
|
|
|
await waitFor(() => {
|
|
expect(mockMarkToolsAsSubmitted).toHaveBeenCalledTimes(1);
|
|
expect(mockSendMessageStream).toHaveBeenCalledTimes(1);
|
|
});
|
|
|
|
const expectedMergedResponse = [
|
|
...toolCall1ResponseParts,
|
|
...toolCall2ResponseParts,
|
|
];
|
|
expect(mockSendMessageStream).toHaveBeenCalledWith(
|
|
expectedMergedResponse,
|
|
expect.any(AbortSignal),
|
|
'prompt-id-2',
|
|
undefined,
|
|
false,
|
|
expectedMergedResponse,
|
|
);
|
|
});
|
|
|
|
it('should inject steering hint prompt for continuation', async () => {
|
|
const toolCallResponseParts: Part[] = [{ text: 'tool final response' }];
|
|
const completedToolCalls: TrackedToolCall[] = [
|
|
{
|
|
request: {
|
|
callId: 'call1',
|
|
name: 'tool1',
|
|
args: {},
|
|
isClientInitiated: false,
|
|
prompt_id: 'prompt-id-ack',
|
|
},
|
|
status: 'success',
|
|
responseSubmittedToGemini: false,
|
|
response: {
|
|
callId: 'call1',
|
|
responseParts: toolCallResponseParts,
|
|
errorType: undefined,
|
|
},
|
|
tool: {
|
|
displayName: 'MockTool',
|
|
},
|
|
invocation: {
|
|
getDescription: () => `Mock description`,
|
|
} as unknown as AnyToolInvocation,
|
|
} as TrackedCompletedToolCall,
|
|
];
|
|
|
|
mockSendMessageStream.mockReturnValue(
|
|
(async function* () {
|
|
yield {
|
|
type: ServerGeminiEventType.Content,
|
|
value: 'Applied the requested adjustment.',
|
|
};
|
|
})(),
|
|
);
|
|
|
|
let capturedOnComplete:
|
|
| ((completedTools: TrackedToolCall[]) => Promise<void>)
|
|
| null = null;
|
|
mockUseToolScheduler.mockImplementation((onComplete) => {
|
|
capturedOnComplete = onComplete;
|
|
return [
|
|
[],
|
|
mockScheduleToolCalls,
|
|
mockMarkToolsAsSubmitted,
|
|
vi.fn(),
|
|
mockCancelAllToolCalls,
|
|
0,
|
|
];
|
|
});
|
|
|
|
renderHookWithProviders(() =>
|
|
useGeminiStream(
|
|
new MockedGeminiClientClass(mockConfig),
|
|
[],
|
|
mockAddItem,
|
|
mockConfig,
|
|
mockLoadedSettings,
|
|
mockOnDebugMessage,
|
|
mockHandleSlashCommand,
|
|
false,
|
|
() => 'vscode' as EditorType,
|
|
() => {},
|
|
() => Promise.resolve(),
|
|
false,
|
|
() => {},
|
|
() => {},
|
|
() => {},
|
|
80,
|
|
24,
|
|
undefined,
|
|
() => 'focus on tests only',
|
|
),
|
|
);
|
|
|
|
await act(async () => {
|
|
if (capturedOnComplete) {
|
|
await new Promise((resolve) => setTimeout(resolve, 0));
|
|
await capturedOnComplete(completedToolCalls);
|
|
}
|
|
});
|
|
|
|
await waitFor(() => {
|
|
expect(mockSendMessageStream).toHaveBeenCalledTimes(1);
|
|
});
|
|
|
|
const sentParts = mockSendMessageStream.mock.calls[0][0] as Part[];
|
|
const injectedHintPart = sentParts[0] as { text?: string };
|
|
expect(injectedHintPart.text).toContain('User steering update:');
|
|
expect(injectedHintPart.text).toContain(
|
|
'<user_input>\nfocus on tests only\n</user_input>',
|
|
);
|
|
expect(injectedHintPart.text).toContain(
|
|
'Classify it as ADD_TASK, MODIFY_TASK, CANCEL_TASK, or EXTRA_CONTEXT.',
|
|
);
|
|
expect(injectedHintPart.text).toContain(
|
|
'Do not cancel/skip tasks unless the user explicitly cancels them.',
|
|
);
|
|
|
|
expect(mockRunInDevTraceSpan).toHaveBeenCalledWith(
|
|
expect.objectContaining({
|
|
operation: GeminiCliOperation.SystemPrompt,
|
|
}),
|
|
expect.any(Function),
|
|
);
|
|
|
|
const spanArgs = mockRunInDevTraceSpan.mock.calls[0];
|
|
const fn = spanArgs[1];
|
|
const metadata = { attributes: {} };
|
|
await act(async () => {
|
|
await fn({ metadata, endSpan: vi.fn() });
|
|
});
|
|
expect(metadata).toMatchObject({
|
|
input: sentParts,
|
|
});
|
|
});
|
|
|
|
it('should handle all tool calls being cancelled', async () => {
|
|
const cancelledToolCalls: TrackedToolCall[] = [
|
|
{
|
|
request: {
|
|
callId: '1',
|
|
name: 'testTool',
|
|
args: {},
|
|
isClientInitiated: false,
|
|
prompt_id: 'prompt-id-3',
|
|
},
|
|
status: CoreToolCallStatus.Cancelled,
|
|
response: {
|
|
callId: '1',
|
|
responseParts: [{ text: CoreToolCallStatus.Cancelled }],
|
|
errorType: undefined, // FIX: Added missing property
|
|
},
|
|
responseSubmittedToGemini: false,
|
|
tool: {
|
|
displayName: 'mock tool',
|
|
},
|
|
invocation: {
|
|
getDescription: () => `Mock description`,
|
|
} as unknown as AnyToolInvocation,
|
|
} as TrackedCancelledToolCall,
|
|
];
|
|
const client = new MockedGeminiClientClass(mockConfig);
|
|
|
|
// Capture the onComplete callback
|
|
let capturedOnComplete:
|
|
| ((completedTools: TrackedToolCall[]) => Promise<void>)
|
|
| null = null;
|
|
|
|
mockUseToolScheduler.mockImplementation((onComplete) => {
|
|
capturedOnComplete = onComplete;
|
|
return [
|
|
[],
|
|
mockScheduleToolCalls,
|
|
mockMarkToolsAsSubmitted,
|
|
vi.fn(),
|
|
mockCancelAllToolCalls,
|
|
0,
|
|
];
|
|
});
|
|
|
|
renderHookWithProviders(() =>
|
|
useGeminiStream(
|
|
client,
|
|
[],
|
|
mockAddItem,
|
|
mockConfig,
|
|
mockLoadedSettings,
|
|
mockOnDebugMessage,
|
|
mockHandleSlashCommand,
|
|
false,
|
|
() => 'vscode' as EditorType,
|
|
() => {},
|
|
() => Promise.resolve(),
|
|
false,
|
|
() => {},
|
|
() => {},
|
|
() => {},
|
|
80,
|
|
24,
|
|
),
|
|
);
|
|
|
|
// Trigger the onComplete callback with cancelled tools
|
|
await act(async () => {
|
|
if (capturedOnComplete) {
|
|
// Wait a tick for refs to be set up
|
|
await new Promise((resolve) => setTimeout(resolve, 0));
|
|
await capturedOnComplete(cancelledToolCalls);
|
|
}
|
|
});
|
|
|
|
await waitFor(() => {
|
|
expect(mockMarkToolsAsSubmitted).toHaveBeenCalledWith(['1']);
|
|
expect(client.addHistory).toHaveBeenCalledWith({
|
|
role: 'user',
|
|
parts: [{ text: CoreToolCallStatus.Cancelled }],
|
|
});
|
|
// Ensure we do NOT call back to the API
|
|
expect(mockSendMessageStream).not.toHaveBeenCalled();
|
|
});
|
|
});
|
|
|
|
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',
|
|
},
|
|
status: CoreToolCallStatus.Error,
|
|
response: {
|
|
callId: 'stop-call',
|
|
responseParts: [{ text: 'error occurred' }],
|
|
errorType: ToolErrorType.STOP_EXECUTION,
|
|
error: new Error('Stop reason from hook'),
|
|
resultDisplay: undefined,
|
|
},
|
|
responseSubmittedToGemini: false,
|
|
tool: {
|
|
displayName: 'stop tool',
|
|
},
|
|
invocation: {
|
|
getDescription: () => `Mock description`,
|
|
} as unknown as AnyToolInvocation,
|
|
} as unknown as TrackedCompletedToolCall,
|
|
];
|
|
const client = new MockedGeminiClientClass(mockConfig);
|
|
|
|
const { result } = renderTestHook([], client);
|
|
|
|
// Trigger the onComplete callback with STOP_EXECUTION tool
|
|
await act(async () => {
|
|
if (capturedOnComplete) {
|
|
await capturedOnComplete(stopExecutionToolCalls);
|
|
}
|
|
});
|
|
|
|
await waitFor(() => {
|
|
expect(mockMarkToolsAsSubmitted).toHaveBeenCalledWith(['stop-call']);
|
|
// Should add an info message to history
|
|
expect(mockAddItem).toHaveBeenCalledWith(
|
|
expect.objectContaining({
|
|
type: MessageType.INFO,
|
|
text: expect.stringContaining(
|
|
'Agent execution stopped: Stop reason from hook',
|
|
),
|
|
}),
|
|
);
|
|
// Ensure we do NOT call back to the API
|
|
expect(mockSendMessageStream).not.toHaveBeenCalled();
|
|
// Streaming state should be Idle
|
|
expect(result.current.streamingState).toBe(StreamingState.Idle);
|
|
});
|
|
|
|
const infoTexts = mockAddItem.mock.calls.map(
|
|
([item]) => (item as { text?: string }).text ?? '',
|
|
);
|
|
expect(
|
|
infoTexts.some((text) =>
|
|
text.includes(
|
|
'Some internal tool attempts failed before this final error',
|
|
),
|
|
),
|
|
).toBe(false);
|
|
expect(
|
|
infoTexts.some((text) =>
|
|
text.includes('This request failed. Press F12 for diagnostics'),
|
|
),
|
|
).toBe(false);
|
|
});
|
|
|
|
it('should add a compact suppressed-error note before STOP_EXECUTION terminal info in low verbosity mode', async () => {
|
|
const stopExecutionToolCalls: 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);
|
|
// The failure hint should NOT be present if the suppressed error note was shown
|
|
expect(failureHintIndex).toBe(-1);
|
|
expect(noteIndex).toBeLessThan(stopIndex);
|
|
});
|
|
|
|
it('should group multiple cancelled tool call responses into a single history entry', async () => {
|
|
const cancelledToolCall1: TrackedCancelledToolCall = {
|
|
request: {
|
|
callId: 'cancel-1',
|
|
name: 'toolA',
|
|
args: {},
|
|
isClientInitiated: false,
|
|
prompt_id: 'prompt-id-7',
|
|
},
|
|
tool: {
|
|
name: 'toolA',
|
|
displayName: 'toolA',
|
|
description: 'descA',
|
|
build: vi.fn(),
|
|
} as any,
|
|
invocation: {
|
|
getDescription: () => `Mock description`,
|
|
} as unknown as AnyToolInvocation,
|
|
status: CoreToolCallStatus.Cancelled,
|
|
response: {
|
|
callId: 'cancel-1',
|
|
responseParts: [
|
|
{ functionResponse: { name: 'toolA', id: 'cancel-1' } },
|
|
],
|
|
resultDisplay: undefined,
|
|
error: undefined,
|
|
errorType: undefined, // FIX: Added missing property
|
|
},
|
|
responseSubmittedToGemini: false,
|
|
};
|
|
const cancelledToolCall2: TrackedCancelledToolCall = {
|
|
request: {
|
|
callId: 'cancel-2',
|
|
name: 'toolB',
|
|
args: {},
|
|
isClientInitiated: false,
|
|
prompt_id: 'prompt-id-8',
|
|
},
|
|
tool: {
|
|
name: 'toolB',
|
|
displayName: 'toolB',
|
|
description: 'descB',
|
|
build: vi.fn(),
|
|
} as any,
|
|
invocation: {
|
|
getDescription: () => `Mock description`,
|
|
} as unknown as AnyToolInvocation,
|
|
status: CoreToolCallStatus.Cancelled,
|
|
response: {
|
|
callId: 'cancel-2',
|
|
responseParts: [
|
|
{ functionResponse: { name: 'toolB', id: 'cancel-2' } },
|
|
],
|
|
resultDisplay: undefined,
|
|
error: undefined,
|
|
errorType: undefined, // FIX: Added missing property
|
|
},
|
|
responseSubmittedToGemini: false,
|
|
};
|
|
const allCancelledTools = [cancelledToolCall1, cancelledToolCall2];
|
|
const client = new MockedGeminiClientClass(mockConfig);
|
|
|
|
let capturedOnComplete:
|
|
| ((completedTools: TrackedToolCall[]) => Promise<void>)
|
|
| null = null;
|
|
|
|
mockUseToolScheduler.mockImplementation((onComplete) => {
|
|
capturedOnComplete = onComplete;
|
|
return [
|
|
[],
|
|
mockScheduleToolCalls,
|
|
mockMarkToolsAsSubmitted,
|
|
vi.fn(),
|
|
mockCancelAllToolCalls,
|
|
0,
|
|
];
|
|
});
|
|
|
|
renderHookWithProviders(() =>
|
|
useGeminiStream(
|
|
client,
|
|
[],
|
|
mockAddItem,
|
|
mockConfig,
|
|
mockLoadedSettings,
|
|
mockOnDebugMessage,
|
|
mockHandleSlashCommand,
|
|
false,
|
|
() => 'vscode' as EditorType,
|
|
() => {},
|
|
() => Promise.resolve(),
|
|
false,
|
|
() => {},
|
|
() => {},
|
|
() => {},
|
|
80,
|
|
24,
|
|
),
|
|
);
|
|
|
|
// Trigger the onComplete callback with multiple cancelled tools
|
|
await act(async () => {
|
|
if (capturedOnComplete) {
|
|
// Wait a tick for refs to be set up
|
|
await new Promise((resolve) => setTimeout(resolve, 0));
|
|
await capturedOnComplete(allCancelledTools);
|
|
}
|
|
});
|
|
|
|
await waitFor(() => {
|
|
// The tools should be marked as submitted locally
|
|
expect(mockMarkToolsAsSubmitted).toHaveBeenCalledWith([
|
|
'cancel-1',
|
|
'cancel-2',
|
|
]);
|
|
|
|
// Crucially, addHistory should be called only ONCE
|
|
expect(client.addHistory).toHaveBeenCalledTimes(1);
|
|
|
|
// And that single call should contain BOTH function responses
|
|
expect(client.addHistory).toHaveBeenCalledWith({
|
|
role: 'user',
|
|
parts: [
|
|
...cancelledToolCall1.response.responseParts,
|
|
...cancelledToolCall2.response.responseParts,
|
|
],
|
|
});
|
|
|
|
// No message should be sent back to the API for a turn with only cancellations
|
|
expect(mockSendMessageStream).not.toHaveBeenCalled();
|
|
});
|
|
});
|
|
|
|
it('should not flicker streaming state to Idle between tool completion and submission', async () => {
|
|
const toolCallResponseParts: PartListUnion = [
|
|
{ text: 'tool 1 final response' },
|
|
];
|
|
|
|
const initialToolCalls: TrackedToolCall[] = [
|
|
{
|
|
request: {
|
|
callId: 'call1',
|
|
name: 'tool1',
|
|
args: {},
|
|
isClientInitiated: false,
|
|
prompt_id: 'prompt-id-4',
|
|
},
|
|
status: CoreToolCallStatus.Executing,
|
|
responseSubmittedToGemini: false,
|
|
tool: {
|
|
name: 'tool1',
|
|
displayName: 'tool1',
|
|
description: 'desc',
|
|
build: vi.fn(),
|
|
} as any,
|
|
invocation: {
|
|
getDescription: () => `Mock description`,
|
|
} as unknown as AnyToolInvocation,
|
|
startTime: Date.now(),
|
|
} as TrackedExecutingToolCall,
|
|
];
|
|
|
|
const completedToolCalls: TrackedToolCall[] = [
|
|
{
|
|
...(initialToolCalls[0] as TrackedExecutingToolCall),
|
|
status: CoreToolCallStatus.Success,
|
|
response: {
|
|
callId: 'call1',
|
|
responseParts: toolCallResponseParts,
|
|
error: undefined,
|
|
errorType: undefined, // FIX: Added missing property
|
|
resultDisplay: 'Tool 1 success display',
|
|
},
|
|
endTime: Date.now(),
|
|
} as TrackedCompletedToolCall,
|
|
];
|
|
|
|
// Capture the onComplete callback
|
|
let capturedOnComplete:
|
|
| ((completedTools: TrackedToolCall[]) => Promise<void>)
|
|
| null = null;
|
|
let currentToolCalls = initialToolCalls;
|
|
|
|
mockUseToolScheduler.mockImplementation((onComplete) => {
|
|
capturedOnComplete = onComplete;
|
|
return [
|
|
currentToolCalls,
|
|
mockScheduleToolCalls,
|
|
mockMarkToolsAsSubmitted,
|
|
vi.fn(), // setToolCallsForDisplay
|
|
mockCancelAllToolCalls,
|
|
0,
|
|
];
|
|
});
|
|
|
|
const { result, rerender } = renderHookWithProviders(() =>
|
|
useGeminiStream(
|
|
new MockedGeminiClientClass(mockConfig),
|
|
[],
|
|
mockAddItem,
|
|
mockConfig,
|
|
mockLoadedSettings,
|
|
mockOnDebugMessage,
|
|
mockHandleSlashCommand,
|
|
false,
|
|
() => 'vscode' as EditorType,
|
|
() => {},
|
|
() => Promise.resolve(),
|
|
false,
|
|
() => {},
|
|
() => {},
|
|
() => {},
|
|
80,
|
|
24,
|
|
),
|
|
);
|
|
|
|
// 1. Initial state should be Responding because a tool is executing.
|
|
expect(result.current.streamingState).toBe(StreamingState.Responding);
|
|
|
|
// 2. Update the tool calls to completed state and rerender
|
|
currentToolCalls = completedToolCalls;
|
|
mockUseToolScheduler.mockImplementation((onComplete) => {
|
|
capturedOnComplete = onComplete;
|
|
return [
|
|
completedToolCalls,
|
|
mockScheduleToolCalls,
|
|
mockMarkToolsAsSubmitted,
|
|
vi.fn(), // setToolCallsForDisplay
|
|
mockCancelAllToolCalls,
|
|
0,
|
|
];
|
|
});
|
|
|
|
act(() => {
|
|
rerender();
|
|
});
|
|
|
|
// 3. The state should *still* be Responding, not Idle.
|
|
// This is because the completed tool's response has not been submitted yet.
|
|
expect(result.current.streamingState).toBe(StreamingState.Responding);
|
|
|
|
// 4. Trigger the onComplete callback to simulate tool completion
|
|
await act(async () => {
|
|
if (capturedOnComplete) {
|
|
// Wait a tick for refs to be set up
|
|
await new Promise((resolve) => setTimeout(resolve, 0));
|
|
await capturedOnComplete(completedToolCalls);
|
|
}
|
|
});
|
|
|
|
// 5. Wait for submitQuery to be called
|
|
await waitFor(() => {
|
|
expect(mockSendMessageStream).toHaveBeenCalledWith(
|
|
toolCallResponseParts,
|
|
expect.any(AbortSignal),
|
|
'prompt-id-4',
|
|
undefined,
|
|
false,
|
|
toolCallResponseParts,
|
|
);
|
|
});
|
|
|
|
// 6. After submission, the state should remain Responding until the stream completes.
|
|
expect(result.current.streamingState).toBe(StreamingState.Responding);
|
|
});
|
|
|
|
describe('User Cancellation', () => {
|
|
let keypressCallback: (key: any) => void;
|
|
const mockUseKeypress = useKeypress as Mock;
|
|
|
|
beforeEach(() => {
|
|
// Capture the callback passed to useKeypress
|
|
mockUseKeypress.mockImplementation((callback, options) => {
|
|
if (options.isActive) {
|
|
keypressCallback = callback;
|
|
} else {
|
|
keypressCallback = () => {};
|
|
}
|
|
});
|
|
});
|
|
|
|
const simulateEscapeKeyPress = () => {
|
|
act(() => {
|
|
keypressCallback({ name: 'escape' });
|
|
});
|
|
};
|
|
|
|
it('should cancel an in-progress stream when escape is pressed', async () => {
|
|
const mockStream = (async function* () {
|
|
yield { type: 'content', value: 'Part 1' };
|
|
// Keep the stream open
|
|
await new Promise(() => {});
|
|
})();
|
|
mockSendMessageStream.mockReturnValue(mockStream);
|
|
|
|
const { result } = renderTestHook();
|
|
|
|
// Start a query
|
|
await act(async () => {
|
|
// eslint-disable-next-line @typescript-eslint/no-floating-promises
|
|
result.current.submitQuery('test query');
|
|
});
|
|
|
|
// Wait for the first part of the response
|
|
await waitFor(() => {
|
|
expect(result.current.streamingState).toBe(StreamingState.Responding);
|
|
});
|
|
|
|
// Simulate escape key press
|
|
simulateEscapeKeyPress();
|
|
|
|
// Verify cancellation message is added
|
|
await waitFor(() => {
|
|
expect(mockAddItem).toHaveBeenCalledWith({
|
|
type: MessageType.INFO,
|
|
text: 'Request cancelled.',
|
|
});
|
|
});
|
|
|
|
// Verify state is reset
|
|
expect(result.current.streamingState).toBe(StreamingState.Idle);
|
|
});
|
|
|
|
it('should call onCancelSubmit handler when escape is pressed', async () => {
|
|
const cancelSubmitSpy = vi.fn();
|
|
const mockStream = (async function* () {
|
|
yield { type: 'content', value: 'Part 1' };
|
|
// Keep the stream open
|
|
await new Promise(() => {});
|
|
})();
|
|
mockSendMessageStream.mockReturnValue(mockStream);
|
|
|
|
const { result } = renderHookWithProviders(() =>
|
|
useGeminiStream(
|
|
mockConfig.getGeminiClient(),
|
|
[],
|
|
mockAddItem,
|
|
mockConfig,
|
|
mockLoadedSettings,
|
|
mockOnDebugMessage,
|
|
mockHandleSlashCommand,
|
|
false,
|
|
() => 'vscode' as EditorType,
|
|
() => {},
|
|
() => Promise.resolve(),
|
|
false,
|
|
() => {},
|
|
cancelSubmitSpy,
|
|
() => {},
|
|
80,
|
|
24,
|
|
),
|
|
);
|
|
|
|
// Start a query
|
|
await act(async () => {
|
|
// eslint-disable-next-line @typescript-eslint/no-floating-promises
|
|
result.current.submitQuery('test query');
|
|
});
|
|
|
|
simulateEscapeKeyPress();
|
|
|
|
expect(cancelSubmitSpy).toHaveBeenCalledWith(false);
|
|
});
|
|
|
|
it('should call setShellInputFocused(false) when escape is pressed', async () => {
|
|
const setShellInputFocusedSpy = vi.fn();
|
|
const mockStream = (async function* () {
|
|
yield { type: 'content', value: 'Part 1' };
|
|
await new Promise(() => {}); // Keep stream open
|
|
})();
|
|
mockSendMessageStream.mockReturnValue(mockStream);
|
|
|
|
const { result } = renderHookWithProviders(() =>
|
|
useGeminiStream(
|
|
mockConfig.getGeminiClient(),
|
|
[],
|
|
mockAddItem,
|
|
mockConfig,
|
|
mockLoadedSettings,
|
|
mockOnDebugMessage,
|
|
mockHandleSlashCommand,
|
|
false,
|
|
() => 'vscode' as EditorType,
|
|
() => {},
|
|
() => Promise.resolve(),
|
|
false,
|
|
() => {},
|
|
vi.fn(),
|
|
setShellInputFocusedSpy, // Pass the spy here
|
|
80,
|
|
24,
|
|
),
|
|
);
|
|
|
|
// Start a query
|
|
await act(async () => {
|
|
// eslint-disable-next-line @typescript-eslint/no-floating-promises
|
|
result.current.submitQuery('test query');
|
|
});
|
|
|
|
simulateEscapeKeyPress();
|
|
|
|
expect(setShellInputFocusedSpy).toHaveBeenCalledWith(false);
|
|
});
|
|
|
|
it('should not do anything if escape is pressed when not responding', () => {
|
|
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 () => {
|
|
// eslint-disable-next-line @typescript-eslint/no-floating-promises
|
|
result.current.submitQuery('long running query');
|
|
});
|
|
|
|
await waitFor(() => {
|
|
expect(result.current.streamingState).toBe(StreamingState.Responding);
|
|
});
|
|
|
|
// Cancel the request
|
|
simulateEscapeKeyPress();
|
|
|
|
// Allow the stream to continue
|
|
await act(async () => {
|
|
continueStream();
|
|
// Wait a bit to see if the second part is processed
|
|
await new Promise((resolve) => setTimeout(resolve, 50));
|
|
});
|
|
|
|
// The text should not have been updated with " Canceled"
|
|
const lastCall = mockAddItem.mock.calls.find(
|
|
(call) => call[0].type === 'gemini',
|
|
);
|
|
expect(lastCall?.[0].text).toBe('Initial');
|
|
|
|
// The final state should be idle after cancellation
|
|
expect(result.current.streamingState).toBe(StreamingState.Idle);
|
|
});
|
|
|
|
it('should cancel if a tool call is in progress', async () => {
|
|
const toolCalls: TrackedToolCall[] = [
|
|
{
|
|
request: { callId: 'call1', name: 'tool1', args: {} },
|
|
status: CoreToolCallStatus.Executing,
|
|
responseSubmittedToGemini: false,
|
|
tool: {
|
|
name: 'tool1',
|
|
description: 'desc1',
|
|
build: vi.fn().mockImplementation((_) => ({
|
|
getDescription: () => `Mock description`,
|
|
})),
|
|
} as any,
|
|
invocation: {
|
|
getDescription: () => `Mock description`,
|
|
},
|
|
startTime: Date.now(),
|
|
liveOutput: '...',
|
|
} as TrackedExecutingToolCall,
|
|
];
|
|
|
|
const { result } = renderTestHook(toolCalls);
|
|
|
|
// State is `Responding` because a tool is running
|
|
expect(result.current.streamingState).toBe(StreamingState.Responding);
|
|
|
|
// Try to cancel
|
|
simulateEscapeKeyPress();
|
|
|
|
// The cancel function should be called
|
|
expect(mockCancelAllToolCalls).toHaveBeenCalled();
|
|
});
|
|
|
|
it('should cancel a request when a tool is awaiting confirmation', async () => {
|
|
const mockOnConfirm = vi.fn().mockResolvedValue(undefined);
|
|
const toolCalls: TrackedToolCall[] = [
|
|
{
|
|
request: {
|
|
callId: 'confirm-call',
|
|
name: 'some_tool',
|
|
args: {},
|
|
isClientInitiated: false,
|
|
prompt_id: 'prompt-id-1',
|
|
},
|
|
status: CoreToolCallStatus.AwaitingApproval,
|
|
responseSubmittedToGemini: false,
|
|
tool: {
|
|
name: 'some_tool',
|
|
description: 'a tool',
|
|
build: vi.fn().mockImplementation((_) => ({
|
|
getDescription: () => `Mock description`,
|
|
})),
|
|
} as any,
|
|
invocation: {
|
|
getDescription: () => `Mock description`,
|
|
} as unknown as AnyToolInvocation,
|
|
confirmationDetails: {
|
|
type: 'edit',
|
|
title: 'Confirm Edit',
|
|
onConfirm: mockOnConfirm,
|
|
fileName: 'file.txt',
|
|
filePath: '/test/file.txt',
|
|
fileDiff: 'fake diff',
|
|
originalContent: 'old',
|
|
newContent: 'new',
|
|
},
|
|
} as TrackedWaitingToolCall,
|
|
];
|
|
|
|
const { result } = renderTestHook(toolCalls);
|
|
|
|
// State is `WaitingForConfirmation` because a tool is awaiting approval
|
|
expect(result.current.streamingState).toBe(
|
|
StreamingState.WaitingForConfirmation,
|
|
);
|
|
|
|
// Try to cancel
|
|
simulateEscapeKeyPress();
|
|
|
|
// The imperative cancel function should be called on the scheduler
|
|
expect(mockCancelAllToolCalls).toHaveBeenCalled();
|
|
|
|
// A cancellation message should be added to history
|
|
await waitFor(() => {
|
|
expect(mockAddItem).toHaveBeenCalledWith(
|
|
expect.objectContaining({
|
|
text: 'Request cancelled.',
|
|
}),
|
|
);
|
|
});
|
|
|
|
// The final state should be idle
|
|
expect(result.current.streamingState).toBe(StreamingState.Idle);
|
|
});
|
|
});
|
|
|
|
describe('Retry Handling', () => {
|
|
it('should 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();
|
|
});
|
|
});
|
|
|
|
describe('Slash Command Handling', () => {
|
|
it('should schedule a tool call when the command processor returns a schedule_tool action', async () => {
|
|
const clientToolRequest: SlashCommandProcessorResult = {
|
|
type: 'schedule_tool',
|
|
toolName: 'save_memory',
|
|
toolArgs: { fact: 'test fact' },
|
|
};
|
|
mockHandleSlashCommand.mockResolvedValue(clientToolRequest);
|
|
|
|
const { result } = renderTestHook();
|
|
|
|
await act(async () => {
|
|
await result.current.submitQuery('/memory add "test fact"');
|
|
});
|
|
|
|
await waitFor(() => {
|
|
expect(mockScheduleToolCalls).toHaveBeenCalledWith(
|
|
[
|
|
expect.objectContaining({
|
|
name: 'save_memory',
|
|
args: { fact: 'test fact' },
|
|
isClientInitiated: true,
|
|
}),
|
|
],
|
|
expect.any(AbortSignal),
|
|
);
|
|
expect(mockSendMessageStream).not.toHaveBeenCalled();
|
|
});
|
|
});
|
|
|
|
it('should stop processing and not call Gemini when a command is handled without a tool call', async () => {
|
|
const uiOnlyCommandResult: SlashCommandProcessorResult = {
|
|
type: 'handled',
|
|
};
|
|
mockHandleSlashCommand.mockResolvedValue(uiOnlyCommandResult);
|
|
|
|
const { result } = renderTestHook();
|
|
|
|
await act(async () => {
|
|
await result.current.submitQuery('/help');
|
|
});
|
|
|
|
await waitFor(() => {
|
|
expect(mockHandleSlashCommand).toHaveBeenCalledWith('/help');
|
|
expect(mockScheduleToolCalls).not.toHaveBeenCalled();
|
|
expect(mockSendMessageStream).not.toHaveBeenCalled(); // No LLM call made
|
|
});
|
|
});
|
|
|
|
it('should call Gemini with prompt content when slash command returns a `submit_prompt` action', async () => {
|
|
const customCommandResult: SlashCommandProcessorResult = {
|
|
type: 'submit_prompt',
|
|
content: 'This is the actual prompt from the command file.',
|
|
};
|
|
mockHandleSlashCommand.mockResolvedValue(customCommandResult);
|
|
|
|
const { result, mockSendMessageStream: localMockSendMessageStream } =
|
|
renderTestHook();
|
|
|
|
await act(async () => {
|
|
await result.current.submitQuery('/my-custom-command');
|
|
});
|
|
|
|
await waitFor(() => {
|
|
expect(mockHandleSlashCommand).toHaveBeenCalledWith(
|
|
'/my-custom-command',
|
|
);
|
|
|
|
expect(localMockSendMessageStream).not.toHaveBeenCalledWith(
|
|
'/my-custom-command',
|
|
expect.anything(),
|
|
expect.anything(),
|
|
);
|
|
|
|
expect(localMockSendMessageStream).toHaveBeenCalledWith(
|
|
'This is the actual prompt from the command file.',
|
|
expect.any(AbortSignal),
|
|
expect.any(String),
|
|
undefined,
|
|
false,
|
|
'/my-custom-command',
|
|
);
|
|
|
|
expect(mockScheduleToolCalls).not.toHaveBeenCalled();
|
|
});
|
|
});
|
|
|
|
it('should correctly handle a submit_prompt action with empty content', async () => {
|
|
const emptyPromptResult: SlashCommandProcessorResult = {
|
|
type: 'submit_prompt',
|
|
content: '',
|
|
};
|
|
mockHandleSlashCommand.mockResolvedValue(emptyPromptResult);
|
|
|
|
const { result, mockSendMessageStream: localMockSendMessageStream } =
|
|
renderTestHook();
|
|
|
|
await act(async () => {
|
|
await result.current.submitQuery('/emptycmd');
|
|
});
|
|
|
|
await waitFor(() => {
|
|
expect(mockHandleSlashCommand).toHaveBeenCalledWith('/emptycmd');
|
|
expect(localMockSendMessageStream).toHaveBeenCalledWith(
|
|
'',
|
|
expect.any(AbortSignal),
|
|
expect.any(String),
|
|
undefined,
|
|
false,
|
|
'/emptycmd',
|
|
);
|
|
});
|
|
});
|
|
|
|
it('should not call handleSlashCommand for line comments', async () => {
|
|
const { result, mockSendMessageStream: localMockSendMessageStream } =
|
|
renderTestHook();
|
|
|
|
await act(async () => {
|
|
await result.current.submitQuery('// This is a line comment');
|
|
});
|
|
|
|
await waitFor(() => {
|
|
expect(mockHandleSlashCommand).not.toHaveBeenCalled();
|
|
expect(localMockSendMessageStream).toHaveBeenCalledWith(
|
|
'// This is a line comment',
|
|
expect.any(AbortSignal),
|
|
expect.any(String),
|
|
undefined,
|
|
false,
|
|
'// This is a line comment',
|
|
);
|
|
});
|
|
});
|
|
|
|
it('should not call handleSlashCommand for block comments', async () => {
|
|
const { result, mockSendMessageStream: localMockSendMessageStream } =
|
|
renderTestHook();
|
|
|
|
await act(async () => {
|
|
await result.current.submitQuery('/* This is a block comment */');
|
|
});
|
|
|
|
await waitFor(() => {
|
|
expect(mockHandleSlashCommand).not.toHaveBeenCalled();
|
|
expect(localMockSendMessageStream).toHaveBeenCalledWith(
|
|
'/* This is a block comment */',
|
|
expect.any(AbortSignal),
|
|
expect.any(String),
|
|
undefined,
|
|
false,
|
|
'/* This is a block comment */',
|
|
);
|
|
});
|
|
});
|
|
|
|
it('should not call handleSlashCommand is shell mode is active', async () => {
|
|
const { result } = renderHookWithProviders(() =>
|
|
useGeminiStream(
|
|
new MockedGeminiClientClass(mockConfig),
|
|
[],
|
|
mockAddItem,
|
|
mockConfig,
|
|
mockLoadedSettings,
|
|
() => {},
|
|
mockHandleSlashCommand,
|
|
true,
|
|
() => 'vscode' as EditorType,
|
|
() => {},
|
|
() => Promise.resolve(),
|
|
false,
|
|
() => {},
|
|
() => {},
|
|
() => {},
|
|
80,
|
|
24,
|
|
),
|
|
);
|
|
|
|
await act(async () => {
|
|
await result.current.submitQuery('/about');
|
|
});
|
|
|
|
await waitFor(() => {
|
|
expect(mockHandleSlashCommand).not.toHaveBeenCalled();
|
|
});
|
|
});
|
|
});
|
|
|
|
describe('Memory Refresh on save_memory', () => {
|
|
it('should call performMemoryRefresh when a save_memory tool call completes successfully', async () => {
|
|
const mockPerformMemoryRefresh = vi.fn();
|
|
const completedToolCall: TrackedCompletedToolCall = {
|
|
request: {
|
|
callId: 'save-mem-call-1',
|
|
name: 'save_memory',
|
|
args: { fact: 'test' },
|
|
isClientInitiated: true,
|
|
prompt_id: 'prompt-id-6',
|
|
},
|
|
status: CoreToolCallStatus.Success,
|
|
responseSubmittedToGemini: false,
|
|
response: {
|
|
callId: 'save-mem-call-1',
|
|
responseParts: [{ text: 'Memory saved' }],
|
|
resultDisplay: 'Success: Memory saved',
|
|
error: undefined,
|
|
errorType: undefined, // FIX: Added missing property
|
|
},
|
|
tool: {
|
|
name: 'save_memory',
|
|
displayName: 'save_memory',
|
|
description: 'Saves memory',
|
|
build: vi.fn(),
|
|
} as any,
|
|
invocation: {
|
|
getDescription: () => `Mock description`,
|
|
} as unknown as AnyToolInvocation,
|
|
};
|
|
|
|
// Capture the onComplete callback
|
|
let capturedOnComplete:
|
|
| ((completedTools: TrackedToolCall[]) => Promise<void>)
|
|
| null = null;
|
|
|
|
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,
|
|
() => {},
|
|
mockPerformMemoryRefresh,
|
|
false,
|
|
() => {},
|
|
() => {},
|
|
() => {},
|
|
80,
|
|
24,
|
|
),
|
|
);
|
|
|
|
// Trigger the onComplete callback with the completed save_memory tool
|
|
await act(async () => {
|
|
if (capturedOnComplete) {
|
|
// Wait a tick for refs to be set up
|
|
await new Promise((resolve) => setTimeout(resolve, 0));
|
|
await capturedOnComplete([completedToolCall]);
|
|
}
|
|
});
|
|
|
|
await waitFor(() => {
|
|
expect(mockPerformMemoryRefresh).toHaveBeenCalledTimes(1);
|
|
});
|
|
});
|
|
});
|
|
|
|
describe('Error Handling', () => {
|
|
it('should call parseAndFormatApiError with the correct authType on stream initialization failure', async () => {
|
|
// 1. Setup
|
|
const mockError = new Error('Rate limit exceeded');
|
|
const mockAuthType = AuthType.LOGIN_WITH_GOOGLE;
|
|
mockParseAndFormatApiError.mockClear();
|
|
mockSendMessageStream.mockReturnValue(
|
|
(async function* () {
|
|
yield { type: 'content', value: '' };
|
|
throw mockError;
|
|
})(),
|
|
);
|
|
|
|
const testConfig = {
|
|
...mockConfig,
|
|
getContentGenerator: vi.fn(),
|
|
getContentGeneratorConfig: vi.fn(() => ({
|
|
authType: mockAuthType,
|
|
})),
|
|
getModel: vi.fn(() => 'gemini-2.5-pro'),
|
|
} as unknown as Config;
|
|
|
|
const { result } = renderHookWithProviders(() =>
|
|
useGeminiStream(
|
|
new MockedGeminiClientClass(testConfig),
|
|
[],
|
|
mockAddItem,
|
|
testConfig,
|
|
mockLoadedSettings,
|
|
mockOnDebugMessage,
|
|
mockHandleSlashCommand,
|
|
false,
|
|
() => 'vscode' as EditorType,
|
|
() => {},
|
|
() => Promise.resolve(),
|
|
false,
|
|
() => {},
|
|
() => {},
|
|
() => {},
|
|
80,
|
|
24,
|
|
),
|
|
);
|
|
|
|
// 2. Action
|
|
await act(async () => {
|
|
await result.current.submitQuery('test query');
|
|
});
|
|
|
|
// 3. Assertion
|
|
await waitFor(() => {
|
|
expect(mockParseAndFormatApiError).toHaveBeenCalledWith(
|
|
'Rate limit exceeded',
|
|
mockAuthType,
|
|
undefined,
|
|
'gemini-2.5-pro',
|
|
'gemini-2.5-flash',
|
|
);
|
|
});
|
|
});
|
|
});
|
|
|
|
describe('handleApprovalModeChange', () => {
|
|
it('should auto-approve all pending tool calls when switching to YOLO mode', async () => {
|
|
const awaitingApprovalToolCalls: TrackedToolCall[] = [
|
|
createMockToolCall('replace', 'call1', 'edit'),
|
|
createMockToolCall('read_file', 'call2', 'info'),
|
|
];
|
|
|
|
const { result } = renderTestHook(awaitingApprovalToolCalls);
|
|
|
|
await act(async () => {
|
|
await result.current.handleApprovalModeChange(ApprovalMode.YOLO);
|
|
});
|
|
|
|
// Both tool calls should be auto-approved
|
|
expect(mockMessageBus.publish).toHaveBeenCalledTimes(2);
|
|
expect(mockMessageBus.publish).toHaveBeenCalledWith(
|
|
expect.objectContaining({
|
|
type: MessageBusType.TOOL_CONFIRMATION_RESPONSE,
|
|
correlationId: 'corr-call1',
|
|
outcome: ToolConfirmationOutcome.ProceedOnce,
|
|
}),
|
|
);
|
|
expect(mockMessageBus.publish).toHaveBeenCalledWith(
|
|
expect.objectContaining({
|
|
correlationId: 'corr-call2',
|
|
outcome: ToolConfirmationOutcome.ProceedOnce,
|
|
}),
|
|
);
|
|
});
|
|
|
|
it('should only auto-approve edit tools when switching to AUTO_EDIT mode', async () => {
|
|
const awaitingApprovalToolCalls: TrackedToolCall[] = [
|
|
createMockToolCall('replace', 'call1', 'edit'),
|
|
createMockToolCall('write_file', 'call2', 'edit'),
|
|
createMockToolCall('read_file', 'call3', 'info'),
|
|
];
|
|
|
|
const { result } = renderTestHook(awaitingApprovalToolCalls);
|
|
|
|
await act(async () => {
|
|
await result.current.handleApprovalModeChange(ApprovalMode.AUTO_EDIT);
|
|
});
|
|
|
|
// Only replace and write_file should be auto-approved
|
|
expect(mockMessageBus.publish).toHaveBeenCalledTimes(2);
|
|
expect(mockMessageBus.publish).toHaveBeenCalledWith(
|
|
expect.objectContaining({ correlationId: 'corr-call1' }),
|
|
);
|
|
expect(mockMessageBus.publish).toHaveBeenCalledWith(
|
|
expect.objectContaining({ correlationId: 'corr-call2' }),
|
|
);
|
|
expect(mockMessageBus.publish).not.toHaveBeenCalledWith(
|
|
expect.objectContaining({ correlationId: 'corr-call3' }),
|
|
);
|
|
});
|
|
|
|
it('should not auto-approve any tools when switching to REQUIRE_CONFIRMATION mode', async () => {
|
|
const awaitingApprovalToolCalls: TrackedToolCall[] = [
|
|
createMockToolCall('replace', 'call1', 'edit'),
|
|
];
|
|
|
|
const { result } = renderTestHook(awaitingApprovalToolCalls);
|
|
|
|
await act(async () => {
|
|
await result.current.handleApprovalModeChange(ApprovalMode.DEFAULT);
|
|
});
|
|
|
|
// No tools should be auto-approved
|
|
expect(mockMessageBus.publish).not.toHaveBeenCalled();
|
|
});
|
|
|
|
it('should handle errors gracefully when auto-approving tool calls', async () => {
|
|
const debuggerSpy = vi
|
|
.spyOn(debugLogger, 'warn')
|
|
.mockImplementation(() => {});
|
|
|
|
mockMessageBus.publish.mockRejectedValueOnce(new Error('Bus error'));
|
|
|
|
const awaitingApprovalToolCalls: TrackedToolCall[] = [
|
|
createMockToolCall('replace', 'call1', 'edit'),
|
|
createMockToolCall('write_file', 'call2', 'edit'),
|
|
];
|
|
|
|
const { result } = renderTestHook(awaitingApprovalToolCalls);
|
|
|
|
await act(async () => {
|
|
await result.current.handleApprovalModeChange(ApprovalMode.YOLO);
|
|
});
|
|
|
|
// Both should be attempted despite first error
|
|
expect(mockMessageBus.publish).toHaveBeenCalledTimes(2);
|
|
expect(debuggerSpy).toHaveBeenCalledWith(
|
|
'Failed to auto-approve tool call call1:',
|
|
expect.any(Error),
|
|
);
|
|
|
|
debuggerSpy.mockRestore();
|
|
});
|
|
|
|
it('should skip tool calls without confirmationDetails', async () => {
|
|
const awaitingApprovalToolCalls: TrackedToolCall[] = [
|
|
{
|
|
request: {
|
|
callId: 'call1',
|
|
name: 'replace',
|
|
args: { old_string: 'old', new_string: 'new' },
|
|
isClientInitiated: false,
|
|
prompt_id: 'prompt-id-1',
|
|
},
|
|
status: CoreToolCallStatus.AwaitingApproval,
|
|
responseSubmittedToGemini: false,
|
|
// No confirmationDetails
|
|
tool: {
|
|
name: 'replace',
|
|
displayName: 'replace',
|
|
description: 'Replace text',
|
|
build: vi.fn(),
|
|
} as any,
|
|
invocation: {
|
|
getDescription: () => 'Mock description',
|
|
} as unknown as AnyToolInvocation,
|
|
correlationId: 'corr-1',
|
|
} as unknown as TrackedWaitingToolCall,
|
|
];
|
|
|
|
const { result } = renderTestHook(awaitingApprovalToolCalls);
|
|
|
|
// Should not throw an error
|
|
await act(async () => {
|
|
await result.current.handleApprovalModeChange(ApprovalMode.YOLO);
|
|
});
|
|
});
|
|
|
|
it('should only process tool calls with awaiting_approval status', async () => {
|
|
const mockOnConfirmAwaiting = vi.fn().mockResolvedValue(undefined);
|
|
const mixedStatusToolCalls: TrackedToolCall[] = [
|
|
createMockToolCall(
|
|
'replace',
|
|
'call1',
|
|
'edit',
|
|
CoreToolCallStatus.AwaitingApproval,
|
|
mockOnConfirmAwaiting,
|
|
),
|
|
{
|
|
request: {
|
|
callId: 'call2',
|
|
name: 'write_file',
|
|
args: { path: '/test/file.txt', content: 'content' },
|
|
isClientInitiated: false,
|
|
prompt_id: 'prompt-id-1',
|
|
},
|
|
status: CoreToolCallStatus.Executing,
|
|
responseSubmittedToGemini: false,
|
|
tool: {
|
|
name: 'write_file',
|
|
displayName: 'write_file',
|
|
description: 'Write file',
|
|
build: vi.fn(),
|
|
} as any,
|
|
invocation: {
|
|
getDescription: () => 'Mock description',
|
|
} as unknown as AnyToolInvocation,
|
|
startTime: Date.now(),
|
|
liveOutput: 'Writing...',
|
|
correlationId: 'corr-call2',
|
|
} as TrackedExecutingToolCall,
|
|
];
|
|
|
|
const { result } = renderTestHook(mixedStatusToolCalls);
|
|
|
|
await act(async () => {
|
|
await result.current.handleApprovalModeChange(ApprovalMode.YOLO);
|
|
});
|
|
|
|
// Only the awaiting_approval tool should be processed.
|
|
expect(mockMessageBus.publish).toHaveBeenCalledTimes(1);
|
|
expect(mockMessageBus.publish).toHaveBeenCalledWith(
|
|
expect.objectContaining({ correlationId: 'corr-call1' }),
|
|
);
|
|
expect(mockMessageBus.publish).not.toHaveBeenCalledWith(
|
|
expect.objectContaining({ correlationId: 'corr-call2' }),
|
|
);
|
|
});
|
|
|
|
it('should inject a notification message when manually exiting Plan Mode', async () => {
|
|
// Setup mockConfig to return PLAN mode initially
|
|
(mockConfig.getApprovalMode as Mock).mockReturnValue(ApprovalMode.PLAN);
|
|
|
|
// Render the hook, which will initialize the previousApprovalModeRef with PLAN
|
|
const { result, client } = renderTestHook([]);
|
|
|
|
// Update mockConfig to return DEFAULT mode (new mode)
|
|
(mockConfig.getApprovalMode as Mock).mockReturnValue(
|
|
ApprovalMode.DEFAULT,
|
|
);
|
|
|
|
await act(async () => {
|
|
// Trigger manual exit from Plan Mode
|
|
await result.current.handleApprovalModeChange(ApprovalMode.DEFAULT);
|
|
});
|
|
|
|
// Verify that addHistory was called with the notification message
|
|
expect(client.addHistory).toHaveBeenCalledWith({
|
|
role: 'user',
|
|
parts: [
|
|
{
|
|
text: getPlanModeExitMessage(ApprovalMode.DEFAULT, true),
|
|
},
|
|
],
|
|
});
|
|
});
|
|
});
|
|
|
|
describe('handleFinishedEvent', () => {
|
|
it('should add info message for MAX_TOKENS finish reason', async () => {
|
|
// Setup mock to return a stream with MAX_TOKENS finish reason
|
|
mockSendMessageStream.mockReturnValue(
|
|
(async function* () {
|
|
yield {
|
|
type: ServerGeminiEventType.Content,
|
|
value: 'This is a truncated response...',
|
|
};
|
|
yield {
|
|
type: ServerGeminiEventType.Finished,
|
|
value: { reason: 'MAX_TOKENS', usageMetadata: undefined },
|
|
};
|
|
})(),
|
|
);
|
|
|
|
const { result } = renderHookWithProviders(() =>
|
|
useGeminiStream(
|
|
new MockedGeminiClientClass(mockConfig),
|
|
[],
|
|
mockAddItem,
|
|
mockConfig,
|
|
mockLoadedSettings,
|
|
mockOnDebugMessage,
|
|
mockHandleSlashCommand,
|
|
false,
|
|
() => 'vscode' as EditorType,
|
|
() => {},
|
|
() => Promise.resolve(),
|
|
false,
|
|
() => {},
|
|
() => {},
|
|
() => {},
|
|
80,
|
|
24,
|
|
),
|
|
);
|
|
|
|
// Submit a query
|
|
await act(async () => {
|
|
await result.current.submitQuery('Generate long text');
|
|
});
|
|
|
|
// Check that the info message was added
|
|
await waitFor(() => {
|
|
expect(mockAddItem).toHaveBeenCalledWith(
|
|
{
|
|
type: 'info',
|
|
text: '⚠️ Response truncated due to token limits.',
|
|
},
|
|
expect.any(Number),
|
|
);
|
|
});
|
|
});
|
|
|
|
describe('ContextWindowWillOverflow event', () => {
|
|
beforeEach(() => {
|
|
vi.mocked(tokenLimit).mockReturnValue(100);
|
|
});
|
|
|
|
it.each([
|
|
{
|
|
name: 'without suggestion when remaining tokens are > 75% of limit',
|
|
requestTokens: 20,
|
|
remainingTokens: 80,
|
|
expectedMessage:
|
|
'Sending this message (20 tokens) might exceed the context window limit (80 tokens left).',
|
|
},
|
|
{
|
|
name: 'with suggestion when remaining tokens are < 75% of limit',
|
|
requestTokens: 30,
|
|
remainingTokens: 70,
|
|
expectedMessage:
|
|
'Sending this message (30 tokens) might exceed the context window limit (70 tokens left). Please try reducing the size of your message or use the `/compress` command to compress the chat history.',
|
|
},
|
|
])(
|
|
'should add message $name',
|
|
async ({ requestTokens, remainingTokens, expectedMessage }) => {
|
|
mockSendMessageStream.mockReturnValue(
|
|
(async function* () {
|
|
yield {
|
|
type: ServerGeminiEventType.ContextWindowWillOverflow,
|
|
value: {
|
|
estimatedRequestTokenCount: requestTokens,
|
|
remainingTokenCount: remainingTokens,
|
|
},
|
|
};
|
|
})(),
|
|
);
|
|
|
|
const { result } = renderHookWithDefaults();
|
|
|
|
await act(async () => {
|
|
await result.current.submitQuery('Test overflow');
|
|
});
|
|
|
|
await waitFor(() => {
|
|
expect(mockAddItem).toHaveBeenCalledWith({
|
|
type: 'info',
|
|
text: expectedMessage,
|
|
});
|
|
});
|
|
},
|
|
);
|
|
});
|
|
|
|
it('should call onCancelSubmit when ContextWindowWillOverflow event is received', async () => {
|
|
const onCancelSubmitSpy = vi.fn();
|
|
// Setup mock to return a stream with ContextWindowWillOverflow event
|
|
mockSendMessageStream.mockReturnValue(
|
|
(async function* () {
|
|
yield {
|
|
type: ServerGeminiEventType.ContextWindowWillOverflow,
|
|
value: {
|
|
estimatedRequestTokenCount: 100,
|
|
remainingTokenCount: 50,
|
|
},
|
|
};
|
|
})(),
|
|
);
|
|
|
|
const { result } = renderHookWithProviders(() =>
|
|
useGeminiStream(
|
|
new MockedGeminiClientClass(mockConfig),
|
|
[],
|
|
mockAddItem,
|
|
mockConfig,
|
|
mockLoadedSettings,
|
|
mockOnDebugMessage,
|
|
mockHandleSlashCommand,
|
|
false,
|
|
() => 'vscode' as EditorType,
|
|
() => {},
|
|
() => Promise.resolve(),
|
|
false,
|
|
() => {},
|
|
onCancelSubmitSpy,
|
|
() => {},
|
|
80,
|
|
24,
|
|
),
|
|
);
|
|
|
|
// Submit a query
|
|
await act(async () => {
|
|
await result.current.submitQuery('Test overflow');
|
|
});
|
|
|
|
// Check that onCancelSubmit was called
|
|
await waitFor(() => {
|
|
expect(onCancelSubmitSpy).toHaveBeenCalledWith(true);
|
|
});
|
|
});
|
|
|
|
it('should add informational messages when ChatCompressed event is received', async () => {
|
|
vi.mocked(tokenLimit).mockReturnValue(10000);
|
|
// Setup mock to return a stream with ChatCompressed event
|
|
mockSendMessageStream.mockReturnValue(
|
|
(async function* () {
|
|
yield {
|
|
type: ServerGeminiEventType.ChatCompressed,
|
|
value: {
|
|
originalTokenCount: 1000,
|
|
newTokenCount: 500,
|
|
compressionStatus: 'compressed',
|
|
},
|
|
};
|
|
})(),
|
|
);
|
|
|
|
const { result } = renderHookWithDefaults();
|
|
|
|
// Submit a query
|
|
await act(async () => {
|
|
await result.current.submitQuery('Test compression');
|
|
});
|
|
|
|
// Check that the succinct info message was added
|
|
await waitFor(() => {
|
|
expect(mockAddItem).toHaveBeenCalledWith(
|
|
expect.objectContaining({
|
|
type: MessageType.INFO,
|
|
text: 'Context compressed from 10% to 5%.',
|
|
secondaryText: 'Change threshold in /settings.',
|
|
color: theme.status.warning,
|
|
}),
|
|
expect.any(Number),
|
|
);
|
|
});
|
|
});
|
|
|
|
it.each([
|
|
{
|
|
reason: 'STOP',
|
|
shouldAddMessage: false,
|
|
},
|
|
{
|
|
reason: 'FINISH_REASON_UNSPECIFIED',
|
|
shouldAddMessage: false,
|
|
},
|
|
{
|
|
reason: 'SAFETY',
|
|
message: '⚠️ Response stopped due to safety reasons.',
|
|
},
|
|
{
|
|
reason: 'RECITATION',
|
|
message: '⚠️ Response stopped due to recitation policy.',
|
|
},
|
|
{
|
|
reason: 'LANGUAGE',
|
|
message: '⚠️ Response stopped due to unsupported language.',
|
|
},
|
|
{
|
|
reason: 'BLOCKLIST',
|
|
message: '⚠️ Response stopped due to forbidden terms.',
|
|
},
|
|
{
|
|
reason: 'PROHIBITED_CONTENT',
|
|
message: '⚠️ Response stopped due to prohibited content.',
|
|
},
|
|
{
|
|
reason: 'SPII',
|
|
message:
|
|
'⚠️ Response stopped due to sensitive personally identifiable information.',
|
|
},
|
|
{
|
|
reason: 'OTHER',
|
|
message: '⚠️ Response stopped for other reasons.',
|
|
},
|
|
{
|
|
reason: 'MALFORMED_FUNCTION_CALL',
|
|
message: '⚠️ Response stopped due to malformed function call.',
|
|
},
|
|
{
|
|
reason: 'IMAGE_SAFETY',
|
|
message: '⚠️ Response stopped due to image safety violations.',
|
|
},
|
|
{
|
|
reason: 'UNEXPECTED_TOOL_CALL',
|
|
message: '⚠️ Response stopped due to unexpected tool call.',
|
|
},
|
|
])(
|
|
'should handle $reason finish reason correctly',
|
|
async ({ reason, shouldAddMessage = true, message }) => {
|
|
mockSendMessageStream.mockReturnValue(
|
|
(async function* () {
|
|
yield {
|
|
type: ServerGeminiEventType.Content,
|
|
value: `Response for ${reason}`,
|
|
};
|
|
yield {
|
|
type: ServerGeminiEventType.Finished,
|
|
value: { reason, usageMetadata: undefined },
|
|
};
|
|
})(),
|
|
);
|
|
|
|
const { result } = renderHookWithDefaults();
|
|
|
|
await act(async () => {
|
|
await result.current.submitQuery(`Test ${reason}`);
|
|
});
|
|
|
|
if (shouldAddMessage) {
|
|
await waitFor(() => {
|
|
expect(mockAddItem).toHaveBeenCalledWith(
|
|
{
|
|
type: 'info',
|
|
text: message,
|
|
},
|
|
expect.any(Number),
|
|
);
|
|
});
|
|
} else {
|
|
// Verify state returns to idle without any info messages
|
|
await waitFor(() => {
|
|
expect(result.current.streamingState).toBe(StreamingState.Idle);
|
|
});
|
|
|
|
const infoMessages = mockAddItem.mock.calls.filter(
|
|
(call) => call[0].type === 'info',
|
|
);
|
|
expect(infoMessages).toHaveLength(0);
|
|
}
|
|
},
|
|
);
|
|
});
|
|
|
|
it('should flush pending text rationale before scheduling tool calls to ensure correct history order', async () => {
|
|
const addItemOrder: string[] = [];
|
|
let capturedOnComplete: 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,
|
|
status: CoreToolCallStatus.Success,
|
|
tool: { displayName: r.name, name: r.name },
|
|
invocation: { getDescription: () => 'desc' },
|
|
response: { responseParts: [], resultDisplay: 'done' },
|
|
startTime: Date.now(),
|
|
endTime: Date.now(),
|
|
}));
|
|
// Wait a tick for refs to be set up
|
|
await new Promise((resolve) => setTimeout(resolve, 0));
|
|
await capturedOnComplete(tools);
|
|
addItemOrder.push('scheduleToolCalls_END');
|
|
});
|
|
|
|
mockAddItem.mockImplementation((item: any) => {
|
|
addItemOrder.push(`addItem:${item.type}`);
|
|
});
|
|
|
|
// We need to capture the onComplete callback from useToolScheduler
|
|
mockUseToolScheduler.mockImplementation((onComplete) => {
|
|
capturedOnComplete = onComplete;
|
|
return [
|
|
[], // toolCalls
|
|
mockScheduleToolCalls,
|
|
vi.fn(), // markToolsAsSubmitted
|
|
vi.fn(), // setToolCallsForDisplay
|
|
vi.fn(), // cancelAllToolCalls
|
|
0, // lastToolOutputTime
|
|
];
|
|
});
|
|
|
|
const { result } = renderHookWithProviders(() =>
|
|
useGeminiStream(
|
|
new MockedGeminiClientClass(mockConfig),
|
|
[],
|
|
mockAddItem,
|
|
mockConfig,
|
|
mockLoadedSettings,
|
|
vi.fn(),
|
|
vi.fn(),
|
|
false,
|
|
() => 'vscode' as EditorType,
|
|
vi.fn(),
|
|
vi.fn(),
|
|
false,
|
|
vi.fn(),
|
|
vi.fn(),
|
|
vi.fn(),
|
|
80,
|
|
24,
|
|
),
|
|
);
|
|
|
|
const mockStream = (async function* () {
|
|
yield {
|
|
type: ServerGeminiEventType.Content,
|
|
value: 'Rationale rationale.',
|
|
};
|
|
yield {
|
|
type: ServerGeminiEventType.ToolCallRequest,
|
|
value: { callId: '1', name: 'test_tool', args: {} },
|
|
};
|
|
})();
|
|
mockSendMessageStream.mockReturnValue(mockStream);
|
|
|
|
await act(async () => {
|
|
await result.current.submitQuery('test input');
|
|
});
|
|
|
|
// Expectation: addItem:gemini (rationale) MUST happen before scheduleToolCalls_START
|
|
const rationaleIndex = addItemOrder.indexOf('addItem:gemini');
|
|
const scheduleIndex = addItemOrder.indexOf('scheduleToolCalls_START');
|
|
const toolGroupIndex = addItemOrder.indexOf('addItem:tool_group');
|
|
|
|
expect(rationaleIndex).toBeGreaterThan(-1);
|
|
expect(scheduleIndex).toBeGreaterThan(-1);
|
|
expect(toolGroupIndex).toBeGreaterThan(-1);
|
|
|
|
// This is the core fix validation: Rationale comes before tools are even scheduled (awaited)
|
|
expect(rationaleIndex).toBeLessThan(scheduleIndex);
|
|
expect(rationaleIndex).toBeLessThan(toolGroupIndex);
|
|
|
|
// Ensure all state updates from recursive submitQuery are settled
|
|
await waitFor(() => {
|
|
expect(result.current.streamingState).toBe(StreamingState.Idle);
|
|
});
|
|
});
|
|
|
|
it('should process @include commands, adding user turn after processing to prevent race conditions', async () => {
|
|
const rawQuery = '@include file.txt Summarize this.';
|
|
const processedQueryParts = [
|
|
{ text: 'Summarize this with content from @file.txt' },
|
|
{ text: 'File content...' },
|
|
];
|
|
const userMessageTimestamp = Date.now();
|
|
vi.spyOn(Date, 'now').mockReturnValue(userMessageTimestamp);
|
|
|
|
handleAtCommandSpy.mockResolvedValue({
|
|
processedQuery: processedQueryParts,
|
|
shouldProceed: true,
|
|
});
|
|
|
|
const { result } = renderHookWithProviders(() =>
|
|
useGeminiStream(
|
|
mockConfig.getGeminiClient(),
|
|
[],
|
|
mockAddItem,
|
|
mockConfig,
|
|
mockLoadedSettings,
|
|
mockOnDebugMessage,
|
|
mockHandleSlashCommand,
|
|
false, // shellModeActive
|
|
vi.fn(), // getPreferredEditor
|
|
vi.fn(), // onAuthError
|
|
vi.fn(), // performMemoryRefresh
|
|
false, // modelSwitched
|
|
vi.fn(), // setModelSwitched
|
|
vi.fn(), // onCancelSubmit
|
|
vi.fn(), // setShellInputFocused
|
|
80, // terminalWidth
|
|
24, // terminalHeight
|
|
),
|
|
);
|
|
|
|
await act(async () => {
|
|
await result.current.submitQuery(rawQuery);
|
|
});
|
|
|
|
expect(handleAtCommandSpy).toHaveBeenCalledWith(
|
|
expect.objectContaining({
|
|
query: rawQuery,
|
|
}),
|
|
);
|
|
|
|
expect(mockAddItem).toHaveBeenCalledWith(
|
|
{
|
|
type: MessageType.USER,
|
|
text: rawQuery,
|
|
},
|
|
userMessageTimestamp,
|
|
);
|
|
|
|
// FIX: The expectation now matches the actual call signature.
|
|
expect(mockSendMessageStream).toHaveBeenCalledWith(
|
|
processedQueryParts, // Argument 1: The parts array directly
|
|
expect.any(AbortSignal), // Argument 2: An AbortSignal
|
|
expect.any(String), // Argument 3: The prompt_id string
|
|
undefined,
|
|
false,
|
|
rawQuery,
|
|
);
|
|
});
|
|
|
|
it('should display user query, then tool execution, then model response', async () => {
|
|
const userQuery = 'read this @file(test.txt)';
|
|
const toolExecutionMessage = 'Reading file: test.txt';
|
|
const modelResponseContent = 'The content of test.txt is: Hello World!';
|
|
|
|
// Mock handleAtCommand to simulate a tool call and add a tool_group message
|
|
handleAtCommandSpy.mockImplementation(
|
|
async ({ addItem: atCommandAddItem, messageId }) => {
|
|
atCommandAddItem(
|
|
{
|
|
type: 'tool_group',
|
|
tools: [
|
|
{
|
|
callId: 'client-read-123',
|
|
name: 'read_file',
|
|
description: toolExecutionMessage,
|
|
status: CoreToolCallStatus.Success,
|
|
resultDisplay: toolExecutionMessage,
|
|
confirmationDetails: undefined,
|
|
},
|
|
],
|
|
},
|
|
messageId,
|
|
);
|
|
return { shouldProceed: true, processedQuery: userQuery };
|
|
},
|
|
);
|
|
|
|
// Mock the Gemini stream to return a model response after the tool
|
|
mockSendMessageStream.mockReturnValue(
|
|
(async function* () {
|
|
yield {
|
|
type: ServerGeminiEventType.Content,
|
|
value: modelResponseContent,
|
|
};
|
|
yield {
|
|
type: ServerGeminiEventType.Finished,
|
|
value: { reason: 'STOP' },
|
|
};
|
|
})(),
|
|
);
|
|
|
|
const { result } = renderTestHook();
|
|
|
|
await act(async () => {
|
|
await result.current.submitQuery(userQuery);
|
|
});
|
|
|
|
// Assert the order of messages added to the history
|
|
await waitFor(() => {
|
|
expect(mockAddItem).toHaveBeenCalledTimes(3); // User prompt + tool execution + model response
|
|
|
|
// 1. User's prompt
|
|
expect(mockAddItem).toHaveBeenNthCalledWith(
|
|
1,
|
|
expect.objectContaining({
|
|
type: MessageType.USER,
|
|
text: userQuery,
|
|
}),
|
|
expect.any(Number),
|
|
);
|
|
|
|
// 2. Tool execution message
|
|
expect(mockAddItem).toHaveBeenNthCalledWith(
|
|
2,
|
|
expect.objectContaining({
|
|
type: 'tool_group',
|
|
tools: expect.arrayContaining([
|
|
expect.objectContaining({
|
|
name: 'read_file',
|
|
status: CoreToolCallStatus.Success,
|
|
}),
|
|
]),
|
|
}),
|
|
expect.any(Number),
|
|
);
|
|
|
|
// 3. Model's response
|
|
expect(mockAddItem).toHaveBeenNthCalledWith(
|
|
3,
|
|
expect.objectContaining({
|
|
type: 'gemini',
|
|
text: modelResponseContent,
|
|
}),
|
|
expect.any(Number),
|
|
);
|
|
});
|
|
});
|
|
describe('Thought Reset', () => {
|
|
it('should keep full thinking entries in history when mode is full', async () => {
|
|
const fullThinkingSettings: LoadedSettings = {
|
|
...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),
|
|
);
|
|
});
|
|
|
|
it('should reset thought to null when starting a new prompt', async () => {
|
|
// First, simulate a response with a thought
|
|
mockSendMessageStream.mockReturnValue(
|
|
(async function* () {
|
|
yield {
|
|
type: ServerGeminiEventType.Thought,
|
|
value: {
|
|
subject: 'Previous thought',
|
|
description: 'Old description',
|
|
},
|
|
};
|
|
yield {
|
|
type: ServerGeminiEventType.Content,
|
|
value: 'Some response content',
|
|
};
|
|
yield {
|
|
type: ServerGeminiEventType.Finished,
|
|
value: { reason: 'STOP', usageMetadata: undefined },
|
|
};
|
|
})(),
|
|
);
|
|
|
|
const { result } = renderHookWithProviders(() =>
|
|
useGeminiStream(
|
|
new MockedGeminiClientClass(mockConfig),
|
|
[],
|
|
mockAddItem,
|
|
mockConfig,
|
|
mockLoadedSettings,
|
|
mockOnDebugMessage,
|
|
mockHandleSlashCommand,
|
|
false,
|
|
() => 'vscode' as EditorType,
|
|
() => {},
|
|
() => Promise.resolve(),
|
|
false,
|
|
() => {},
|
|
() => {},
|
|
() => {},
|
|
80,
|
|
24,
|
|
),
|
|
);
|
|
|
|
// Submit first query to set a thought
|
|
await act(async () => {
|
|
await result.current.submitQuery('First query');
|
|
});
|
|
|
|
// Wait for the first response to complete
|
|
await waitFor(() => {
|
|
expect(mockAddItem).toHaveBeenCalledWith(
|
|
expect.objectContaining({
|
|
type: 'gemini',
|
|
text: 'Some response content',
|
|
}),
|
|
expect.any(Number),
|
|
);
|
|
});
|
|
|
|
// Now simulate a new response without a thought
|
|
mockSendMessageStream.mockReturnValue(
|
|
(async function* () {
|
|
yield {
|
|
type: ServerGeminiEventType.Content,
|
|
value: 'New response content',
|
|
};
|
|
yield {
|
|
type: ServerGeminiEventType.Finished,
|
|
value: { reason: 'STOP', usageMetadata: undefined },
|
|
};
|
|
})(),
|
|
);
|
|
|
|
// Submit second query - thought should be reset
|
|
await act(async () => {
|
|
await result.current.submitQuery('Second query');
|
|
});
|
|
|
|
// The thought should be reset to null when starting the new prompt
|
|
// We can verify this by checking that the LoadingIndicator would not show the previous thought
|
|
// The actual thought state is internal to the hook, but we can verify the behavior
|
|
// by ensuring the second response doesn't show the previous thought
|
|
await waitFor(() => {
|
|
expect(mockAddItem).toHaveBeenCalledWith(
|
|
expect.objectContaining({
|
|
type: 'gemini',
|
|
text: 'New response content',
|
|
}),
|
|
expect.any(Number),
|
|
);
|
|
});
|
|
});
|
|
|
|
it('should memoize pendingHistoryItems', () => {
|
|
mockUseToolScheduler.mockReturnValue([
|
|
[],
|
|
mockScheduleToolCalls,
|
|
mockMarkToolsAsSubmitted,
|
|
vi.fn(),
|
|
mockCancelAllToolCalls,
|
|
0,
|
|
]);
|
|
|
|
const { result, rerender } = renderHookWithProviders(() =>
|
|
useGeminiStream(
|
|
mockConfig.getGeminiClient(),
|
|
[],
|
|
mockAddItem,
|
|
mockConfig,
|
|
mockLoadedSettings,
|
|
mockOnDebugMessage,
|
|
mockHandleSlashCommand,
|
|
false,
|
|
() => 'vscode' as EditorType,
|
|
() => {},
|
|
() => Promise.resolve(),
|
|
false,
|
|
() => {},
|
|
() => {},
|
|
() => {},
|
|
80,
|
|
24,
|
|
),
|
|
);
|
|
|
|
const firstResult = result.current.pendingHistoryItems;
|
|
rerender();
|
|
const secondResult = result.current.pendingHistoryItems;
|
|
|
|
expect(firstResult).toStrictEqual(secondResult);
|
|
|
|
const newToolCalls: TrackedToolCall[] = [
|
|
{
|
|
request: { callId: 'call1', name: 'tool1', args: {} },
|
|
status: CoreToolCallStatus.Executing,
|
|
tool: {
|
|
name: 'tool1',
|
|
displayName: 'tool1',
|
|
description: 'desc1',
|
|
build: vi.fn(),
|
|
},
|
|
invocation: {
|
|
getDescription: () => 'Mock description',
|
|
},
|
|
} as unknown as TrackedExecutingToolCall,
|
|
];
|
|
|
|
mockUseToolScheduler.mockReturnValue([
|
|
newToolCalls,
|
|
mockScheduleToolCalls,
|
|
mockMarkToolsAsSubmitted,
|
|
vi.fn(),
|
|
mockCancelAllToolCalls,
|
|
0,
|
|
]);
|
|
|
|
rerender();
|
|
const thirdResult = result.current.pendingHistoryItems;
|
|
|
|
expect(thirdResult).not.toStrictEqual(secondResult);
|
|
});
|
|
|
|
it('should reset thought to null when user cancels', async () => {
|
|
// Mock a stream that yields a thought then gets cancelled
|
|
mockSendMessageStream.mockReturnValue(
|
|
(async function* () {
|
|
yield {
|
|
type: ServerGeminiEventType.Thought,
|
|
value: { subject: 'Some thought', description: 'Description' },
|
|
};
|
|
yield { type: ServerGeminiEventType.UserCancelled };
|
|
})(),
|
|
);
|
|
|
|
const { result } = renderHookWithProviders(() =>
|
|
useGeminiStream(
|
|
new MockedGeminiClientClass(mockConfig),
|
|
[],
|
|
mockAddItem,
|
|
mockConfig,
|
|
mockLoadedSettings,
|
|
mockOnDebugMessage,
|
|
mockHandleSlashCommand,
|
|
false,
|
|
() => 'vscode' as EditorType,
|
|
() => {},
|
|
() => Promise.resolve(),
|
|
false,
|
|
() => {},
|
|
() => {},
|
|
() => {},
|
|
80,
|
|
24,
|
|
),
|
|
);
|
|
|
|
// Submit query
|
|
await act(async () => {
|
|
await result.current.submitQuery('Test query');
|
|
});
|
|
|
|
// Verify cancellation message was added
|
|
await waitFor(() => {
|
|
expect(mockAddItem).toHaveBeenCalledWith(
|
|
expect.objectContaining({
|
|
type: 'info',
|
|
text: 'User cancelled the request.',
|
|
}),
|
|
expect.any(Number),
|
|
);
|
|
});
|
|
|
|
// Verify state is reset to idle
|
|
expect(result.current.streamingState).toBe(StreamingState.Idle);
|
|
});
|
|
|
|
it('should reset thought to null when there is an error', async () => {
|
|
// Mock a stream that yields a thought then encounters an error
|
|
mockSendMessageStream.mockReturnValue(
|
|
(async function* () {
|
|
yield {
|
|
type: ServerGeminiEventType.Thought,
|
|
value: { subject: 'Some thought', description: 'Description' },
|
|
};
|
|
yield {
|
|
type: ServerGeminiEventType.Error,
|
|
value: { error: { message: 'Test error' } },
|
|
};
|
|
})(),
|
|
);
|
|
|
|
const { result } = renderHookWithProviders(() =>
|
|
useGeminiStream(
|
|
new MockedGeminiClientClass(mockConfig),
|
|
[],
|
|
mockAddItem,
|
|
mockConfig,
|
|
mockLoadedSettings,
|
|
mockOnDebugMessage,
|
|
mockHandleSlashCommand,
|
|
false,
|
|
() => 'vscode' as EditorType,
|
|
() => {},
|
|
() => Promise.resolve(),
|
|
false,
|
|
() => {},
|
|
() => {},
|
|
() => {},
|
|
80,
|
|
24,
|
|
),
|
|
);
|
|
|
|
// Submit query
|
|
await act(async () => {
|
|
await result.current.submitQuery('Test query');
|
|
});
|
|
|
|
// Verify error message was added
|
|
await waitFor(() => {
|
|
expect(mockAddItem).toHaveBeenCalledWith(
|
|
expect.objectContaining({
|
|
type: CoreToolCallStatus.Error,
|
|
}),
|
|
expect.any(Number),
|
|
);
|
|
});
|
|
|
|
// Verify parseAndFormatApiError was called
|
|
expect(mockParseAndFormatApiError).toHaveBeenCalledWith(
|
|
{ message: 'Test error' },
|
|
expect.any(String),
|
|
undefined,
|
|
'gemini-2.5-pro',
|
|
'gemini-2.5-flash',
|
|
);
|
|
});
|
|
|
|
it('should update lastOutputTime on Gemini thought and content events', async () => {
|
|
vi.useFakeTimers();
|
|
const startTime = 1000000;
|
|
vi.setSystemTime(startTime);
|
|
|
|
// Mock a stream that yields a thought then content
|
|
mockSendMessageStream.mockReturnValue(
|
|
(async function* () {
|
|
yield {
|
|
type: ServerGeminiEventType.Thought,
|
|
value: { subject: 'Thinking...', description: '' },
|
|
};
|
|
// Advance time for the next event
|
|
vi.advanceTimersByTime(1000);
|
|
yield {
|
|
type: ServerGeminiEventType.Content,
|
|
value: 'Hello',
|
|
};
|
|
})(),
|
|
);
|
|
|
|
const { result } = renderHookWithProviders(() =>
|
|
useGeminiStream(
|
|
new MockedGeminiClientClass(mockConfig),
|
|
[],
|
|
mockAddItem,
|
|
mockConfig,
|
|
mockLoadedSettings,
|
|
mockOnDebugMessage,
|
|
mockHandleSlashCommand,
|
|
false,
|
|
() => 'vscode' as EditorType,
|
|
() => {},
|
|
() => Promise.resolve(),
|
|
false,
|
|
() => {},
|
|
() => {},
|
|
() => {},
|
|
80,
|
|
24,
|
|
),
|
|
);
|
|
|
|
// Submit query
|
|
await act(async () => {
|
|
await result.current.submitQuery('Test query');
|
|
});
|
|
|
|
// Verify lastOutputTime was updated
|
|
// It should be the time of the last event (startTime + 1000)
|
|
expect(result.current.lastOutputTime).toBe(startTime + 1000);
|
|
|
|
vi.useRealTimers();
|
|
});
|
|
});
|
|
|
|
describe('Loop Detection Confirmation', () => {
|
|
beforeEach(() => {
|
|
// Add mock for getLoopDetectionService to the config
|
|
const mockLoopDetectionService = {
|
|
disableForSession: vi.fn(),
|
|
};
|
|
mockConfig.getGeminiClient = vi.fn().mockReturnValue({
|
|
...new MockedGeminiClientClass(mockConfig),
|
|
getLoopDetectionService: () => mockLoopDetectionService,
|
|
});
|
|
});
|
|
|
|
it('should set loopDetectionConfirmationRequest when LoopDetected event is received', async () => {
|
|
mockSendMessageStream.mockReturnValue(
|
|
(async function* () {
|
|
yield {
|
|
type: ServerGeminiEventType.Content,
|
|
value: 'Some content',
|
|
};
|
|
yield {
|
|
type: ServerGeminiEventType.LoopDetected,
|
|
};
|
|
})(),
|
|
);
|
|
|
|
const { result } = renderTestHook();
|
|
|
|
await act(async () => {
|
|
await result.current.submitQuery('test query');
|
|
});
|
|
|
|
await waitFor(() => {
|
|
expect(result.current.loopDetectionConfirmationRequest).not.toBeNull();
|
|
expect(
|
|
typeof result.current.loopDetectionConfirmationRequest?.onComplete,
|
|
).toBe('function');
|
|
});
|
|
});
|
|
|
|
it('should disable loop detection and show message when user selects "disable"', async () => {
|
|
const mockLoopDetectionService = {
|
|
disableForSession: vi.fn(),
|
|
};
|
|
const mockClient = {
|
|
...new MockedGeminiClientClass(mockConfig),
|
|
getLoopDetectionService: () => mockLoopDetectionService,
|
|
};
|
|
mockConfig.getGeminiClient = vi.fn().mockReturnValue(mockClient);
|
|
|
|
// Mock for the initial request
|
|
mockSendMessageStream.mockReturnValueOnce(
|
|
(async function* () {
|
|
yield {
|
|
type: ServerGeminiEventType.LoopDetected,
|
|
};
|
|
})(),
|
|
);
|
|
|
|
// Mock for the retry request
|
|
mockSendMessageStream.mockReturnValueOnce(
|
|
(async function* () {
|
|
yield {
|
|
type: ServerGeminiEventType.Content,
|
|
value: 'Retry successful',
|
|
};
|
|
yield {
|
|
type: ServerGeminiEventType.Finished,
|
|
value: { reason: 'STOP' },
|
|
};
|
|
})(),
|
|
);
|
|
|
|
const { result } = renderTestHook();
|
|
|
|
await act(async () => {
|
|
await result.current.submitQuery('test query');
|
|
});
|
|
|
|
// Wait for confirmation request to be set
|
|
await waitFor(() => {
|
|
expect(result.current.loopDetectionConfirmationRequest).not.toBeNull();
|
|
});
|
|
|
|
// Simulate user selecting "disable"
|
|
await act(async () => {
|
|
result.current.loopDetectionConfirmationRequest?.onComplete({
|
|
userSelection: 'disable',
|
|
});
|
|
});
|
|
|
|
// Verify loop detection was disabled
|
|
expect(mockLoopDetectionService.disableForSession).toHaveBeenCalledTimes(
|
|
1,
|
|
);
|
|
|
|
// Verify confirmation request was cleared
|
|
expect(result.current.loopDetectionConfirmationRequest).toBeNull();
|
|
|
|
// Verify appropriate message was added
|
|
expect(mockAddItem).toHaveBeenCalledWith({
|
|
type: 'info',
|
|
text: 'Loop detection has been disabled for this session. Retrying request...',
|
|
});
|
|
|
|
// Verify that the request was retried
|
|
await waitFor(() => {
|
|
expect(mockSendMessageStream).toHaveBeenCalledTimes(2);
|
|
expect(mockSendMessageStream).toHaveBeenNthCalledWith(
|
|
2,
|
|
'test query',
|
|
expect.any(AbortSignal),
|
|
expect.any(String),
|
|
undefined,
|
|
false,
|
|
'test query',
|
|
);
|
|
});
|
|
});
|
|
|
|
it('should keep loop detection enabled and show message when user selects "keep"', async () => {
|
|
const mockLoopDetectionService = {
|
|
disableForSession: vi.fn(),
|
|
};
|
|
const mockClient = {
|
|
...new MockedGeminiClientClass(mockConfig),
|
|
getLoopDetectionService: () => mockLoopDetectionService,
|
|
};
|
|
mockConfig.getGeminiClient = vi.fn().mockReturnValue(mockClient);
|
|
|
|
mockSendMessageStream.mockReturnValue(
|
|
(async function* () {
|
|
yield {
|
|
type: ServerGeminiEventType.LoopDetected,
|
|
};
|
|
})(),
|
|
);
|
|
|
|
const { result } = renderTestHook();
|
|
|
|
await act(async () => {
|
|
await result.current.submitQuery('test query');
|
|
});
|
|
|
|
// Wait for confirmation request to be set
|
|
await waitFor(() => {
|
|
expect(result.current.loopDetectionConfirmationRequest).not.toBeNull();
|
|
});
|
|
|
|
// Simulate user selecting "keep"
|
|
await act(async () => {
|
|
result.current.loopDetectionConfirmationRequest?.onComplete({
|
|
userSelection: 'keep',
|
|
});
|
|
});
|
|
|
|
// Verify loop detection was NOT disabled
|
|
expect(mockLoopDetectionService.disableForSession).not.toHaveBeenCalled();
|
|
|
|
// Verify confirmation request was cleared
|
|
expect(result.current.loopDetectionConfirmationRequest).toBeNull();
|
|
|
|
// Verify appropriate message was added
|
|
expect(mockAddItem).toHaveBeenCalledWith({
|
|
type: 'info',
|
|
text: 'A potential loop was detected. This can happen due to repetitive tool calls or other model behavior. The request has been halted.',
|
|
});
|
|
|
|
// Verify that the request was NOT retried
|
|
expect(mockSendMessageStream).toHaveBeenCalledTimes(1);
|
|
});
|
|
|
|
it('should handle multiple loop detection events properly', async () => {
|
|
const { result } = renderTestHook();
|
|
|
|
// First loop detection - set up fresh mock for first call
|
|
mockSendMessageStream.mockReturnValueOnce(
|
|
(async function* () {
|
|
yield {
|
|
type: ServerGeminiEventType.LoopDetected,
|
|
};
|
|
})(),
|
|
);
|
|
|
|
// First loop detection
|
|
await act(async () => {
|
|
await result.current.submitQuery('first query');
|
|
});
|
|
|
|
await waitFor(() => {
|
|
expect(result.current.loopDetectionConfirmationRequest).not.toBeNull();
|
|
});
|
|
|
|
// Simulate user selecting "keep" for first request
|
|
await act(async () => {
|
|
result.current.loopDetectionConfirmationRequest?.onComplete({
|
|
userSelection: 'keep',
|
|
});
|
|
});
|
|
|
|
expect(result.current.loopDetectionConfirmationRequest).toBeNull();
|
|
|
|
// Verify first message was added
|
|
expect(mockAddItem).toHaveBeenCalledWith({
|
|
type: 'info',
|
|
text: 'A potential loop was detected. This can happen due to repetitive tool calls or other model behavior. The request has been halted.',
|
|
});
|
|
|
|
// Second loop detection - set up fresh mock for second call
|
|
mockSendMessageStream.mockReturnValueOnce(
|
|
(async function* () {
|
|
yield {
|
|
type: ServerGeminiEventType.LoopDetected,
|
|
};
|
|
})(),
|
|
);
|
|
|
|
// Mock for the retry request
|
|
mockSendMessageStream.mockReturnValueOnce(
|
|
(async function* () {
|
|
yield {
|
|
type: ServerGeminiEventType.Content,
|
|
value: 'Retry successful',
|
|
};
|
|
yield {
|
|
type: ServerGeminiEventType.Finished,
|
|
value: { reason: 'STOP' },
|
|
};
|
|
})(),
|
|
);
|
|
|
|
// Second loop detection
|
|
await act(async () => {
|
|
await result.current.submitQuery('second query');
|
|
});
|
|
|
|
await waitFor(() => {
|
|
expect(result.current.loopDetectionConfirmationRequest).not.toBeNull();
|
|
});
|
|
|
|
// Simulate user selecting "disable" for second request
|
|
await act(async () => {
|
|
result.current.loopDetectionConfirmationRequest?.onComplete({
|
|
userSelection: 'disable',
|
|
});
|
|
});
|
|
|
|
expect(result.current.loopDetectionConfirmationRequest).toBeNull();
|
|
|
|
// Verify second message was added
|
|
expect(mockAddItem).toHaveBeenCalledWith({
|
|
type: 'info',
|
|
text: 'Loop detection has been disabled for this session. Retrying request...',
|
|
});
|
|
|
|
// Verify that the request was retried
|
|
await waitFor(() => {
|
|
expect(mockSendMessageStream).toHaveBeenCalledTimes(3); // 1st query, 2nd query, retry of 2nd query
|
|
expect(mockSendMessageStream).toHaveBeenNthCalledWith(
|
|
3,
|
|
'second query',
|
|
expect.any(AbortSignal),
|
|
expect.any(String),
|
|
undefined,
|
|
false,
|
|
'second query',
|
|
);
|
|
});
|
|
});
|
|
|
|
it('should process LoopDetected event after moving pending history to history', async () => {
|
|
mockSendMessageStream.mockReturnValue(
|
|
(async function* () {
|
|
yield {
|
|
type: ServerGeminiEventType.Content,
|
|
value: 'Some response content',
|
|
};
|
|
yield {
|
|
type: ServerGeminiEventType.LoopDetected,
|
|
};
|
|
})(),
|
|
);
|
|
|
|
const { result } = renderTestHook();
|
|
|
|
await act(async () => {
|
|
await result.current.submitQuery('test query');
|
|
});
|
|
|
|
// Verify that the content was added to history before the loop detection dialog
|
|
await waitFor(() => {
|
|
expect(mockAddItem).toHaveBeenCalledWith(
|
|
expect.objectContaining({
|
|
type: 'gemini',
|
|
text: 'Some response content',
|
|
}),
|
|
expect.any(Number),
|
|
);
|
|
});
|
|
|
|
// Then verify loop detection confirmation request was set
|
|
await waitFor(() => {
|
|
expect(result.current.loopDetectionConfirmationRequest).not.toBeNull();
|
|
});
|
|
});
|
|
|
|
describe('Race Condition Prevention', () => {
|
|
it('should reject concurrent submitQuery when already responding', async () => {
|
|
// Stream that stays open (simulates "still responding")
|
|
mockSendMessageStream.mockReturnValue(
|
|
(async function* () {
|
|
yield {
|
|
type: ServerGeminiEventType.Content,
|
|
value: 'First response',
|
|
};
|
|
// Keep the stream open
|
|
await new Promise(() => {});
|
|
})(),
|
|
);
|
|
|
|
const { result } = 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',
|
|
);
|
|
});
|
|
});
|
|
});
|
|
});
|
|
|
|
describe('Agent Execution Events', () => {
|
|
it('should handle AgentExecutionStopped event with systemMessage', async () => {
|
|
mockSendMessageStream.mockReturnValue(
|
|
(async function* () {
|
|
yield {
|
|
type: ServerGeminiEventType.AgentExecutionStopped,
|
|
value: {
|
|
reason: 'hook-reason',
|
|
systemMessage: 'Custom stop message',
|
|
},
|
|
};
|
|
})(),
|
|
);
|
|
|
|
const { result } = renderTestHook();
|
|
|
|
await act(async () => {
|
|
await result.current.submitQuery('test stop');
|
|
});
|
|
|
|
await waitFor(() => {
|
|
expect(mockAddItem).toHaveBeenCalledWith(
|
|
{
|
|
type: MessageType.INFO,
|
|
text: 'Agent execution stopped: Custom stop message',
|
|
},
|
|
expect.any(Number),
|
|
);
|
|
expect(result.current.streamingState).toBe(StreamingState.Idle);
|
|
});
|
|
});
|
|
|
|
it('should handle AgentExecutionStopped event by falling back to reason when systemMessage is missing', async () => {
|
|
mockSendMessageStream.mockReturnValue(
|
|
(async function* () {
|
|
yield {
|
|
type: ServerGeminiEventType.AgentExecutionStopped,
|
|
value: { reason: 'Stopped by hook' },
|
|
};
|
|
})(),
|
|
);
|
|
|
|
const { result } = renderTestHook();
|
|
|
|
await act(async () => {
|
|
await result.current.submitQuery('test stop');
|
|
});
|
|
|
|
await waitFor(() => {
|
|
expect(mockAddItem).toHaveBeenCalledWith(
|
|
{
|
|
type: MessageType.INFO,
|
|
text: 'Agent execution stopped: Stopped by hook',
|
|
},
|
|
expect.any(Number),
|
|
);
|
|
expect(result.current.streamingState).toBe(StreamingState.Idle);
|
|
});
|
|
});
|
|
|
|
it('should handle AgentExecutionBlocked event with systemMessage', async () => {
|
|
mockSendMessageStream.mockReturnValue(
|
|
(async function* () {
|
|
yield {
|
|
type: ServerGeminiEventType.AgentExecutionBlocked,
|
|
value: {
|
|
reason: 'hook-reason',
|
|
systemMessage: 'Custom block message',
|
|
},
|
|
};
|
|
})(),
|
|
);
|
|
|
|
const { result } = renderTestHook();
|
|
|
|
await act(async () => {
|
|
await result.current.submitQuery('test block');
|
|
});
|
|
|
|
await waitFor(() => {
|
|
expect(mockAddItem).toHaveBeenCalledWith(
|
|
{
|
|
type: MessageType.WARNING,
|
|
text: 'Agent execution blocked: Custom block message',
|
|
},
|
|
expect.any(Number),
|
|
);
|
|
});
|
|
});
|
|
|
|
it('should handle AgentExecutionBlocked event by falling back to reason when systemMessage is missing', async () => {
|
|
mockSendMessageStream.mockReturnValue(
|
|
(async function* () {
|
|
yield {
|
|
type: ServerGeminiEventType.AgentExecutionBlocked,
|
|
value: { reason: 'Blocked by hook' },
|
|
};
|
|
})(),
|
|
);
|
|
|
|
const { result } = renderTestHook();
|
|
|
|
await act(async () => {
|
|
await result.current.submitQuery('test block');
|
|
});
|
|
|
|
await waitFor(() => {
|
|
expect(mockAddItem).toHaveBeenCalledWith(
|
|
{
|
|
type: MessageType.WARNING,
|
|
text: 'Agent execution blocked: Blocked by hook',
|
|
},
|
|
expect.any(Number),
|
|
);
|
|
});
|
|
});
|
|
});
|
|
|
|
describe('Stream Splitting', () => {
|
|
it('should not add empty history item when splitting message results in empty or whitespace-only beforeText', async () => {
|
|
// Mock split point to always be 0, causing beforeText to be empty
|
|
vi.mocked(findLastSafeSplitPoint).mockReturnValue(0);
|
|
|
|
mockSendMessageStream.mockReturnValue(
|
|
(async function* () {
|
|
yield { type: ServerGeminiEventType.Content, value: 'test content' };
|
|
})(),
|
|
);
|
|
|
|
const { result } = 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);
|
|
});
|
|
});
|
|
|
|
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 });
|
|
});
|
|
expect(spanMetadata.input).toBe('telemetry test query');
|
|
});
|
|
});
|