feat(core): add policy chain support for Gemini 3.1 (#19991)

packages/core/src/availability/policyCatalog.test.ts

@@ -12,6 +12,8 @@ import {
 } from './policyCatalog.js';
 import {
   DEFAULT_GEMINI_MODEL,
+  PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL,
+  PREVIEW_GEMINI_3_1_MODEL,
   PREVIEW_GEMINI_MODEL,
 } from '../config/models.js';
 
@@ -22,6 +24,27 @@ describe('policyCatalog', () => {
     expect(chain).toHaveLength(2);
   });
 
+  it('returns Gemini 3.1 chain when useGemini31 is true', () => {
+    const chain = getModelPolicyChain({
+      previewEnabled: true,
+      useGemini31: true,
+    });
+    expect(chain[0]?.model).toBe(PREVIEW_GEMINI_3_1_MODEL);
+    expect(chain).toHaveLength(2);
+    expect(chain[1]?.model).toBe('gemini-3-flash-preview');
+  });
+
+  it('returns Gemini 3.1 Custom Tools chain when useGemini31 and useCustomToolModel are true', () => {
+    const chain = getModelPolicyChain({
+      previewEnabled: true,
+      useGemini31: true,
+      useCustomToolModel: true,
+    });
+    expect(chain[0]?.model).toBe(PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL);
+    expect(chain).toHaveLength(2);
+    expect(chain[1]?.model).toBe('gemini-3-flash-preview');
+  });
+
   it('returns default chain when preview disabled', () => {
     const chain = getModelPolicyChain({ previewEnabled: false });
     expect(chain[0]?.model).toBe(DEFAULT_GEMINI_MODEL);

packages/core/src/availability/policyCatalog.ts

@@ -16,6 +16,7 @@ import {
   DEFAULT_GEMINI_MODEL,
   PREVIEW_GEMINI_FLASH_MODEL,
   PREVIEW_GEMINI_MODEL,
+  resolveModel,
 } from '../config/models.js';
 import type { UserTierId } from '../code_assist/types.js';
 
@@ -28,6 +29,8 @@ type PolicyConfig = Omit<ModelPolicy, 'actions' | 'stateTransitions'> & {
 export interface ModelPolicyOptions {
   previewEnabled: boolean;
   userTier?: UserTierId;
+  useGemini31?: boolean;
+  useCustomToolModel?: boolean;
 }
 
 const DEFAULT_ACTIONS: ModelPolicyActionMap = {
@@ -56,11 +59,6 @@ const DEFAULT_CHAIN: ModelPolicyChain = [
   definePolicy({ model: DEFAULT_GEMINI_FLASH_MODEL, isLastResort: true }),
 ];
 
-const PREVIEW_CHAIN: ModelPolicyChain = [
-  definePolicy({ model: PREVIEW_GEMINI_MODEL }),
-  definePolicy({ model: PREVIEW_GEMINI_FLASH_MODEL, isLastResort: true }),
-];
-
 const FLASH_LITE_CHAIN: ModelPolicyChain = [
   definePolicy({
     model: DEFAULT_GEMINI_FLASH_LITE_MODEL,
@@ -84,7 +82,15 @@ export function getModelPolicyChain(
   options: ModelPolicyOptions,
 ): ModelPolicyChain {
   if (options.previewEnabled) {
-    return cloneChain(PREVIEW_CHAIN);
+    const previewModel = resolveModel(
+      PREVIEW_GEMINI_MODEL,
+      options.useGemini31,
+      options.useCustomToolModel,
+    );
+    return [
+      definePolicy({ model: previewModel }),
+      definePolicy({ model: PREVIEW_GEMINI_FLASH_MODEL, isLastResort: true }),
+    ];
   }
 
   return cloneChain(DEFAULT_CHAIN);

packages/core/src/availability/policyHelpers.test.ts

@@ -15,12 +15,17 @@ import type { Config } from '../config/config.js';
 import {
   DEFAULT_GEMINI_FLASH_LITE_MODEL,
   DEFAULT_GEMINI_MODEL_AUTO,
+  PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL,
+  PREVIEW_GEMINI_3_1_MODEL,
 } from '../config/models.js';
+import { AuthType } from '../core/contentGenerator.js';
 
 const createMockConfig = (overrides: Partial<Config> = {}): Config =>
   ({
     getUserTier: () => undefined,
     getModel: () => 'gemini-2.5-pro',
+    getGemini31LaunchedSync: () => false,
+    getContentGeneratorConfig: () => ({ authType: undefined }),
     ...overrides,
   }) as unknown as Config;
 
@@ -128,6 +133,27 @@ describe('policyHelpers', () => {
       expect(chain[0]?.model).toBe('gemini-2.5-pro');
       expect(chain[1]?.model).toBe('gemini-2.5-flash');
     });
+
+    it('returns Gemini 3.1 Pro chain when launched and auto-gemini-3 requested', () => {
+      const config = createMockConfig({
+        getModel: () => 'auto-gemini-3',
+        getGemini31LaunchedSync: () => true,
+      });
+      const chain = resolvePolicyChain(config);
+      expect(chain[0]?.model).toBe(PREVIEW_GEMINI_3_1_MODEL);
+      expect(chain[1]?.model).toBe('gemini-3-flash-preview');
+    });
+
+    it('returns Gemini 3.1 Pro Custom Tools chain when launched, auth is Gemini, and auto-gemini-3 requested', () => {
+      const config = createMockConfig({
+        getModel: () => 'auto-gemini-3',
+        getGemini31LaunchedSync: () => true,
+        getContentGeneratorConfig: () => ({ authType: AuthType.USE_GEMINI }),
+      });
+      const chain = resolvePolicyChain(config);
+      expect(chain[0]?.model).toBe(PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL);
+      expect(chain[1]?.model).toBe('gemini-3-flash-preview');
+    });
   });
 
   describe('buildFallbackPolicyContext', () => {

packages/core/src/availability/policyHelpers.ts

@@ -6,6 +6,7 @@
 
 import type { GenerateContentConfig } from '@google/genai';
 import type { Config } from '../config/config.js';
+import { AuthType } from '../core/contentGenerator.js';
 import type {
   FailureKind,
   FallbackAction,
@@ -44,9 +45,15 @@ export function resolvePolicyChain(
   const configuredModel = config.getModel();
 
   let chain;
+  const useGemini31 = config.getGemini31LaunchedSync?.() ?? false;
+  const useCustomToolModel =
+    useGemini31 &&
+    config.getContentGeneratorConfig?.()?.authType === AuthType.USE_GEMINI;
+
   const resolvedModel = resolveModel(
     modelFromConfig,
-    config.getGemini31LaunchedSync?.() ?? false,
+    useGemini31,
+    useCustomToolModel,
   );
   const isAutoPreferred = preferredModel ? isAutoModel(preferredModel) : false;
   const isAutoConfigured = isAutoModel(configuredModel);
@@ -67,6 +74,8 @@ export function resolvePolicyChain(
       chain = getModelPolicyChain({
         previewEnabled,
         userTier: config.getUserTier(),
+        useGemini31,
+        useCustomToolModel,
       });
     } else {
       // User requested Gemini 3 but has no access. Proactively downgrade
@@ -74,6 +83,8 @@ export function resolvePolicyChain(
       return getModelPolicyChain({
         previewEnabled: false,
         userTier: config.getUserTier(),
+        useGemini31,
+        useCustomToolModel,
       });
     }
   } else {