Respect previewFeatures value from the remote flag if undefined (#15214)

Author: sehoon38

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

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 {