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

@@ -483,7 +483,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 ?? [],