feat(core,ui): support Gemini 3.1 Pro Preview and active model filtering (#125)

* feat(core,ui): support Gemini 3.1 Pro Preview and active model filtering

* fix(core,ui): use optional chaining for config methods to support mocks

* fix(core): clear stale authType in refreshAuth to avoid incorrect model resolution

* do not show gemini 3.1 model when users do not have access to gemini 3.1 in stats

packages/cli/src/ui/components/ModelDialog.tsx

@@ -9,6 +9,7 @@ import { useCallback, useContext, useMemo, useState } from 'react';
 import { Box, Text } from 'ink';
 import {
   PREVIEW_GEMINI_MODEL,
+  PREVIEW_GEMINI_3_1_MODEL,
   PREVIEW_GEMINI_FLASH_MODEL,
   PREVIEW_GEMINI_MODEL_AUTO,
   DEFAULT_GEMINI_MODEL,
@@ -37,6 +38,7 @@ export function ModelDialog({ onClose }: ModelDialogProps): React.JSX.Element {
   const preferredModel = config?.getModel() || DEFAULT_GEMINI_MODEL_AUTO;
 
   const shouldShowPreviewModels = config?.getHasAccessToPreviewModel();
+  const useGemini31 = config?.getGemini31LaunchedSync?.() ?? false;
 
   const manualModelSelected = useMemo(() => {
     const manualModels = [
@@ -44,6 +46,7 @@ export function ModelDialog({ onClose }: ModelDialogProps): React.JSX.Element {
       DEFAULT_GEMINI_FLASH_MODEL,
       DEFAULT_GEMINI_FLASH_LITE_MODEL,
       PREVIEW_GEMINI_MODEL,
+      PREVIEW_GEMINI_3_1_MODEL,
       PREVIEW_GEMINI_FLASH_MODEL,
     ];
     if (manualModels.includes(preferredModel)) {
@@ -94,13 +97,14 @@ export function ModelDialog({ onClose }: ModelDialogProps): React.JSX.Element {
       list.unshift({
         value: PREVIEW_GEMINI_MODEL_AUTO,
         title: getDisplayString(PREVIEW_GEMINI_MODEL_AUTO),
-        description:
-          'Let Gemini CLI decide the best model for the task: gemini-3-pro, gemini-3-flash',
+        description: useGemini31
+          ? 'Let Gemini CLI decide the best model for the task: gemini-3.1-pro, gemini-3-flash'
+          : 'Let Gemini CLI decide the best model for the task: gemini-3-pro, gemini-3-flash',
         key: PREVIEW_GEMINI_MODEL_AUTO,
       });
     }
     return list;
-  }, [shouldShowPreviewModels, manualModelSelected]);
+  }, [shouldShowPreviewModels, manualModelSelected, useGemini31]);
 
   const manualOptions = useMemo(() => {
     const list = [
@@ -124,9 +128,9 @@ export function ModelDialog({ onClose }: ModelDialogProps): React.JSX.Element {
     if (shouldShowPreviewModels) {
       list.unshift(
         {
-          value: PREVIEW_GEMINI_MODEL,
-          title: PREVIEW_GEMINI_MODEL,
-          key: PREVIEW_GEMINI_MODEL,
+          value: useGemini31 ? PREVIEW_GEMINI_3_1_MODEL : PREVIEW_GEMINI_MODEL,
+          title: useGemini31 ? PREVIEW_GEMINI_3_1_MODEL : PREVIEW_GEMINI_MODEL,
+          key: useGemini31 ? PREVIEW_GEMINI_3_1_MODEL : PREVIEW_GEMINI_MODEL,
         },
         {
           value: PREVIEW_GEMINI_FLASH_MODEL,
@@ -136,7 +140,7 @@ export function ModelDialog({ onClose }: ModelDialogProps): React.JSX.Element {
       );
     }
     return list;
-  }, [shouldShowPreviewModels]);
+  }, [shouldShowPreviewModels, useGemini31]);
 
   const options = view === 'main' ? mainOptions : manualOptions;
 

packages/cli/src/ui/components/StatsDisplay.tsx

@@ -23,11 +23,12 @@ import {
 import { computeSessionStats } from '../utils/computeStats.js';
 import {
   type RetrieveUserQuotaResponse,
-  VALID_GEMINI_MODELS,
+  isActiveModel,
   getDisplayString,
   isAutoModel,
 } from '@google/gemini-cli-core';
 import { useSettings } from '../contexts/SettingsContext.js';
+import { useConfig } from '../contexts/ConfigContext.js';
 import type { QuotaStats } from '../types.js';
 import { QuotaStatsInfo } from './QuotaStatsInfo.js';
 
@@ -82,6 +83,7 @@ const Section: React.FC<SectionProps> = ({ title, children }) => (
 const buildModelRows = (
   models: Record<string, ModelMetrics>,
   quotas?: RetrieveUserQuotaResponse,
+  useGemini3_1 = false,
 ) => {
   const getBaseModelName = (name: string) => name.replace('-001', '');
   const usedModelNames = new Set(Object.keys(models).map(getBaseModelName));
@@ -109,7 +111,7 @@ const buildModelRows = (
       ?.filter(
         (b) =>
           b.modelId &&
-          VALID_GEMINI_MODELS.has(b.modelId) &&
+          isActiveModel(b.modelId, useGemini3_1) &&
           !usedModelNames.has(b.modelId),
       )
       .map((bucket) => ({
@@ -135,6 +137,7 @@ const ModelUsageTable: React.FC<{
   pooledRemaining?: number;
   pooledLimit?: number;
   pooledResetTime?: string;
+  useGemini3_1?: boolean;
 }> = ({
   models,
   quotas,
@@ -144,8 +147,9 @@ const ModelUsageTable: React.FC<{
   pooledRemaining,
   pooledLimit,
   pooledResetTime,
+  useGemini3_1 = false,
 }) => {
-  const rows = buildModelRows(models, quotas);
+  const rows = buildModelRows(models, quotas, useGemini3_1);
 
   if (rows.length === 0) {
     return null;
@@ -401,6 +405,8 @@ export const StatsDisplay: React.FC<StatsDisplayProps> = ({
   const { models, tools, files } = metrics;
   const computed = computeSessionStats(metrics);
   const settings = useSettings();
+  const config = useConfig();
+  const useGemini3_1 = config.getGemini31LaunchedSync?.() ?? false;
 
   const pooledRemaining = quotaStats?.remaining;
   const pooledLimit = quotaStats?.limit;
@@ -535,6 +541,7 @@ export const StatsDisplay: React.FC<StatsDisplayProps> = ({
         pooledRemaining={pooledRemaining}
         pooledLimit={pooledLimit}
         pooledResetTime={pooledResetTime}
+        useGemini3_1={useGemini3_1}
       />
     </Box>
   );

packages/cli/src/ui/hooks/useQuotaAndFallback.ts

@@ -14,9 +14,8 @@ import {
   TerminalQuotaError,
   ModelNotFoundError,
   type UserTierId,
-  PREVIEW_GEMINI_MODEL,
-  DEFAULT_GEMINI_MODEL,
   VALID_GEMINI_MODELS,
+  isProModel,
 } from '@google/gemini-cli-core';
 import { useCallback, useEffect, useRef, useState } from 'react';
 import { type UseHistoryManagerReturn } from './useHistoryManager.js';
@@ -67,11 +66,9 @@ export function useQuotaAndFallback({
       let message: string;
       let isTerminalQuotaError = false;
       let isModelNotFoundError = false;
-      const usageLimitReachedModel =
-        failedModel === DEFAULT_GEMINI_MODEL ||
-        failedModel === PREVIEW_GEMINI_MODEL
-          ? 'all Pro models'
-          : failedModel;
+      const usageLimitReachedModel = isProModel(failedModel)
+        ? 'all Pro models'
+        : failedModel;
       if (error instanceof TerminalQuotaError) {
         isTerminalQuotaError = true;
         // Common part of the message for both tiers

packages/cli/src/zed-integration/zedIntegration.ts

@@ -488,7 +488,10 @@ export class Session {
       const functionCalls: FunctionCall[] = [];
 
       try {
-        const model = resolveModel(this.config.getModel());
+        const model = resolveModel(
+          this.config.getModel(),
+          (await this.config.getGemini31Launched?.()) ?? false,
+        );
         const responseStream = await chat.sendMessageStream(
           { model },
           nextMessage?.parts ?? [],

packages/core/src/availability/policyHelpers.ts

@@ -44,7 +44,10 @@ export function resolvePolicyChain(
   const configuredModel = config.getModel();
 
   let chain;
-  const resolvedModel = resolveModel(modelFromConfig);
+  const resolvedModel = resolveModel(
+    modelFromConfig,
+    config.getGemini31LaunchedSync?.() ?? false,
+  );
   const isAutoPreferred = preferredModel ? isAutoModel(preferredModel) : false;
   const isAutoConfigured = isAutoModel(configuredModel);
   const hasAccessToPreview = config.getHasAccessToPreviewModel?.() ?? true;

packages/core/src/code_assist/experiments/flagNames.ts

@@ -16,6 +16,7 @@ export const ExperimentFlags = {
   MASKING_PROTECTION_THRESHOLD: 45758817,
   MASKING_PRUNABLE_THRESHOLD: 45758818,
   MASKING_PROTECT_LATEST_TURN: 45758819,
+  GEMINI_3_1_PRO_LAUNCHED: 45760185,
 } as const;
 
 export type ExperimentFlagName =

packages/core/src/config/config.ts

@@ -1041,6 +1041,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,
@@ -1320,7 +1326,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;
   }
 
@@ -1329,7 +1338,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;
   }
 
@@ -1338,7 +1350,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;
   }
 
@@ -2175,6 +2190,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,6 +25,7 @@ import {
   PREVIEW_GEMINI_FLASH_MODEL,
   PREVIEW_GEMINI_MODEL_AUTO,
   DEFAULT_GEMINI_MODEL_AUTO,
+  isActiveModel,
 } from './models.js';
 
 describe('isCustomModel', () => {
@@ -240,3 +241,24 @@ describe('resolveClassifierModel', () => {
     ).toBe(PREVIEW_GEMINI_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 false for invalid models', () => {
+    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);
+  });
+});

packages/core/src/config/models.ts

@@ -5,6 +5,7 @@
  */
 
 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_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 +13,7 @@ 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_FLASH_MODEL,
   DEFAULT_GEMINI_MODEL,
   DEFAULT_GEMINI_FLASH_MODEL,
@@ -37,19 +39,24 @@ 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,
+): string {
   switch (requestedModel) {
+    case PREVIEW_GEMINI_MODEL:
     case PREVIEW_GEMINI_MODEL_AUTO: {
-      return PREVIEW_GEMINI_MODEL;
+      return useGemini3_1 ? PREVIEW_GEMINI_3_1_MODEL : 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;
+      return useGemini3_1 ? PREVIEW_GEMINI_3_1_MODEL : PREVIEW_GEMINI_MODEL;
     }
     case GEMINI_MODEL_ALIAS_FLASH: {
       return PREVIEW_GEMINI_FLASH_MODEL;
@@ -115,11 +122,26 @@ 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 === PREVIEW_GEMINI_MODEL ||
+    model === PREVIEW_GEMINI_3_1_MODEL ||
+    model === DEFAULT_GEMINI_MODEL
+  );
+}
+
 /**
  * Checks if the model is a Gemini 3 model.
  *
@@ -188,3 +210,24 @@ 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,
+): boolean {
+  if (!VALID_GEMINI_MODELS.has(model)) {
+    return false;
+  }
+  if (useGemini3_1) {
+    return model !== PREVIEW_GEMINI_MODEL;
+  } else {
+    return model !== PREVIEW_GEMINI_3_1_MODEL;
+  }
+}

packages/core/src/core/client.ts

@@ -542,7 +542,10 @@ export class GeminiClient {
 
     // Availability logic: The configured model is the source of truth,
     // including any permanent fallbacks (config.setModel) or manual overrides.
-    return resolveModel(this.config.getActiveModel());
+    return resolveModel(
+      this.config.getActiveModel(),
+      this.config.getGemini31LaunchedSync?.() ?? false,
+    );
   }
 
   private async *processTurn(

packages/core/src/core/contentGenerator.ts

@@ -146,7 +146,12 @@ export async function createContentGenerator(
       return new LoggingContentGenerator(fakeGenerator, gcConfig);
     }
     const version = await getVersion();
-    const model = resolveModel(gcConfig.getModel());
+    const model = resolveModel(
+      gcConfig.getModel(),
+      config.authType === AuthType.USE_GEMINI ||
+        config.authType === AuthType.USE_VERTEX_AI ||
+        ((await gcConfig.getGemini31Launched?.()) ?? false),
+    );
     const customHeadersEnv =
       process.env['GEMINI_CLI_CUSTOM_HEADERS'] || undefined;
     const userAgent = `GeminiCLI/${version}/${model} (${process.platform}; ${process.arch})`;

packages/core/src/core/geminiChat.ts

@@ -496,13 +496,14 @@ export class GeminiChat {
     const initialActiveModel = this.config.getActiveModel();
 
     const apiCall = async () => {
+      const useGemini3_1 = (await this.config.getGemini31Launched?.()) ?? false;
       // Default to the last used model (which respects arguments/availability selection)
-      let modelToUse = resolveModel(lastModelToUse);
+      let modelToUse = resolveModel(lastModelToUse, useGemini3_1);
 
       // If the active model has changed (e.g. due to a fallback updating the config),
       // we switch to the new active model.
       if (this.config.getActiveModel() !== initialActiveModel) {
-        modelToUse = resolveModel(this.config.getActiveModel());
+        modelToUse = resolveModel(this.config.getActiveModel(), useGemini3_1);
       }
 
       if (modelToUse !== lastModelToUse) {

packages/core/src/prompts/promptProvider.ts

@@ -58,7 +58,10 @@ export class PromptProvider {
     const enabledToolNames = new Set(toolNames);
     const approvedPlanPath = config.getApprovedPlanPath();
 
-    const desiredModel = resolveModel(config.getActiveModel());
+    const desiredModel = resolveModel(
+      config.getActiveModel(),
+      config.getGemini31LaunchedSync?.() ?? false,
+    );
     const isModernModel = supportsModernFeatures(desiredModel);
     const activeSnippets = isModernModel ? snippets : legacySnippets;
     const contextFilenames = getAllGeminiMdFilenames();
@@ -231,7 +234,10 @@ export class PromptProvider {
   }
 
   getCompressionPrompt(config: Config): string {
-    const desiredModel = resolveModel(config.getActiveModel());
+    const desiredModel = resolveModel(
+      config.getActiveModel(),
+      config.getGemini31LaunchedSync?.() ?? false,
+    );
     const isModernModel = supportsModernFeatures(desiredModel);
     const activeSnippets = isModernModel ? snippets : legacySnippets;
     return activeSnippets.getCompressionPrompt();

packages/core/src/routing/strategies/defaultStrategy.ts

@@ -21,7 +21,10 @@ export class DefaultStrategy implements TerminalStrategy {
     config: Config,
     _baseLlmClient: BaseLlmClient,
   ): Promise<RoutingDecision> {
-    const defaultModel = resolveModel(config.getModel());
+    const defaultModel = resolveModel(
+      config.getModel(),
+      config.getGemini31LaunchedSync?.() ?? false,
+    );
     return {
       model: defaultModel,
       metadata: {

packages/core/src/routing/strategies/fallbackStrategy.ts

@@ -23,7 +23,10 @@ export class FallbackStrategy implements RoutingStrategy {
     _baseLlmClient: BaseLlmClient,
   ): Promise<RoutingDecision | null> {
     const requestedModel = context.requestedModel ?? config.getModel();
-    const resolvedModel = resolveModel(requestedModel);
+    const resolvedModel = resolveModel(
+      requestedModel,
+      config.getGemini31LaunchedSync?.() ?? false,
+    );
     const service = config.getModelAvailabilityService();
     const snapshot = service.snapshot(resolvedModel);
 

packages/core/src/routing/strategies/overrideStrategy.ts

@@ -33,7 +33,10 @@ export class OverrideStrategy implements RoutingStrategy {
 
     // Return the overridden model name.
     return {
-      model: resolveModel(overrideModel),
+      model: resolveModel(
+        overrideModel,
+        config.getGemini31LaunchedSync?.() ?? false,
+      ),
       metadata: {
         source: this.name,
         latencyMs: 0,

packages/core/src/services/chatCompressionService.ts

@@ -29,6 +29,7 @@ import {
   DEFAULT_GEMINI_MODEL,
   PREVIEW_GEMINI_MODEL,
   PREVIEW_GEMINI_FLASH_MODEL,
+  PREVIEW_GEMINI_3_1_MODEL,
 } from '../config/models.js';
 import { PreCompressTrigger } from '../hooks/types.js';
 import { LlmRole } from '../telemetry/types.js';
@@ -101,6 +102,7 @@ export function findCompressSplitPoint(
 export function modelStringToModelConfigAlias(model: string): string {
   switch (model) {
     case PREVIEW_GEMINI_MODEL:
+    case PREVIEW_GEMINI_3_1_MODEL:
       return 'chat-compression-3-pro';
     case PREVIEW_GEMINI_FLASH_MODEL:
       return 'chat-compression-3-flash';