File size: 4,251 Bytes
7a1ad33 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | /**
* @license
* Copyright 2025 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { describe, it, expect, beforeEach, vi } from 'vitest';
import { Config } from './config.js';
import { DEFAULT_GEMINI_MODEL, DEFAULT_GEMINI_FLASH_MODEL } from './models.js';
import { logFlashFallback } from '../telemetry/loggers.js';
import { FlashFallbackEvent } from '../telemetry/types.js';
import fs from 'node:fs';
vi.mock('node:fs');
vi.mock('../telemetry/loggers.js', () => ({
logFlashFallback: vi.fn(),
logRipgrepFallback: vi.fn(),
}));
describe('Flash Model Fallback Configuration', () => {
let config: Config;
beforeEach(() => {
vi.mocked(fs.existsSync).mockReturnValue(true);
vi.mocked(fs.statSync).mockReturnValue({
isDirectory: () => true,
} as fs.Stats);
config = new Config({
sessionId: 'test-session',
targetDir: '/test',
debugMode: false,
cwd: '/test',
model: DEFAULT_GEMINI_MODEL,
});
// Initialize contentGeneratorConfig for testing
(
config as unknown as { contentGeneratorConfig: unknown }
).contentGeneratorConfig = {
model: DEFAULT_GEMINI_MODEL,
authType: 'oauth-personal',
};
});
describe('getModel', () => {
it('should return contentGeneratorConfig model if available', () => {
// Simulate initialized content generator config
config.setModel(DEFAULT_GEMINI_FLASH_MODEL);
expect(config.getModel()).toBe(DEFAULT_GEMINI_FLASH_MODEL);
});
it('should fall back to initial model if contentGeneratorConfig is not available', () => {
// Test with fresh config where contentGeneratorConfig might not be set
const newConfig = new Config({
sessionId: 'test-session-2',
targetDir: '/test',
debugMode: false,
cwd: '/test',
model: 'custom-model',
});
expect(newConfig.getModel()).toBe('custom-model');
});
});
describe('activateFallbackMode', () => {
it('should set model to fallback and log event', () => {
config.activateFallbackMode(DEFAULT_GEMINI_FLASH_MODEL);
expect(config.getModel()).toBe(DEFAULT_GEMINI_FLASH_MODEL);
expect(logFlashFallback).toHaveBeenCalledWith(
config,
expect.any(FlashFallbackEvent),
);
});
it('should set fallback override when failedModel is provided and register runtime override', () => {
config.activateFallbackMode(
DEFAULT_GEMINI_FLASH_MODEL,
DEFAULT_GEMINI_MODEL,
);
expect(config.getModel()).toBe(DEFAULT_GEMINI_FLASH_MODEL);
expect(config.getFallbackOverride(DEFAULT_GEMINI_MODEL)).toBe(
DEFAULT_GEMINI_FLASH_MODEL,
);
// Verify it registers the runtime model override with ModelConfigService
expect(
config
.getModelConfigService()
.getResolvedConfig({ model: DEFAULT_GEMINI_MODEL }).model,
).toBe(DEFAULT_GEMINI_FLASH_MODEL);
});
it('should flatten override chains when a model that was previously a target fails', () => {
// 1. Initial fallback: A -> B
config.activateFallbackMode('model-B', 'model-A');
expect(config.getFallbackOverride('model-A')).toBe('model-B');
expect(
config.getModelConfigService().getResolvedConfig({ model: 'model-A' })
.model,
).toBe('model-B');
// 2. Chained fallback: B fails, fallback to C
// This should update A -> C as well.
config.activateFallbackMode('model-C', 'model-B');
expect(config.getFallbackOverride('model-A')).toBe('model-C');
expect(config.getFallbackOverride('model-B')).toBe('model-C');
expect(
config.getModelConfigService().getResolvedConfig({ model: 'model-A' })
.model,
).toBe('model-C');
expect(
config.getModelConfigService().getResolvedConfig({ model: 'model-B' })
.model,
).toBe('model-C');
});
it('should not reset availability service if model has not changed', () => {
const resetSpy = vi.spyOn(config.getModelAvailabilityService(), 'reset');
const currentModel = config.getActiveModel();
config.activateFallbackMode(currentModel);
expect(resetSpy).not.toHaveBeenCalled();
});
});
});
|