diff --git a/packages/core/src/config/config.test.ts b/packages/core/src/config/config.test.ts index 5fd89916c56..20212e46bb5 100644 --- a/packages/core/src/config/config.test.ts +++ b/packages/core/src/config/config.test.ts @@ -38,6 +38,7 @@ import { DEFAULT_GEMINI_MODEL, DEFAULT_GEMINI_MODEL_AUTO, PREVIEW_GEMINI_MODEL, + PREVIEW_GEMINI_MODEL_AUTO, } from './models.js'; vi.mock('fs', async (importOriginal) => { @@ -1862,6 +1863,15 @@ describe('Config Quota & Preview Model Access', () => { expect(config.getModel()).toBe(nonPreviewModel); }); + it('should switch to preview auto model if enabling preview features while using default auto model', () => { + config.setPreviewFeatures(false); + config.setModel(DEFAULT_GEMINI_MODEL_AUTO); + + config.setPreviewFeatures(true); + + expect(config.getModel()).toBe(PREVIEW_GEMINI_MODEL_AUTO); + }); + it('should NOT reset model if enabling preview features', () => { config.setPreviewFeatures(false); config.setModel(PREVIEW_GEMINI_MODEL); // Just pretending it was set somehow diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index a0ad61c41ff..711c11edb93 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -52,6 +52,7 @@ import { DEFAULT_THINKING_MODE, isPreviewModel, PREVIEW_GEMINI_MODEL, + PREVIEW_GEMINI_MODEL_AUTO, } from './models.js'; import { shouldAttemptBrowserLaunch } from '../utils/browser.js'; import type { MCPOAuthConfig } from '../mcp/oauth-provider.js'; @@ -738,8 +739,6 @@ export class Config { // Initialize BaseLlmClient now that the ContentGenerator is available this.baseLlmClient = new BaseLlmClient(this.contentGenerator, this); - const previewFeatures = this.getPreviewFeatures(); - const codeAssistServer = getCodeAssistServer(this); if (codeAssistServer) { if (codeAssistServer.projectId) { @@ -751,7 +750,7 @@ export class Config { this.setExperiments(experiments); // If preview features have not been set and the user authenticated through Google, we enable preview based on remote config only if it's true - if (previewFeatures === undefined) { + if (this.getPreviewFeatures() === undefined) { const remotePreviewFeatures = experiments.flags[ExperimentFlags.ENABLE_PREVIEW]?.boolValue; if (remotePreviewFeatures === true) { @@ -975,14 +974,22 @@ export class Config { } setPreviewFeatures(previewFeatures: boolean) { - // If it's using a preview model and it's turning off previewFeatures, - // switch the model to the default auto mode. - if (this.previewFeatures && !previewFeatures) { - if (isPreviewModel(this.getModel())) { - this.setModel(DEFAULT_GEMINI_MODEL_AUTO); - } + // No change in state, no action needed + if (this.previewFeatures === previewFeatures) { + return; } this.previewFeatures = previewFeatures; + const currentModel = this.getModel(); + + // Case 1: Disabling preview features while on a preview model + if (!previewFeatures && isPreviewModel(currentModel)) { + this.setModel(DEFAULT_GEMINI_MODEL_AUTO); + } + + // Case 2: Enabling preview features while on the default auto model + else if (previewFeatures && currentModel === DEFAULT_GEMINI_MODEL_AUTO) { + this.setModel(PREVIEW_GEMINI_MODEL_AUTO); + } } getHasAccessToPreviewModel(): boolean {