From 16e06adb461d135d5a5b490062439c145005e5fb Mon Sep 17 00:00:00 2001 From: Sehoon Shon Date: Fri, 12 Dec 2025 09:43:46 -0500 Subject: [PATCH] =?UTF-8?q?Use=20Preview=20Flash=20model=20if=20the=20main?= =?UTF-8?q?=20model=20is=20a=20preview=20model=20for=20code=E2=80=A6=20(#8?= =?UTF-8?q?3)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Use Preview Flash model if the main model is a preview model for codebase investigator * fix test * Generate settings doc --- docs/get-started/configuration.md | 2 +- packages/cli/src/config/settingsSchema.ts | 4 +- packages/core/src/agents/registry.test.ts | 60 +++++++++++++++++++++-- packages/core/src/agents/registry.ts | 39 +++++++-------- packages/core/src/config/config.test.ts | 4 +- packages/core/src/config/config.ts | 1 - packages/core/src/config/models.ts | 16 ++++++ schemas/settings.schema.json | 4 +- 8 files changed, 96 insertions(+), 34 deletions(-) diff --git a/docs/get-started/configuration.md b/docs/get-started/configuration.md index 85aa188406..07024ff5dc 100644 --- a/docs/get-started/configuration.md +++ b/docs/get-started/configuration.md @@ -832,7 +832,7 @@ their corresponding top-level category object in your `settings.json` file. - **`experimental.codebaseInvestigatorSettings.model`** (string): - **Description:** The model to use for the Codebase Investigator agent. - - **Default:** `"gemini-2.5-pro"` + - **Default:** `"auto"` - **Requires restart:** Yes #### `hooks` diff --git a/packages/cli/src/config/settingsSchema.ts b/packages/cli/src/config/settingsSchema.ts index a2e7860d54..5b8f75af60 100644 --- a/packages/cli/src/config/settingsSchema.ts +++ b/packages/cli/src/config/settingsSchema.ts @@ -19,7 +19,7 @@ import { DEFAULT_TRUNCATE_TOOL_OUTPUT_LINES, DEFAULT_TRUNCATE_TOOL_OUTPUT_THRESHOLD, DEFAULT_MODEL_CONFIGS, - DEFAULT_GEMINI_MODEL, + GEMINI_MODEL_ALIAS_AUTO, } from '@google/gemini-cli-core'; import type { CustomTheme } from '../ui/themes/theme.js'; import type { SessionRetentionSettings } from './settings.js'; @@ -1394,7 +1394,7 @@ const SETTINGS_SCHEMA = { label: 'Model', category: 'Experimental', requiresRestart: true, - default: DEFAULT_GEMINI_MODEL, + default: GEMINI_MODEL_ALIAS_AUTO, description: 'The model to use for the Codebase Investigator agent.', showInDialog: false, diff --git a/packages/core/src/agents/registry.test.ts b/packages/core/src/agents/registry.test.ts index 27e04ad229..ab0291b2cd 100644 --- a/packages/core/src/agents/registry.test.ts +++ b/packages/core/src/agents/registry.test.ts @@ -10,7 +10,13 @@ import { makeFakeConfig } from '../test-utils/config.js'; import type { AgentDefinition } from './types.js'; import type { Config } from '../config/config.js'; import { debugLogger } from '../utils/debugLogger.js'; -import { DEFAULT_GEMINI_MODEL } from '../config/models.js'; +import { + DEFAULT_GEMINI_FLASH_LITE_MODEL, + GEMINI_MODEL_ALIAS_AUTO, + PREVIEW_GEMINI_FLASH_MODEL, + PREVIEW_GEMINI_MODEL, + PREVIEW_GEMINI_MODEL_AUTO, +} from '../config/models.js'; // A test-only subclass to expose the protected `registerAgent` method. class TestableAgentRegistry extends AgentRegistry { @@ -74,12 +80,12 @@ describe('AgentRegistry', () => { ); }); - it('should use preview model for codebase investigator if main model is preview', async () => { + it('should use preview flash model for codebase investigator if main model is preview pro', async () => { const previewConfig = makeFakeConfig({ - model: 'gemini-3-pro-preview', + model: PREVIEW_GEMINI_MODEL, codebaseInvestigatorSettings: { enabled: true, - model: DEFAULT_GEMINI_MODEL, + model: GEMINI_MODEL_ALIAS_AUTO, }, }); const previewRegistry = new TestableAgentRegistry(previewConfig); @@ -90,7 +96,51 @@ describe('AgentRegistry', () => { 'codebase_investigator', ); expect(investigatorDef).toBeDefined(); - expect(investigatorDef?.modelConfig.model).toBe('gemini-3-pro-preview'); + expect(investigatorDef?.modelConfig.model).toBe( + PREVIEW_GEMINI_FLASH_MODEL, + ); + }); + + it('should use preview flash model for codebase investigator if main model is preview auto', async () => { + const previewConfig = makeFakeConfig({ + model: PREVIEW_GEMINI_MODEL_AUTO, + codebaseInvestigatorSettings: { + enabled: true, + model: GEMINI_MODEL_ALIAS_AUTO, + }, + }); + const previewRegistry = new TestableAgentRegistry(previewConfig); + + await previewRegistry.initialize(); + + const investigatorDef = previewRegistry.getDefinition( + 'codebase_investigator', + ); + expect(investigatorDef).toBeDefined(); + expect(investigatorDef?.modelConfig.model).toBe( + PREVIEW_GEMINI_FLASH_MODEL, + ); + }); + + it('should use the model from the investigator settings', async () => { + const previewConfig = makeFakeConfig({ + model: PREVIEW_GEMINI_MODEL, + codebaseInvestigatorSettings: { + enabled: true, + model: DEFAULT_GEMINI_FLASH_LITE_MODEL, + }, + }); + const previewRegistry = new TestableAgentRegistry(previewConfig); + + await previewRegistry.initialize(); + + const investigatorDef = previewRegistry.getDefinition( + 'codebase_investigator', + ); + expect(investigatorDef).toBeDefined(); + expect(investigatorDef?.modelConfig.model).toBe( + DEFAULT_GEMINI_FLASH_LITE_MODEL, + ); }); }); diff --git a/packages/core/src/agents/registry.ts b/packages/core/src/agents/registry.ts index ffaf41a0f2..86ec004e5a 100644 --- a/packages/core/src/agents/registry.ts +++ b/packages/core/src/agents/registry.ts @@ -11,11 +11,12 @@ import { type z } from 'zod'; import { debugLogger } from '../utils/debugLogger.js'; import { DEFAULT_GEMINI_MODEL, - DEFAULT_GEMINI_MODEL_AUTO, - PREVIEW_GEMINI_MODEL, - PREVIEW_GEMINI_MODEL_AUTO, + GEMINI_MODEL_ALIAS_AUTO, + PREVIEW_GEMINI_FLASH_MODEL, + isPreviewModel, } from '../config/models.js'; import type { ModelConfigAlias } from '../services/modelConfigService.js'; +import { coreEvents, CoreEvent } from '../utils/events.js'; /** * Returns the model config alias for a given agent definition. @@ -42,6 +43,10 @@ export class AgentRegistry { async initialize(): Promise { this.loadBuiltInAgents(); + coreEvents.on(CoreEvent.ModelChanged, () => { + this.loadBuiltInAgents(); + }); + if (this.config.getDebugMode()) { debugLogger.log( `[AgentRegistry] Initialized with ${this.agents.size} agents.`, @@ -54,23 +59,17 @@ export class AgentRegistry { // Only register the agent if it's enabled in the settings. if (investigatorSettings?.enabled) { - let model = - investigatorSettings.model ?? - CodebaseInvestigatorAgent.modelConfig.model; - - // If the user is using the preview model for the main agent, force the sub-agent to use it too - // if it's configured to use 'pro' or 'auto'. - if ( - this.config.getModel() === PREVIEW_GEMINI_MODEL || - this.config.getModel() === PREVIEW_GEMINI_MODEL_AUTO - ) { - if ( - model === PREVIEW_GEMINI_MODEL_AUTO || - model === DEFAULT_GEMINI_MODEL_AUTO || - model === DEFAULT_GEMINI_MODEL - ) { - model = PREVIEW_GEMINI_MODEL; - } + let model; + const settingsModel = investigatorSettings.model; + // Check if the user explicitly set a model in the settings. + if (settingsModel && settingsModel !== GEMINI_MODEL_ALIAS_AUTO) { + model = settingsModel; + } else { + // Use Preview Flash model if the main model is any of the preview models + // If the main model is not preview model, use default pro model. + model = isPreviewModel(this.config.getModel()) + ? PREVIEW_GEMINI_FLASH_MODEL + : DEFAULT_GEMINI_MODEL; } const agentDef = { diff --git a/packages/core/src/config/config.test.ts b/packages/core/src/config/config.test.ts index 3a4f2c2657..0b9395a916 100644 --- a/packages/core/src/config/config.test.ts +++ b/packages/core/src/config/config.test.ts @@ -1704,18 +1704,16 @@ describe('Availability Service Integration', () => { cwd: '.', }; - it('setActiveModel updates active model and emits event', async () => { + it('setActiveModel updates active model', async () => { const config = new Config(baseParams); const model1 = 'model1'; const model2 = 'model2'; config.setActiveModel(model1); expect(config.getActiveModel()).toBe(model1); - expect(mockCoreEvents.emitModelChanged).toHaveBeenCalledWith(model1); config.setActiveModel(model2); expect(config.getActiveModel()).toBe(model2); - expect(mockCoreEvents.emitModelChanged).toHaveBeenCalledWith(model2); }); it('getActiveModel defaults to configured model if not set', () => { diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 81bdec649e..aa2d5d6925 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -843,7 +843,6 @@ export class Config { setActiveModel(model: string): void { if (this._activeModel !== model) { this._activeModel = model; - coreEvents.emitModelChanged(model); } } diff --git a/packages/core/src/config/models.ts b/packages/core/src/config/models.ts index de0ff76f3e..ccaa46f059 100644 --- a/packages/core/src/config/models.ts +++ b/packages/core/src/config/models.ts @@ -12,6 +12,7 @@ export const DEFAULT_GEMINI_FLASH_LITE_MODEL = 'gemini-2.5-flash-lite'; export const VALID_GEMINI_MODELS = new Set([ PREVIEW_GEMINI_MODEL, + PREVIEW_GEMINI_FLASH_MODEL, DEFAULT_GEMINI_MODEL, DEFAULT_GEMINI_FLASH_MODEL, DEFAULT_GEMINI_FLASH_LITE_MODEL, @@ -21,6 +22,7 @@ export const PREVIEW_GEMINI_MODEL_AUTO = 'auto-gemini-3'; export const DEFAULT_GEMINI_MODEL_AUTO = 'auto-gemini-2.5'; // Model aliases for user convenience. +export const GEMINI_MODEL_ALIAS_AUTO = 'auto'; export const GEMINI_MODEL_ALIAS_PRO = 'pro'; export const GEMINI_MODEL_ALIAS_FLASH = 'flash'; export const GEMINI_MODEL_ALIAS_FLASH_LITE = 'flash-lite'; @@ -124,6 +126,20 @@ export function getDisplayString(model: string) { } } +/** + * Checks if the model is a preview model. + * + * @param model The model name to check. + * @returns True if the model is a preview model. + */ +export function isPreviewModel(model: string): boolean { + return ( + model === PREVIEW_GEMINI_MODEL || + model === PREVIEW_GEMINI_FLASH_MODEL || + model === PREVIEW_GEMINI_MODEL_AUTO + ); +} + /** * Checks if the model is a Gemini 2.x model. * diff --git a/schemas/settings.schema.json b/schemas/settings.schema.json index ec161aadfd..bf4def096c 100644 --- a/schemas/settings.schema.json +++ b/schemas/settings.schema.json @@ -1361,8 +1361,8 @@ "model": { "title": "Model", "description": "The model to use for the Codebase Investigator agent.", - "markdownDescription": "The model to use for the Codebase Investigator agent.\n\n- Category: `Experimental`\n- Requires restart: `yes`\n- Default: `gemini-2.5-pro`", - "default": "gemini-2.5-pro", + "markdownDescription": "The model to use for the Codebase Investigator agent.\n\n- Category: `Experimental`\n- Requires restart: `yes`\n- Default: `auto`", + "default": "auto", "type": "string" } },