feat(models): support Gemini 3.1 Pro Preview and fixes (#19676)

packages/core/src/config/config.ts

@@ -1071,6 +1071,12 @@ export class Config {
     // Reset availability status when switching auth (e.g. from limited key to OAuth)
     this.modelAvailabilityService.reset();
 
+    // Clear stale authType to ensure getGemini31LaunchedSync doesn't return stale results
+    // during the transition.
+    if (this.contentGeneratorConfig) {
+      this.contentGeneratorConfig.authType = undefined;
+    }
+
     const newContentGeneratorConfig = await createContentGeneratorConfig(
       this,
       authMethod,
@@ -1350,7 +1356,10 @@ export class Config {
     if (pooled.remaining !== undefined) {
       return pooled.remaining;
     }
-    const primaryModel = resolveModel(this.getModel());
+    const primaryModel = resolveModel(
+      this.getModel(),
+      this.getGemini31LaunchedSync(),
+    );
     return this.modelQuotas.get(primaryModel)?.remaining;
   }
 
@@ -1359,7 +1368,10 @@ export class Config {
     if (pooled.limit !== undefined) {
       return pooled.limit;
     }
-    const primaryModel = resolveModel(this.getModel());
+    const primaryModel = resolveModel(
+      this.getModel(),
+      this.getGemini31LaunchedSync(),
+    );
     return this.modelQuotas.get(primaryModel)?.limit;
   }
 
@@ -1368,7 +1380,10 @@ export class Config {
     if (pooled.resetTime !== undefined) {
       return pooled.resetTime;
     }
-    const primaryModel = resolveModel(this.getModel());
+    const primaryModel = resolveModel(
+      this.getModel(),
+      this.getGemini31LaunchedSync(),
+    );
     return this.modelQuotas.get(primaryModel)?.resetTime;
   }
 
@@ -2253,6 +2268,36 @@ export class Config {
     );
   }
 
+  /**
+   * Returns whether Gemini 3.1 has been launched.
+   * This method is async and ensures that experiments are loaded before returning the result.
+   */
+  async getGemini31Launched(): Promise<boolean> {
+    await this.ensureExperimentsLoaded();
+    return this.getGemini31LaunchedSync();
+  }
+
+  /**
+   * Returns whether Gemini 3.1 has been launched.
+   *
+   * Note: This method should only be called after startup, once experiments have been loaded.
+   * If you need to call this during startup or from an async context, use
+   * getGemini31Launched instead.
+   */
+  getGemini31LaunchedSync(): boolean {
+    const authType = this.contentGeneratorConfig?.authType;
+    if (
+      authType === AuthType.USE_GEMINI ||
+      authType === AuthType.USE_VERTEX_AI
+    ) {
+      return true;
+    }
+    return (
+      this.experiments?.flags[ExperimentFlags.GEMINI_3_1_PRO_LAUNCHED]
+        ?.boolValue ?? false
+    );
+  }
+
   private async ensureExperimentsLoaded(): Promise<void> {
     if (!this.experimentsPromise) {
       return;

packages/core/src/config/models.test.ts

@@ -25,8 +25,41 @@ import {
   PREVIEW_GEMINI_FLASH_MODEL,
   PREVIEW_GEMINI_MODEL_AUTO,
   DEFAULT_GEMINI_MODEL_AUTO,
+  isActiveModel,
+  PREVIEW_GEMINI_3_1_MODEL,
+  PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL,
+  isPreviewModel,
+  isProModel,
 } from './models.js';
 
+describe('isPreviewModel', () => {
+  it('should return true for preview models', () => {
+    expect(isPreviewModel(PREVIEW_GEMINI_MODEL)).toBe(true);
+    expect(isPreviewModel(PREVIEW_GEMINI_3_1_MODEL)).toBe(true);
+    expect(isPreviewModel(PREVIEW_GEMINI_FLASH_MODEL)).toBe(true);
+    expect(isPreviewModel(PREVIEW_GEMINI_MODEL_AUTO)).toBe(true);
+  });
+
+  it('should return false for non-preview models', () => {
+    expect(isPreviewModel(DEFAULT_GEMINI_MODEL)).toBe(false);
+    expect(isPreviewModel('gemini-1.5-pro')).toBe(false);
+  });
+});
+
+describe('isProModel', () => {
+  it('should return true for models containing "pro"', () => {
+    expect(isProModel('gemini-3-pro-preview')).toBe(true);
+    expect(isProModel('gemini-2.5-pro')).toBe(true);
+    expect(isProModel('pro')).toBe(true);
+  });
+
+  it('should return false for models without "pro"', () => {
+    expect(isProModel('gemini-3-flash-preview')).toBe(false);
+    expect(isProModel('gemini-2.5-flash')).toBe(false);
+    expect(isProModel('auto')).toBe(false);
+  });
+});
+
 describe('isCustomModel', () => {
   it('should return true for models not starting with gemini-', () => {
     expect(isCustomModel('testing')).toBe(true);
@@ -115,6 +148,12 @@ describe('getDisplayString', () => {
     );
   });
 
+  it('should return PREVIEW_GEMINI_3_1_MODEL for PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL', () => {
+    expect(getDisplayString(PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL)).toBe(
+      PREVIEW_GEMINI_3_1_MODEL,
+    );
+  });
+
   it('should return the model name as is for other models', () => {
     expect(getDisplayString('custom-model')).toBe('custom-model');
     expect(getDisplayString(DEFAULT_GEMINI_FLASH_LITE_MODEL)).toBe(
@@ -146,6 +185,16 @@ describe('resolveModel', () => {
       expect(model).toBe(PREVIEW_GEMINI_MODEL);
     });
 
+    it('should return Gemini 3.1 Pro when auto-gemini-3 is requested and useGemini3_1 is true', () => {
+      const model = resolveModel(PREVIEW_GEMINI_MODEL_AUTO, true);
+      expect(model).toBe(PREVIEW_GEMINI_3_1_MODEL);
+    });
+
+    it('should return Gemini 3.1 Pro Custom Tools when auto-gemini-3 is requested, useGemini3_1 is true, and useCustomToolModel is true', () => {
+      const model = resolveModel(PREVIEW_GEMINI_MODEL_AUTO, true, true);
+      expect(model).toBe(PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL);
+    });
+
     it('should return the Default Pro model when auto-gemini-2.5 is requested', () => {
       const model = resolveModel(DEFAULT_GEMINI_MODEL_AUTO);
       expect(model).toBe(DEFAULT_GEMINI_MODEL);
@@ -239,4 +288,71 @@ describe('resolveClassifierModel', () => {
       resolveClassifierModel(PREVIEW_GEMINI_MODEL_AUTO, GEMINI_MODEL_ALIAS_PRO),
     ).toBe(PREVIEW_GEMINI_MODEL);
   });
+
+  it('should return Gemini 3.1 Pro when alias is pro and useGemini3_1 is true', () => {
+    expect(
+      resolveClassifierModel(
+        PREVIEW_GEMINI_MODEL_AUTO,
+        GEMINI_MODEL_ALIAS_PRO,
+        true,
+      ),
+    ).toBe(PREVIEW_GEMINI_3_1_MODEL);
+  });
+
+  it('should return Gemini 3.1 Pro Custom Tools when alias is pro, useGemini3_1 is true, and useCustomToolModel is true', () => {
+    expect(
+      resolveClassifierModel(
+        PREVIEW_GEMINI_MODEL_AUTO,
+        GEMINI_MODEL_ALIAS_PRO,
+        true,
+        true,
+      ),
+    ).toBe(PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL);
+  });
+});
+
+describe('isActiveModel', () => {
+  it('should return true for valid models when useGemini3_1 is false', () => {
+    expect(isActiveModel(DEFAULT_GEMINI_MODEL)).toBe(true);
+    expect(isActiveModel(PREVIEW_GEMINI_MODEL)).toBe(true);
+    expect(isActiveModel(DEFAULT_GEMINI_FLASH_MODEL)).toBe(true);
+  });
+
+  it('should return true for unknown models and aliases', () => {
+    expect(isActiveModel('invalid-model')).toBe(false);
+    expect(isActiveModel(GEMINI_MODEL_ALIAS_AUTO)).toBe(false);
+  });
+
+  it('should return false for PREVIEW_GEMINI_MODEL when useGemini3_1 is true', () => {
+    expect(isActiveModel(PREVIEW_GEMINI_MODEL, true)).toBe(false);
+  });
+
+  it('should return true for other valid models when useGemini3_1 is true', () => {
+    expect(isActiveModel(DEFAULT_GEMINI_MODEL, true)).toBe(true);
+  });
+
+  it('should correctly filter Gemini 3.1 models based on useCustomToolModel when useGemini3_1 is true', () => {
+    // When custom tools are preferred, standard 3.1 should be inactive
+    expect(isActiveModel(PREVIEW_GEMINI_3_1_MODEL, true, true)).toBe(false);
+    expect(
+      isActiveModel(PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL, true, true),
+    ).toBe(true);
+
+    // When custom tools are NOT preferred, custom tools 3.1 should be inactive
+    expect(isActiveModel(PREVIEW_GEMINI_3_1_MODEL, true, false)).toBe(true);
+    expect(
+      isActiveModel(PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL, true, false),
+    ).toBe(false);
+  });
+
+  it('should return false for both Gemini 3.1 models when useGemini3_1 is false', () => {
+    expect(isActiveModel(PREVIEW_GEMINI_3_1_MODEL, false, true)).toBe(false);
+    expect(isActiveModel(PREVIEW_GEMINI_3_1_MODEL, false, false)).toBe(false);
+    expect(
+      isActiveModel(PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL, false, true),
+    ).toBe(false);
+    expect(
+      isActiveModel(PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL, false, false),
+    ).toBe(false);
+  });
 });

packages/core/src/config/models.ts

@@ -5,6 +5,9 @@
  */
 
 export const PREVIEW_GEMINI_MODEL = 'gemini-3-pro-preview';
+export const PREVIEW_GEMINI_3_1_MODEL = 'gemini-3.1-pro-preview';
+export const PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL =
+  'gemini-3.1-pro-preview-customtools';
 export const PREVIEW_GEMINI_FLASH_MODEL = 'gemini-3-flash-preview';
 export const DEFAULT_GEMINI_MODEL = 'gemini-2.5-pro';
 export const DEFAULT_GEMINI_FLASH_MODEL = 'gemini-2.5-flash';
@@ -12,6 +15,8 @@ export const DEFAULT_GEMINI_FLASH_LITE_MODEL = 'gemini-2.5-flash-lite';
 
 export const VALID_GEMINI_MODELS = new Set([
   PREVIEW_GEMINI_MODEL,
+  PREVIEW_GEMINI_3_1_MODEL,
+  PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL,
   PREVIEW_GEMINI_FLASH_MODEL,
   DEFAULT_GEMINI_MODEL,
   DEFAULT_GEMINI_FLASH_MODEL,
@@ -37,20 +42,29 @@ export const DEFAULT_THINKING_MODE = 8192;
  * to a concrete model name.
  *
  * @param requestedModel The model alias or concrete model name requested by the user.
+ * @param useGemini3_1 Whether to use Gemini 3.1 Pro Preview for auto/pro aliases.
  * @returns The resolved concrete model name.
  */
-export function resolveModel(requestedModel: string): string {
+export function resolveModel(
+  requestedModel: string,
+  useGemini3_1: boolean = false,
+  useCustomToolModel: boolean = false,
+): string {
   switch (requestedModel) {
-    case PREVIEW_GEMINI_MODEL_AUTO: {
+    case PREVIEW_GEMINI_MODEL:
+    case PREVIEW_GEMINI_MODEL_AUTO:
+    case GEMINI_MODEL_ALIAS_AUTO:
+    case GEMINI_MODEL_ALIAS_PRO: {
+      if (useGemini3_1) {
+        return useCustomToolModel
+          ? PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL
+          : PREVIEW_GEMINI_3_1_MODEL;
+      }
       return PREVIEW_GEMINI_MODEL;
     }
     case DEFAULT_GEMINI_MODEL_AUTO: {
       return DEFAULT_GEMINI_MODEL;
     }
-    case GEMINI_MODEL_ALIAS_AUTO:
-    case GEMINI_MODEL_ALIAS_PRO: {
-      return PREVIEW_GEMINI_MODEL;
-    }
     case GEMINI_MODEL_ALIAS_FLASH: {
       return PREVIEW_GEMINI_FLASH_MODEL;
     }
@@ -73,6 +87,8 @@ export function resolveModel(requestedModel: string): string {
 export function resolveClassifierModel(
   requestedModel: string,
   modelAlias: string,
+  useGemini3_1: boolean = false,
+  useCustomToolModel: boolean = false,
 ): string {
   if (modelAlias === GEMINI_MODEL_ALIAS_FLASH) {
     if (
@@ -89,7 +105,7 @@ export function resolveClassifierModel(
     }
     return resolveModel(GEMINI_MODEL_ALIAS_FLASH);
   }
-  return resolveModel(requestedModel);
+  return resolveModel(requestedModel, useGemini3_1, useCustomToolModel);
 }
 export function getDisplayString(model: string) {
   switch (model) {
@@ -101,6 +117,8 @@ export function getDisplayString(model: string) {
       return PREVIEW_GEMINI_MODEL;
     case GEMINI_MODEL_ALIAS_FLASH:
       return PREVIEW_GEMINI_FLASH_MODEL;
+    case PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL:
+      return PREVIEW_GEMINI_3_1_MODEL;
     default:
       return model;
   }
@@ -115,11 +133,22 @@ export function getDisplayString(model: string) {
 export function isPreviewModel(model: string): boolean {
   return (
     model === PREVIEW_GEMINI_MODEL ||
+    model === PREVIEW_GEMINI_3_1_MODEL ||
     model === PREVIEW_GEMINI_FLASH_MODEL ||
     model === PREVIEW_GEMINI_MODEL_AUTO
   );
 }
 
+/**
+ * Checks if the model is a Pro model.
+ *
+ * @param model The model name to check.
+ * @returns True if the model is a Pro model.
+ */
+export function isProModel(model: string): boolean {
+  return model.toLowerCase().includes('pro');
+}
+
 /**
  * Checks if the model is a Gemini 3 model.
  *
@@ -188,3 +217,35 @@ export function isAutoModel(model: string): boolean {
 export function supportsMultimodalFunctionResponse(model: string): boolean {
   return model.startsWith('gemini-3-');
 }
+
+/**
+ * Checks if the given model is considered active based on the current configuration.
+ *
+ * @param model The model name to check.
+ * @param useGemini3_1 Whether Gemini 3.1 Pro Preview is enabled.
+ * @returns True if the model is active.
+ */
+export function isActiveModel(
+  model: string,
+  useGemini3_1: boolean = false,
+  useCustomToolModel: boolean = false,
+): boolean {
+  if (!VALID_GEMINI_MODELS.has(model)) {
+    return false;
+  }
+  if (useGemini3_1) {
+    if (model === PREVIEW_GEMINI_MODEL) {
+      return false;
+    }
+    if (useCustomToolModel) {
+      return model !== PREVIEW_GEMINI_3_1_MODEL;
+    } else {
+      return model !== PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL;
+    }
+  } else {
+    return (
+      model !== PREVIEW_GEMINI_3_1_MODEL &&
+      model !== PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL
+    );
+  }
+}