From aced40107103ca5f7d81c9265b60a3f0d2d31c83 Mon Sep 17 00:00:00 2001 From: Abhi <43648792+abhipatel12@users.noreply.github.com> Date: Tue, 9 Dec 2025 16:07:50 -0500 Subject: [PATCH] feat: ensure codebase investigator uses preview model when main agent does (#14412) --- docs/get-started/configuration.md | 2 +- packages/cli/src/config/settingsSchema.ts | 4 ++-- .../src/agents/codebase-investigator.test.ts | 4 ++-- .../core/src/agents/codebase-investigator.ts | 4 ++-- packages/core/src/agents/registry.test.ts | 19 +++++++++++++++ packages/core/src/agents/registry.ts | 24 ++++++++++++++++--- packages/core/src/config/config.test.ts | 5 ++++ packages/core/src/config/config.ts | 3 +-- schemas/settings.schema.json | 4 ++-- 9 files changed, 55 insertions(+), 14 deletions(-) diff --git a/docs/get-started/configuration.md b/docs/get-started/configuration.md index b2e52569e0..71167e818a 100644 --- a/docs/get-started/configuration.md +++ b/docs/get-started/configuration.md @@ -816,7 +816,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:** `"pro"` - **Requires restart:** Yes #### `hooks` diff --git a/packages/cli/src/config/settingsSchema.ts b/packages/cli/src/config/settingsSchema.ts index f103cf2db9..5efca6f20f 100644 --- a/packages/cli/src/config/settingsSchema.ts +++ b/packages/cli/src/config/settingsSchema.ts @@ -18,8 +18,8 @@ import type { import { DEFAULT_TRUNCATE_TOOL_OUTPUT_LINES, DEFAULT_TRUNCATE_TOOL_OUTPUT_THRESHOLD, - DEFAULT_GEMINI_MODEL, DEFAULT_MODEL_CONFIGS, + GEMINI_MODEL_ALIAS_PRO, } from '@google/gemini-cli-core'; import type { CustomTheme } from '../ui/themes/theme.js'; import type { SessionRetentionSettings } from './settings.js'; @@ -1384,7 +1384,7 @@ const SETTINGS_SCHEMA = { label: 'Model', category: 'Experimental', requiresRestart: true, - default: DEFAULT_GEMINI_MODEL, + default: GEMINI_MODEL_ALIAS_PRO, description: 'The model to use for the Codebase Investigator agent.', showInDialog: false, diff --git a/packages/core/src/agents/codebase-investigator.test.ts b/packages/core/src/agents/codebase-investigator.test.ts index 3d8453cb97..f5398fc32a 100644 --- a/packages/core/src/agents/codebase-investigator.test.ts +++ b/packages/core/src/agents/codebase-investigator.test.ts @@ -12,7 +12,7 @@ import { LS_TOOL_NAME, READ_FILE_TOOL_NAME, } from '../tools/tool-names.js'; -import { DEFAULT_GEMINI_MODEL } from '../config/models.js'; +import { GEMINI_MODEL_ALIAS_PRO } from '../config/models.js'; describe('CodebaseInvestigatorAgent', () => { it('should have the correct agent definition', () => { @@ -26,7 +26,7 @@ describe('CodebaseInvestigatorAgent', () => { ).toBe(true); expect(CodebaseInvestigatorAgent.outputConfig?.outputName).toBe('report'); expect(CodebaseInvestigatorAgent.modelConfig?.model).toBe( - DEFAULT_GEMINI_MODEL, + GEMINI_MODEL_ALIAS_PRO, ); expect(CodebaseInvestigatorAgent.toolConfig?.tools).toEqual([ LS_TOOL_NAME, diff --git a/packages/core/src/agents/codebase-investigator.ts b/packages/core/src/agents/codebase-investigator.ts index adda3e96a5..947fc1f9c4 100644 --- a/packages/core/src/agents/codebase-investigator.ts +++ b/packages/core/src/agents/codebase-investigator.ts @@ -11,7 +11,7 @@ import { LS_TOOL_NAME, READ_FILE_TOOL_NAME, } from '../tools/tool-names.js'; -import { DEFAULT_GEMINI_MODEL } from '../config/models.js'; +import { GEMINI_MODEL_ALIAS_PRO } from '../config/models.js'; import { z } from 'zod'; // Define a type that matches the outputConfig schema for type safety. @@ -69,7 +69,7 @@ export const CodebaseInvestigatorAgent: AgentDefinition< processOutput: (output) => JSON.stringify(output, null, 2), modelConfig: { - model: DEFAULT_GEMINI_MODEL, + model: GEMINI_MODEL_ALIAS_PRO, temp: 0.1, top_p: 0.95, thinkingBudget: -1, diff --git a/packages/core/src/agents/registry.test.ts b/packages/core/src/agents/registry.test.ts index 23af430c62..4df0b206a9 100644 --- a/packages/core/src/agents/registry.test.ts +++ b/packages/core/src/agents/registry.test.ts @@ -72,6 +72,25 @@ describe('AgentRegistry', () => { `[AgentRegistry] Initialized with ${agentCount} agents.`, ); }); + + it('should use preview model for codebase investigator if main model is preview', async () => { + const previewConfig = makeFakeConfig({ + model: 'gemini-3-pro-preview', + codebaseInvestigatorSettings: { + enabled: true, + model: 'pro', + }, + }); + const previewRegistry = new TestableAgentRegistry(previewConfig); + + await previewRegistry.initialize(); + + const investigatorDef = previewRegistry.getDefinition( + 'codebase_investigator', + ); + expect(investigatorDef).toBeDefined(); + expect(investigatorDef?.modelConfig.model).toBe('gemini-3-pro-preview'); + }); }); describe('registration logic', () => { diff --git a/packages/core/src/agents/registry.ts b/packages/core/src/agents/registry.ts index 1bee30ac29..c7d7c65fae 100644 --- a/packages/core/src/agents/registry.ts +++ b/packages/core/src/agents/registry.ts @@ -9,6 +9,11 @@ import type { AgentDefinition } from './types.js'; import { CodebaseInvestigatorAgent } from './codebase-investigator.js'; import { type z } from 'zod'; import { debugLogger } from '../utils/debugLogger.js'; +import { + DEFAULT_GEMINI_MODEL_AUTO, + GEMINI_MODEL_ALIAS_PRO, + PREVIEW_GEMINI_MODEL, +} from '../config/models.js'; import type { ModelConfigAlias } from '../services/modelConfigService.js'; /** @@ -48,13 +53,26 @@ 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) { + if ( + model === GEMINI_MODEL_ALIAS_PRO || + model === DEFAULT_GEMINI_MODEL_AUTO + ) { + model = PREVIEW_GEMINI_MODEL; + } + } + const agentDef = { ...CodebaseInvestigatorAgent, modelConfig: { ...CodebaseInvestigatorAgent.modelConfig, - model: - investigatorSettings.model ?? - CodebaseInvestigatorAgent.modelConfig.model, + model, thinkingBudget: investigatorSettings.thinkingBudget ?? CodebaseInvestigatorAgent.modelConfig.thinkingBudget, diff --git a/packages/core/src/config/config.test.ts b/packages/core/src/config/config.test.ts index c67f05f888..90247b4441 100644 --- a/packages/core/src/config/config.test.ts +++ b/packages/core/src/config/config.test.ts @@ -888,6 +888,11 @@ describe('Server Config (config.ts)', () => { expect(SubagentToolWrapperMock).not.toHaveBeenCalled(); }); + it('should not set default codebase investigator model in config (defaults in registry)', () => { + const config = new Config(baseParams); + expect(config.getCodebaseInvestigatorSettings()?.model).toBeUndefined(); + }); + describe('with minified tool class names', () => { beforeEach(() => { Object.defineProperty( diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index b9a78ca7b9..678af20cd0 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -48,7 +48,6 @@ import { tokenLimit } from '../core/tokenLimits.js'; import { DEFAULT_GEMINI_EMBEDDING_MODEL, DEFAULT_GEMINI_FLASH_MODEL, - DEFAULT_GEMINI_MODEL, DEFAULT_THINKING_MODE, } from './models.js'; import { shouldAttemptBrowserLaunch } from '../utils/browser.js'; @@ -573,7 +572,7 @@ export class Config { thinkingBudget: params.codebaseInvestigatorSettings?.thinkingBudget ?? DEFAULT_THINKING_MODE, - model: params.codebaseInvestigatorSettings?.model ?? DEFAULT_GEMINI_MODEL, + model: params.codebaseInvestigatorSettings?.model, }; this.continueOnFailedApiCall = params.continueOnFailedApiCall ?? true; this.enableShellOutputEfficiency = diff --git a/schemas/settings.schema.json b/schemas/settings.schema.json index 93c95bb659..9b38000e5d 100644 --- a/schemas/settings.schema.json +++ b/schemas/settings.schema.json @@ -1331,8 +1331,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: `pro`", + "default": "pro", "type": "string" } },