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

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

@@ -9,6 +9,7 @@ import { Box, Text } from 'ink';
 import { theme } from '../semantic-colors.js';
 import { GIT_COMMIT_INFO } from '../../generated/git-commit.js';
 import { useSettings } from '../contexts/SettingsContext.js';
+import { getDisplayString } from '@google/gemini-cli-core';
 
 interface AboutBoxProps {
   cliVersion: string;
@@ -79,7 +80,9 @@ export const AboutBox: React.FC<AboutBoxProps> = ({
           </Text>
         </Box>
         <Box>
-          <Text color={theme.text.primary}>{modelVersion}</Text>
+          <Text color={theme.text.primary}>
+            {getDisplayString(modelVersion)}
+          </Text>
         </Box>
       </Box>
       <Box flexDirection="row">

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

@@ -9,11 +9,17 @@ import { act } from 'react';
 import { ModelDialog } from './ModelDialog.js';
 import { renderWithProviders } from '../../test-utils/render.js';
 import { waitFor } from '../../test-utils/async.js';
+import { createMockSettings } from '../../test-utils/settings.js';
 import {
   DEFAULT_GEMINI_MODEL,
   DEFAULT_GEMINI_MODEL_AUTO,
   DEFAULT_GEMINI_FLASH_MODEL,
   DEFAULT_GEMINI_FLASH_LITE_MODEL,
+  PREVIEW_GEMINI_MODEL,
+  PREVIEW_GEMINI_3_1_MODEL,
+  PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL,
+  PREVIEW_GEMINI_FLASH_MODEL,
+  AuthType,
 } from '@google/gemini-cli-core';
 import type { Config, ModelSlashCommandEvent } from '@google/gemini-cli-core';
 
@@ -42,12 +48,14 @@ describe('<ModelDialog />', () => {
   const mockGetModel = vi.fn();
   const mockOnClose = vi.fn();
   const mockGetHasAccessToPreviewModel = vi.fn();
+  const mockGetGemini31LaunchedSync = vi.fn();
 
   interface MockConfig extends Partial<Config> {
     setModel: (model: string, isTemporary?: boolean) => void;
     getModel: () => string;
     getHasAccessToPreviewModel: () => boolean;
     getIdeMode: () => boolean;
+    getGemini31LaunchedSync: () => boolean;
   }
 
   const mockConfig: MockConfig = {
@@ -55,12 +63,14 @@ describe('<ModelDialog />', () => {
     getModel: mockGetModel,
     getHasAccessToPreviewModel: mockGetHasAccessToPreviewModel,
     getIdeMode: () => false,
+    getGemini31LaunchedSync: mockGetGemini31LaunchedSync,
   };
 
   beforeEach(() => {
     vi.resetAllMocks();
     mockGetModel.mockReturnValue(DEFAULT_GEMINI_MODEL_AUTO);
     mockGetHasAccessToPreviewModel.mockReturnValue(false);
+    mockGetGemini31LaunchedSync.mockReturnValue(false);
 
     // Default implementation for getDisplayString
     mockGetDisplayString.mockImplementation((val: string) => {
@@ -70,9 +80,21 @@ describe('<ModelDialog />', () => {
     });
   });
 
-  const renderComponent = async (configValue = mockConfig as Config) => {
+  const renderComponent = async (
+    configValue = mockConfig as Config,
+    authType = AuthType.LOGIN_WITH_GOOGLE,
+  ) => {
+    const settings = createMockSettings({
+      security: {
+        auth: {
+          selectedType: authType,
+        },
+      },
+    });
+
     const result = renderWithProviders(<ModelDialog onClose={mockOnClose} />, {
       config: configValue,
+      settings,
     });
     await result.waitUntilReady();
     return result;
@@ -241,4 +263,98 @@ describe('<ModelDialog />', () => {
     });
     unmount();
   });
+
+  it('shows the preferred manual model in the main view option', async () => {
+    mockGetModel.mockReturnValue(DEFAULT_GEMINI_MODEL);
+    const { lastFrame, unmount } = await renderComponent();
+
+    expect(lastFrame()).toContain(`Manual (${DEFAULT_GEMINI_MODEL})`);
+    unmount();
+  });
+
+  describe('Preview Models', () => {
+    beforeEach(() => {
+      mockGetHasAccessToPreviewModel.mockReturnValue(true);
+    });
+
+    it('shows Auto (Preview) in main view when access is granted', async () => {
+      const { lastFrame, unmount } = await renderComponent();
+      expect(lastFrame()).toContain('Auto (Preview)');
+      unmount();
+    });
+
+    it('shows Gemini 3 models in manual view when Gemini 3.1 is NOT launched', async () => {
+      mockGetGemini31LaunchedSync.mockReturnValue(false);
+      const { lastFrame, stdin, waitUntilReady, unmount } =
+        await renderComponent();
+
+      // Go to manual view
+      await act(async () => {
+        stdin.write('\u001B[B'); // Manual
+      });
+      await waitUntilReady();
+      await act(async () => {
+        stdin.write('\r');
+      });
+      await waitUntilReady();
+
+      const output = lastFrame();
+      expect(output).toContain(PREVIEW_GEMINI_MODEL);
+      expect(output).toContain(PREVIEW_GEMINI_FLASH_MODEL);
+      unmount();
+    });
+
+    it('shows Gemini 3.1 models in manual view when Gemini 3.1 IS launched', async () => {
+      mockGetGemini31LaunchedSync.mockReturnValue(true);
+      const { lastFrame, stdin, waitUntilReady, unmount } =
+        await renderComponent(mockConfig as Config, AuthType.USE_VERTEX_AI);
+
+      // Go to manual view
+      await act(async () => {
+        stdin.write('\u001B[B'); // Manual
+      });
+      await waitUntilReady();
+      await act(async () => {
+        stdin.write('\r');
+      });
+      await waitUntilReady();
+
+      const output = lastFrame();
+      expect(output).toContain(PREVIEW_GEMINI_3_1_MODEL);
+      expect(output).toContain(PREVIEW_GEMINI_FLASH_MODEL);
+      unmount();
+    });
+
+    it('uses custom tools model when Gemini 3.1 IS launched and auth is Gemini API Key', async () => {
+      mockGetGemini31LaunchedSync.mockReturnValue(true);
+      const { stdin, waitUntilReady, unmount } = await renderComponent(
+        mockConfig as Config,
+        AuthType.USE_GEMINI,
+      );
+
+      // Go to manual view
+      await act(async () => {
+        stdin.write('\u001B[B'); // Manual
+      });
+      await waitUntilReady();
+      await act(async () => {
+        stdin.write('\r');
+      });
+      await waitUntilReady();
+
+      // Select Gemini 3.1 (first item in preview section)
+      await act(async () => {
+        stdin.write('\r');
+      });
+      await waitUntilReady();
+
+      await waitFor(() => {
+        expect(mockSetModel).toHaveBeenCalledWith(
+          PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL,
+          true,
+        );
+      });
+      unmount();
+    });
+  });
 });

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,
@@ -18,11 +19,14 @@ import {
   ModelSlashCommandEvent,
   logModelSlashCommand,
   getDisplayString,
+  AuthType,
+  PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL,
 } from '@google/gemini-cli-core';
 import { useKeypress } from '../hooks/useKeypress.js';
 import { theme } from '../semantic-colors.js';
 import { DescriptiveRadioButtonSelect } from './shared/DescriptiveRadioButtonSelect.js';
 import { ConfigContext } from '../contexts/ConfigContext.js';
+import { useSettings } from '../contexts/SettingsContext.js';
 
 interface ModelDialogProps {
   onClose: () => void;
@@ -30,6 +34,7 @@ interface ModelDialogProps {
 
 export function ModelDialog({ onClose }: ModelDialogProps): React.JSX.Element {
   const config = useContext(ConfigContext);
+  const settings = useSettings();
   const [view, setView] = useState<'main' | 'manual'>('main');
   const [persistMode, setPersistMode] = useState(false);
 
@@ -37,6 +42,10 @@ 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 selectedAuthType = settings.merged.security.auth.selectedType;
+  const useCustomToolModel =
+    useGemini31 && selectedAuthType === AuthType.USE_GEMINI;
 
   const manualModelSelected = useMemo(() => {
     const manualModels = [
@@ -44,6 +53,8 @@ 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_3_1_CUSTOM_TOOLS_MODEL,
       PREVIEW_GEMINI_FLASH_MODEL,
     ];
     if (manualModels.includes(preferredModel)) {
@@ -94,13 +105,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 = [
@@ -122,11 +134,19 @@ export function ModelDialog({ onClose }: ModelDialogProps): React.JSX.Element {
     ];
 
     if (shouldShowPreviewModels) {
+      const previewProModel = useGemini31
+        ? PREVIEW_GEMINI_3_1_MODEL
+        : PREVIEW_GEMINI_MODEL;
+
+      const previewProValue = useCustomToolModel
+        ? PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL
+        : previewProModel;
+
       list.unshift(
         {
-          value: PREVIEW_GEMINI_MODEL,
-          title: PREVIEW_GEMINI_MODEL,
-          key: PREVIEW_GEMINI_MODEL,
+          value: previewProValue,
+          title: previewProModel,
+          key: previewProModel,
         },
         {
           value: PREVIEW_GEMINI_FLASH_MODEL,
@@ -136,7 +156,7 @@ export function ModelDialog({ onClose }: ModelDialogProps): React.JSX.Element {
       );
     }
     return list;
-  }, [shouldShowPreviewModels]);
+  }, [shouldShowPreviewModels, useGemini31, useCustomToolModel]);
 
   const options = view === 'main' ? mainOptions : manualOptions;
 

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

@@ -23,11 +23,13 @@ import {
 import { computeSessionStats } from '../utils/computeStats.js';
 import {
   type RetrieveUserQuotaResponse,
-  VALID_GEMINI_MODELS,
+  isActiveModel,
   getDisplayString,
   isAutoModel,
+  AuthType,
 } 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,9 +84,13 @@ const Section: React.FC<SectionProps> = ({ title, children }) => (
 const buildModelRows = (
   models: Record<string, ModelMetrics>,
   quotas?: RetrieveUserQuotaResponse,
+  useGemini3_1 = false,
+  useCustomToolModel = false,
 ) => {
   const getBaseModelName = (name: string) => name.replace('-001', '');
-  const usedModelNames = new Set(Object.keys(models).map(getBaseModelName));
+  const usedModelNames = new Set(
+    Object.keys(models).map(getBaseModelName).map(getDisplayString),
+  );
 
   // 1. Models with active usage
   const activeRows = Object.entries(models).map(([name, metrics]) => {
@@ -93,7 +99,7 @@ const buildModelRows = (
     const inputTokens = metrics.tokens.input;
     return {
       key: name,
-      modelName,
+      modelName: getDisplayString(modelName),
       requests: metrics.api.totalRequests,
       cachedTokens: cachedTokens.toLocaleString(),
       inputTokens: inputTokens.toLocaleString(),
@@ -109,12 +115,12 @@ const buildModelRows = (
       ?.filter(
         (b) =>
           b.modelId &&
-          VALID_GEMINI_MODELS.has(b.modelId) &&
-          !usedModelNames.has(b.modelId),
+          isActiveModel(b.modelId, useGemini3_1, useCustomToolModel) &&
+          !usedModelNames.has(getDisplayString(b.modelId)),
       )
       .map((bucket) => ({
         key: bucket.modelId!,
-        modelName: bucket.modelId!,
+        modelName: getDisplayString(bucket.modelId!),
         requests: '-',
         cachedTokens: '-',
         inputTokens: '-',
@@ -135,6 +141,8 @@ const ModelUsageTable: React.FC<{
   pooledRemaining?: number;
   pooledLimit?: number;
   pooledResetTime?: string;
+  useGemini3_1?: boolean;
+  useCustomToolModel?: boolean;
 }> = ({
   models,
   quotas,
@@ -144,8 +152,10 @@ const ModelUsageTable: React.FC<{
   pooledRemaining,
   pooledLimit,
   pooledResetTime,
+  useGemini3_1,
+  useCustomToolModel,
 }) => {
-  const rows = buildModelRows(models, quotas);
+  const rows = buildModelRows(models, quotas, useGemini3_1, useCustomToolModel);
 
   if (rows.length === 0) {
     return null;
@@ -403,7 +413,11 @@ 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 useCustomToolModel =
+    useGemini3_1 &&
+    config.getContentGeneratorConfig().authType === AuthType.USE_GEMINI;
   const pooledRemaining = quotaStats?.remaining;
   const pooledLimit = quotaStats?.limit;
   const pooledResetTime = quotaStats?.resetTime;
@@ -544,6 +558,8 @@ export const StatsDisplay: React.FC<StatsDisplayProps> = ({
         pooledRemaining={pooledRemaining}
         pooledLimit={pooledLimit}
         pooledResetTime={pooledResetTime}
+        useGemini3_1={useGemini3_1}
+        useCustomToolModel={useCustomToolModel}
       />
       {renderFooter()}
     </Box>

packages/cli/src/ui/components/messages/ModelMessage.tsx

@@ -7,6 +7,7 @@
 import type React from 'react';
 import { Text, Box } from 'ink';
 import { theme } from '../../semantic-colors.js';
+import { getDisplayString } from '@google/gemini-cli-core';
 
 interface ModelMessageProps {
   model: string;
@@ -15,7 +16,7 @@ interface ModelMessageProps {
 export const ModelMessage: React.FC<ModelMessageProps> = ({ model }) => (
   <Box marginLeft={2}>
     <Text color={theme.ui.comment} italic>
-      Responding with {model}
+      Responding with {getDisplayString(model)}
     </Text>
   </Box>
 );