quantum-cli-core
Version:
Quantum CLI Core - Multi-LLM Collaboration System
671 lines • 32.9 kB
JavaScript
/**
* @license
* Copyright 2025 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
import { GoogleGenAI, } from '@google/genai';
import { GeminiClient } from './client.js';
import { findIndexAfterFraction } from '../utils/model-utils.js';
import { AuthType } from './contentGenerator.js';
import { Config } from '../config/config.js';
import { Turn } from './turn.js';
import { getCoreSystemPrompt } from './prompts.js';
import { DEFAULT_GEMINI_FLASH_MODEL } from '../config/models.js';
import { FileDiscoveryService } from '../services/fileDiscoveryService.js';
import { setSimulate429 } from '../utils/testUtils.js';
import { tokenLimit } from './tokenLimits.js';
// --- Mocks ---
const mockChatCreateFn = vi.fn();
const mockGenerateContentFn = vi.fn();
const mockEmbedContentFn = vi.fn();
const mockTurnRunFn = vi.fn();
vi.mock('@google/genai');
vi.mock('./turn', () => {
// Define a mock class that has the same shape as the real Turn
class MockTurn {
pendingToolCalls = [];
// The run method is a property that holds our mock function
run = mockTurnRunFn;
constructor() {
// The constructor can be empty or do some mock setup
}
}
// Export the mock class as 'Turn'
return { Turn: MockTurn };
});
vi.mock('../config/config.js');
vi.mock('./prompts');
vi.mock('../utils/getFolderStructure', () => ({
getFolderStructure: vi.fn().mockResolvedValue('Mock Folder Structure'),
}));
vi.mock('../utils/errorReporting', () => ({ reportError: vi.fn() }));
vi.mock('../utils/nextSpeakerChecker', () => ({
checkNextSpeaker: vi.fn().mockResolvedValue(null),
}));
vi.mock('../utils/generateContentResponseUtilities', () => ({
getResponseText: (result) => result.candidates?.[0]?.content?.parts?.map((part) => part.text).join('') ||
undefined,
}));
vi.mock('../telemetry/index.js', () => ({
logApiRequest: vi.fn(),
logApiResponse: vi.fn(),
logApiError: vi.fn(),
}));
describe('findIndexAfterFraction', () => {
const history = [
{ role: 'user', parts: [{ text: 'This is the first message.' }] },
{ role: 'model', parts: [{ text: 'This is the second message.' }] },
{ role: 'user', parts: [{ text: 'This is the third message.' }] },
{ role: 'model', parts: [{ text: 'This is the fourth message.' }] },
{ role: 'user', parts: [{ text: 'This is the fifth message.' }] },
];
it('should throw an error for non-positive numbers', () => {
expect(() => findIndexAfterFraction(history, 0)).toThrow('Fraction must be between 0 and 1');
});
it('should throw an error for a fraction greater than or equal to 1', () => {
expect(() => findIndexAfterFraction(history, 1)).toThrow('Fraction must be between 0 and 1');
});
it('should handle a fraction in the middle', () => {
// Total length is 257. 257 * 0.5 = 128.5
// 0: 53
// 1: 53 + 54 = 107
// 2: 107 + 53 = 160
// 160 >= 128.5, so index is 2
expect(findIndexAfterFraction(history, 0.5)).toBe(2);
});
it('should handle an empty history', () => {
expect(findIndexAfterFraction([], 0.5)).toBe(0);
});
it('should handle a history with only one item', () => {
expect(findIndexAfterFraction(history.slice(0, 1), 0.5)).toBe(0);
});
it('should handle history with weird parts', () => {
const historyWithEmptyParts = [
{ role: 'user', parts: [{ text: 'Message 1' }] },
{ role: 'model', parts: [{ fileData: { fileUri: 'derp' } }] },
{ role: 'user', parts: [{ text: 'Message 2' }] },
];
expect(findIndexAfterFraction(historyWithEmptyParts, 0.5)).toBe(1);
});
});
describe('Gemini Client (client.ts)', () => {
let client;
beforeEach(async () => {
vi.resetAllMocks();
// Disable 429 simulation for tests
setSimulate429(false);
// Set up the mock for GoogleGenAI constructor and its methods
const MockedGoogleGenAI = vi.mocked(GoogleGenAI);
MockedGoogleGenAI.mockImplementation(() => {
const mock = {
chats: { create: mockChatCreateFn },
models: {
generateContent: mockGenerateContentFn,
embedContent: mockEmbedContentFn,
},
};
return mock;
});
mockChatCreateFn.mockResolvedValue({});
mockGenerateContentFn.mockResolvedValue({
candidates: [
{
content: {
parts: [{ text: '{"key": "value"}' }],
},
},
],
});
// Because the GeminiClient constructor kicks off an async process (startChat)
// that depends on a fully-formed Config object, we need to mock the
// entire implementation of Config for these tests.
const mockToolRegistry = {
getFunctionDeclarations: vi.fn().mockReturnValue([]),
getTool: vi.fn().mockReturnValue(null),
};
const fileService = new FileDiscoveryService('/test/dir');
const MockedConfig = vi.mocked(Config, true);
const contentGeneratorConfig = {
model: 'test-model',
apiKey: 'test-key',
vertexai: false,
authType: AuthType.USE_GEMINI,
};
MockedConfig.mockImplementation(() => {
const mock = {
getContentGeneratorConfig: vi
.fn()
.mockReturnValue(contentGeneratorConfig),
getToolRegistry: vi.fn().mockResolvedValue(mockToolRegistry),
getModel: vi.fn().mockReturnValue('test-model'),
getEmbeddingModel: vi.fn().mockReturnValue('test-embedding-model'),
getApiKey: vi.fn().mockReturnValue('test-key'),
getVertexAI: vi.fn().mockReturnValue(false),
getUserAgent: vi.fn().mockReturnValue('test-agent'),
getUserMemory: vi.fn().mockReturnValue(''),
getFullContext: vi.fn().mockReturnValue(false),
getSessionId: vi.fn().mockReturnValue('test-session-id'),
getProxy: vi.fn().mockReturnValue(undefined),
getWorkingDir: vi.fn().mockReturnValue('/test/dir'),
getFileService: vi.fn().mockReturnValue(fileService),
getCollaborationConfig: vi.fn().mockReturnValue({ enabled: false }),
};
return mock;
});
// We can instantiate the client here since Config is mocked
// and the constructor will use the mocked GoogleGenAI
const mockConfig = new Config({});
client = new GeminiClient(mockConfig);
await client.initialize(contentGeneratorConfig);
});
afterEach(() => {
vi.restoreAllMocks();
});
// NOTE: The following tests for startChat were removed due to persistent issues with
// the @google/genai mock. Specifically, the mockChatCreateFn (representing instance.chats.create)
// was not being detected as called by the GeminiClient instance.
// This likely points to a subtle issue in how the GoogleGenerativeAI class constructor
// and its instance methods are mocked and then used by the class under test.
// For future debugging, ensure that the `this.client` in `GeminiClient` (which is an
// instance of the mocked GoogleGenerativeAI) correctly has its `chats.create` method
// pointing to `mockChatCreateFn`.
// it('startChat should call getCoreSystemPrompt with userMemory and pass to chats.create', async () => { ... });
// it('startChat should call getCoreSystemPrompt with empty string if userMemory is empty', async () => { ... });
// NOTE: The following tests for generateJson were removed due to persistent issues with
// the @google/genai mock, similar to the startChat tests. The mockGenerateContentFn
// (representing instance.models.generateContent) was not being detected as called, or the mock
// was not preventing an actual API call (leading to API key errors).
// For future debugging, ensure `this.client.models.generateContent` in `GeminiClient` correctly
// uses the `mockGenerateContentFn`.
// it('generateJson should call getCoreSystemPrompt with userMemory and pass to generateContent', async () => { ... });
// it('generateJson should call getCoreSystemPrompt with empty string if userMemory is empty', async () => { ... });
describe('generateEmbedding', () => {
const texts = ['hello world', 'goodbye world'];
const testEmbeddingModel = 'test-embedding-model';
it('should call embedContent with correct parameters and return embeddings', async () => {
const mockEmbeddings = [
[0.1, 0.2, 0.3],
[0.4, 0.5, 0.6],
];
const mockResponse = {
embeddings: [
{ values: mockEmbeddings[0] },
{ values: mockEmbeddings[1] },
],
};
mockEmbedContentFn.mockResolvedValue(mockResponse);
const result = await client.generateEmbedding(texts);
expect(mockEmbedContentFn).toHaveBeenCalledTimes(1);
expect(mockEmbedContentFn).toHaveBeenCalledWith({
model: testEmbeddingModel,
contents: texts,
});
expect(result).toEqual(mockEmbeddings);
});
it('should return an empty array if an empty array is passed', async () => {
const result = await client.generateEmbedding([]);
expect(result).toEqual([]);
expect(mockEmbedContentFn).not.toHaveBeenCalled();
});
it('should throw an error if API response has no embeddings array', async () => {
mockEmbedContentFn.mockResolvedValue({}); // No `embeddings` key
await expect(client.generateEmbedding(texts)).rejects.toThrow('No embeddings found in API response.');
});
it('should throw an error if API response has an empty embeddings array', async () => {
const mockResponse = {
embeddings: [],
};
mockEmbedContentFn.mockResolvedValue(mockResponse);
await expect(client.generateEmbedding(texts)).rejects.toThrow('No embeddings found in API response.');
});
it('should throw an error if API returns a mismatched number of embeddings', async () => {
const mockResponse = {
embeddings: [{ values: [1, 2, 3] }], // Only one for two texts
};
mockEmbedContentFn.mockResolvedValue(mockResponse);
await expect(client.generateEmbedding(texts)).rejects.toThrow('API returned a mismatched number of embeddings. Expected 2, got 1.');
});
it('should throw an error if any embedding has nullish values', async () => {
const mockResponse = {
embeddings: [{ values: [1, 2, 3] }, { values: undefined }], // Second one is bad
};
mockEmbedContentFn.mockResolvedValue(mockResponse);
await expect(client.generateEmbedding(texts)).rejects.toThrow('API returned an empty embedding for input text at index 1: "goodbye world"');
});
it('should throw an error if any embedding has an empty values array', async () => {
const mockResponse = {
embeddings: [{ values: [] }, { values: [1, 2, 3] }], // First one is bad
};
mockEmbedContentFn.mockResolvedValue(mockResponse);
await expect(client.generateEmbedding(texts)).rejects.toThrow('API returned an empty embedding for input text at index 0: "hello world"');
});
it('should propagate errors from the API call', async () => {
const apiError = new Error('API Failure');
mockEmbedContentFn.mockRejectedValue(apiError);
await expect(client.generateEmbedding(texts)).rejects.toThrow('API Failure');
});
});
describe('generateContent', () => {
it('should call generateContent with the correct parameters', async () => {
const contents = [{ role: 'user', parts: [{ text: 'hello' }] }];
const generationConfig = { temperature: 0.5 };
const abortSignal = new AbortController().signal;
// Mock countTokens
const mockGenerator = {
countTokens: vi.fn().mockResolvedValue({ totalTokens: 1 }),
generateContent: mockGenerateContentFn,
};
client['contentGenerator'] = mockGenerator;
await client.generateContent(contents, generationConfig, abortSignal);
expect(mockGenerateContentFn).toHaveBeenCalledWith({
model: 'test-model',
config: {
abortSignal,
systemInstruction: getCoreSystemPrompt(''),
temperature: 0.5,
topP: 1,
},
contents,
});
});
});
describe('generateJson', () => {
it('should call generateContent with the correct parameters', async () => {
const contents = [{ role: 'user', parts: [{ text: 'hello' }] }];
const schema = { type: 'string' };
const abortSignal = new AbortController().signal;
// Mock countTokens
const mockGenerator = {
countTokens: vi.fn().mockResolvedValue({ totalTokens: 1 }),
generateContent: mockGenerateContentFn,
};
client['contentGenerator'] = mockGenerator;
await client.generateJson(contents, schema, abortSignal);
expect(mockGenerateContentFn).toHaveBeenCalledWith({
model: DEFAULT_GEMINI_FLASH_MODEL,
config: {
abortSignal,
systemInstruction: getCoreSystemPrompt(''),
temperature: 0,
topP: 1,
responseSchema: schema,
responseMimeType: 'application/json',
},
contents,
});
});
});
describe('addHistory', () => {
it('should call chat.addHistory with the provided content', async () => {
const mockChat = {
addHistory: vi.fn(),
};
// eslint-disable-next-line @typescript-eslint/no-explicit-any
client['chat'] = mockChat;
const newContent = {
role: 'user',
parts: [{ text: 'New history item' }],
};
await client.addHistory(newContent);
expect(mockChat.addHistory).toHaveBeenCalledWith(newContent);
});
});
describe('resetChat', () => {
it('should create a new chat session, clearing the old history', async () => {
// 1. Get the initial chat instance and add some history.
const initialChat = client.getChat();
const initialHistory = await client.getHistory();
await client.addHistory({
role: 'user',
parts: [{ text: 'some old message' }],
});
const historyWithOldMessage = await client.getHistory();
expect(historyWithOldMessage.length).toBeGreaterThan(initialHistory.length);
// 2. Call resetChat.
await client.resetChat();
// 3. Get the new chat instance and its history.
const newChat = client.getChat();
const newHistory = await client.getHistory();
// 4. Assert that the chat instance is new and the history is reset.
expect(newChat).not.toBe(initialChat);
expect(newHistory.length).toBe(initialHistory.length);
expect(JSON.stringify(newHistory)).not.toContain('some old message');
});
});
describe('tryCompressChat', () => {
const mockCountTokens = vi.fn();
const mockSendMessage = vi.fn();
beforeEach(() => {
vi.mock('./tokenLimits', () => ({
tokenLimit: vi.fn(),
}));
const mockGenerator = {
countTokens: mockCountTokens,
};
client['contentGenerator'] = mockGenerator;
// Mock the chat's sendMessage method
const mockChat = {
getHistory: vi
.fn()
.mockReturnValue([
{ role: 'user', parts: [{ text: '...history...' }] },
]),
addHistory: vi.fn(),
setHistory: vi.fn(),
sendMessage: mockSendMessage,
};
client['chat'] = mockChat;
});
it('should not trigger summarization if token count is below threshold', async () => {
const MOCKED_TOKEN_LIMIT = 1000;
vi.mocked(tokenLimit).mockReturnValue(MOCKED_TOKEN_LIMIT);
mockCountTokens.mockResolvedValue({
totalTokens: MOCKED_TOKEN_LIMIT * 0.699, // TOKEN_THRESHOLD_FOR_SUMMARIZATION = 0.7
});
const initialChat = client.getChat();
const result = await client.tryCompressChat();
const newChat = client.getChat();
expect(tokenLimit).toHaveBeenCalled();
expect(result).toBeNull();
expect(newChat).toBe(initialChat);
});
it('should trigger summarization if token count is at threshold', async () => {
const MOCKED_TOKEN_LIMIT = 1000;
vi.mocked(tokenLimit).mockReturnValue(MOCKED_TOKEN_LIMIT);
const originalTokenCount = 1000 * 0.7;
const newTokenCount = 100;
mockCountTokens
.mockResolvedValueOnce({ totalTokens: originalTokenCount }) // First call for the check
.mockResolvedValueOnce({ totalTokens: newTokenCount }); // Second call for the new history
// Mock the summary response from the chat
mockSendMessage.mockResolvedValue({
role: 'model',
parts: [{ text: 'This is a summary.' }],
});
const initialChat = client.getChat();
const result = await client.tryCompressChat();
const newChat = client.getChat();
expect(tokenLimit).toHaveBeenCalled();
expect(mockSendMessage).toHaveBeenCalled();
// Assert that summarization happened and returned the correct stats
expect(result).toEqual({
originalTokenCount,
newTokenCount,
});
// Assert that the chat was reset
expect(newChat).not.toBe(initialChat);
});
it('should always trigger summarization when force is true, regardless of token count', async () => {
const originalTokenCount = 10; // Well below threshold
const newTokenCount = 5;
mockCountTokens
.mockResolvedValueOnce({ totalTokens: originalTokenCount })
.mockResolvedValueOnce({ totalTokens: newTokenCount });
// Mock the summary response from the chat
mockSendMessage.mockResolvedValue({
role: 'model',
parts: [{ text: 'This is a summary.' }],
});
const initialChat = client.getChat();
const result = await client.tryCompressChat(true); // force = true
const newChat = client.getChat();
expect(mockSendMessage).toHaveBeenCalled();
expect(result).toEqual({
originalTokenCount,
newTokenCount,
});
// Assert that the chat was reset
expect(newChat).not.toBe(initialChat);
});
});
describe('sendMessageStream', () => {
it('should return the turn instance after the stream is complete', async () => {
// Arrange
const mockStream = (async function* () {
yield { type: 'content', value: 'Hello' };
})();
mockTurnRunFn.mockReturnValue(mockStream);
const mockChat = {
addHistory: vi.fn(),
getHistory: vi.fn().mockReturnValue([]),
};
client['chat'] = mockChat;
const mockGenerator = {
countTokens: vi.fn().mockResolvedValue({ totalTokens: 0 }),
};
client['contentGenerator'] = mockGenerator;
// Act
const stream = client.sendMessageStream([{ text: 'Hi' }], new AbortController().signal);
// Consume the stream manually to get the final return value.
let finalResult;
while (true) {
const result = await stream.next();
if (result.done) {
finalResult = result.value;
break;
}
}
// Assert
expect(finalResult).toBeInstanceOf(Turn);
});
it('should stop infinite loop after MAX_TURNS when nextSpeaker always returns model', async () => {
// Get the mocked checkNextSpeaker function and configure it to trigger infinite loop
const { checkNextSpeaker } = await import('../utils/nextSpeakerChecker.js');
const mockCheckNextSpeaker = vi.mocked(checkNextSpeaker);
mockCheckNextSpeaker.mockResolvedValue({
next_speaker: 'model',
reasoning: 'Test case - always continue',
});
// Mock Turn to have no pending tool calls (which would allow nextSpeaker check)
const mockStream = (async function* () {
yield { type: 'content', value: 'Continue...' };
})();
mockTurnRunFn.mockReturnValue(mockStream);
const mockChat = {
addHistory: vi.fn(),
getHistory: vi.fn().mockReturnValue([]),
};
client['chat'] = mockChat;
const mockGenerator = {
countTokens: vi.fn().mockResolvedValue({ totalTokens: 0 }),
};
client['contentGenerator'] = mockGenerator;
// Use a signal that never gets aborted
const abortController = new AbortController();
const signal = abortController.signal;
// Act - Start the stream that should loop
const stream = client.sendMessageStream([{ text: 'Start conversation' }], signal);
// Count how many stream events we get
let eventCount = 0;
let finalResult;
// Consume the stream and count iterations
while (true) {
const result = await stream.next();
if (result.done) {
finalResult = result.value;
break;
}
eventCount++;
// Safety check to prevent actual infinite loop in test
if (eventCount > 200) {
abortController.abort();
throw new Error('Test exceeded expected event limit - possible actual infinite loop');
}
}
// Assert
expect(finalResult).toBeInstanceOf(Turn);
// Debug: Check how many times checkNextSpeaker was called
const callCount = mockCheckNextSpeaker.mock.calls.length;
// If infinite loop protection is working, checkNextSpeaker should be called many times
// but stop at MAX_TURNS (100). Since each recursive call should trigger checkNextSpeaker,
// we expect it to be called multiple times before hitting the limit
expect(mockCheckNextSpeaker).toHaveBeenCalled();
// The test should demonstrate that the infinite loop protection works:
// - If checkNextSpeaker is called many times (close to MAX_TURNS), it shows the loop was happening
// - If it's only called once, the recursive behavior might not be triggered
if (callCount === 0) {
throw new Error('checkNextSpeaker was never called - the recursive condition was not met');
}
else if (callCount === 1) {
// This might be expected behavior if the turn has pending tool calls or other conditions prevent recursion
console.log('checkNextSpeaker called only once - no infinite loop occurred');
}
else {
console.log(`checkNextSpeaker called ${callCount} times - infinite loop protection worked`);
// If called multiple times, we expect it to be stopped before MAX_TURNS
expect(callCount).toBeLessThanOrEqual(100); // Should not exceed MAX_TURNS
}
// The stream should produce events and eventually terminate
expect(eventCount).toBeGreaterThanOrEqual(1);
expect(eventCount).toBeLessThan(200); // Should not exceed our safety limit
});
it('should respect MAX_TURNS limit even when turns parameter is set to a large value', async () => {
// This test verifies that the infinite loop protection works even when
// someone tries to bypass it by calling with a very large turns value
// Get the mocked checkNextSpeaker function and configure it to trigger infinite loop
const { checkNextSpeaker } = await import('../utils/nextSpeakerChecker.js');
const mockCheckNextSpeaker = vi.mocked(checkNextSpeaker);
mockCheckNextSpeaker.mockResolvedValue({
next_speaker: 'model',
reasoning: 'Test case - always continue',
});
// Mock Turn to have no pending tool calls (which would allow nextSpeaker check)
const mockStream = (async function* () {
yield { type: 'content', value: 'Continue...' };
})();
mockTurnRunFn.mockReturnValue(mockStream);
const mockChat = {
addHistory: vi.fn(),
getHistory: vi.fn().mockReturnValue([]),
};
client['chat'] = mockChat;
const mockGenerator = {
countTokens: vi.fn().mockResolvedValue({ totalTokens: 0 }),
};
client['contentGenerator'] = mockGenerator;
// Use a signal that never gets aborted
const abortController = new AbortController();
const signal = abortController.signal;
// Act - Start the stream with an extremely high turns value
// This simulates a case where the turns protection is bypassed
const stream = client.sendMessageStream([{ text: 'Start conversation' }], signal, Number.MAX_SAFE_INTEGER);
// Count how many stream events we get
let eventCount = 0;
const maxTestIterations = 1000; // Higher limit to show the loop continues
// Consume the stream and count iterations
try {
while (true) {
const result = await stream.next();
if (result.done) {
break;
}
eventCount++;
// This test should hit this limit, demonstrating the infinite loop
if (eventCount > maxTestIterations) {
abortController.abort();
// This is the expected behavior - we hit the infinite loop
break;
}
}
}
catch (error) {
// If the test framework times out, that also demonstrates the infinite loop
console.error('Test timed out or errored:', error);
}
// Assert that the fix works - the loop should stop at MAX_TURNS
const callCount = mockCheckNextSpeaker.mock.calls.length;
// With the fix: even when turns is set to a very high value,
// the loop should stop at MAX_TURNS (100)
expect(callCount).toBeLessThanOrEqual(100); // Should not exceed MAX_TURNS
expect(eventCount).toBeLessThanOrEqual(200); // Should have reasonable number of events
console.log(`Infinite loop protection working: checkNextSpeaker called ${callCount} times, ` +
`${eventCount} events generated (properly bounded by MAX_TURNS)`);
});
});
describe('generateContent', () => {
it('should use current model from config for content generation', async () => {
const initialModel = client['config'].getModel();
const contents = [{ role: 'user', parts: [{ text: 'test' }] }];
const currentModel = initialModel + '-changed';
vi.spyOn(client['config'], 'getModel').mockReturnValueOnce(currentModel);
const mockGenerator = {
countTokens: vi.fn().mockResolvedValue({ totalTokens: 1 }),
generateContent: mockGenerateContentFn,
};
client['contentGenerator'] = mockGenerator;
await client.generateContent(contents, {}, new AbortController().signal);
expect(mockGenerateContentFn).not.toHaveBeenCalledWith({
model: initialModel,
config: expect.any(Object),
contents,
});
expect(mockGenerateContentFn).toHaveBeenCalledWith({
model: currentModel,
config: expect.any(Object),
contents,
});
});
});
describe('tryCompressChat', () => {
it('should use current model from config for token counting after sendMessage', async () => {
const initialModel = client['config'].getModel();
const mockCountTokens = vi
.fn()
.mockResolvedValueOnce({ totalTokens: 100000 })
.mockResolvedValueOnce({ totalTokens: 5000 });
const mockSendMessage = vi.fn().mockResolvedValue({ text: 'Summary' });
const mockChatHistory = [
{ role: 'user', parts: [{ text: 'Long conversation' }] },
{ role: 'model', parts: [{ text: 'Long response' }] },
];
const mockChat = {
getHistory: vi.fn().mockReturnValue(mockChatHistory),
setHistory: vi.fn(),
sendMessage: mockSendMessage,
};
const mockGenerator = {
countTokens: mockCountTokens,
};
// mock the model has been changed between calls of `countTokens`
const firstCurrentModel = initialModel + '-changed-1';
const secondCurrentModel = initialModel + '-changed-2';
vi.spyOn(client['config'], 'getModel')
.mockReturnValueOnce(firstCurrentModel)
.mockReturnValueOnce(secondCurrentModel);
client['chat'] = mockChat;
client['contentGenerator'] = mockGenerator;
client['startChat'] = vi.fn().mockResolvedValue(mockChat);
const result = await client.tryCompressChat(true);
expect(mockCountTokens).toHaveBeenCalledTimes(2);
expect(mockCountTokens).toHaveBeenNthCalledWith(1, {
model: firstCurrentModel,
contents: mockChatHistory,
});
expect(mockCountTokens).toHaveBeenNthCalledWith(2, {
model: secondCurrentModel,
contents: expect.any(Array),
});
expect(result).toEqual({
originalTokenCount: 100000,
newTokenCount: 5000,
});
});
});
describe('handleFlashFallback', () => {
it('should use current model from config when checking for fallback', async () => {
const initialModel = client['config'].getModel();
const fallbackModel = DEFAULT_GEMINI_FLASH_MODEL;
// mock config been changed
const currentModel = initialModel + '-changed';
vi.spyOn(client['config'], 'getModel').mockReturnValueOnce(currentModel);
const mockFallbackHandler = vi.fn().mockResolvedValue(true);
client['config'].flashFallbackHandler = mockFallbackHandler;
client['config'].setModel = vi.fn();
const result = await client['handleFlashFallback'](AuthType.LOGIN_WITH_GOOGLE);
expect(result).toBe(fallbackModel);
expect(mockFallbackHandler).toHaveBeenCalledWith(currentModel, fallbackModel);
});
});
});
//# sourceMappingURL=client.test.js.map