From aa9163da60e69f6077c6e9bc2d7142c0bcd7abb4 Mon Sep 17 00:00:00 2001 From: Sehoon Shon Date: Mon, 23 Feb 2026 10:13:48 -0500 Subject: [PATCH] feat(core): add policy chain support for Gemini 3.1 (#19991) --- .../src/availability/policyCatalog.test.ts | 23 ++++++++++++++++ .../core/src/availability/policyCatalog.ts | 18 ++++++++----- .../src/availability/policyHelpers.test.ts | 26 +++++++++++++++++++ .../core/src/availability/policyHelpers.ts | 13 +++++++++- 4 files changed, 73 insertions(+), 7 deletions(-) diff --git a/packages/core/src/availability/policyCatalog.test.ts b/packages/core/src/availability/policyCatalog.test.ts index 8edb2c6ad8..0133308688 100644 --- a/packages/core/src/availability/policyCatalog.test.ts +++ b/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); diff --git a/packages/core/src/availability/policyCatalog.ts b/packages/core/src/availability/policyCatalog.ts index 48713621cf..39dea34a2f 100644 --- a/packages/core/src/availability/policyCatalog.ts +++ b/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 & { 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); diff --git a/packages/core/src/availability/policyHelpers.test.ts b/packages/core/src/availability/policyHelpers.test.ts index 298f17eb74..22b8a62700 100644 --- a/packages/core/src/availability/policyHelpers.test.ts +++ b/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 => ({ 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', () => { diff --git a/packages/core/src/availability/policyHelpers.ts b/packages/core/src/availability/policyHelpers.ts index 6cf22d6388..456c8a855f 100644 --- a/packages/core/src/availability/policyHelpers.ts +++ b/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 {