Check if a user has access to preview model (#87)

packages/cli/src/ui/commands/modelCommand.test.ts

@@ -4,10 +4,11 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { describe, it, expect, beforeEach } from 'vitest';
+import { describe, it, expect, beforeEach, vi } from 'vitest';
 import { modelCommand } from './modelCommand.js';
 import { type CommandContext } from './types.js';
 import { createMockCommandContext } from '../../test-utils/mockCommandContext.js';
+import type { Config } from '@google/gemini-cli-core';
 
 describe('modelCommand', () => {
   let mockContext: CommandContext;
@@ -29,6 +30,21 @@ describe('modelCommand', () => {
     });
   });
 
+  it('should call refreshUserQuota if config is available', async () => {
+    if (!modelCommand.action) {
+      throw new Error('The model command must have an action.');
+    }
+
+    const mockRefreshUserQuota = vi.fn();
+    mockContext.services.config = {
+      refreshUserQuota: mockRefreshUserQuota,
+    } as unknown as Config;
+
+    await modelCommand.action(mockContext, '');
+
+    expect(mockRefreshUserQuota).toHaveBeenCalled();
+  });
+
   it('should have the correct name and description', () => {
     expect(modelCommand.name).toBe('model');
     expect(modelCommand.description).toBe(

packages/cli/src/ui/commands/modelCommand.ts

@@ -4,15 +4,24 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { CommandKind, type SlashCommand } from './types.js';
+import {
+  type CommandContext,
+  CommandKind,
+  type SlashCommand,
+} from './types.js';
 
 export const modelCommand: SlashCommand = {
   name: 'model',
   description: 'Opens a dialog to configure the model',
   kind: CommandKind.BUILT_IN,
   autoExecute: true,
-  action: async () => ({
-    type: 'dialog',
-    dialog: 'model',
-  }),
+  action: async (context: CommandContext) => {
+    if (context.services.config) {
+      await context.services.config.refreshUserQuota();
+    }
+    return {
+      type: 'dialog',
+      dialog: 'model',
+    };
+  },
 };

packages/cli/src/ui/commands/statsCommand.test.ts

@@ -10,6 +10,7 @@ import { type CommandContext } from './types.js';
 import { createMockCommandContext } from '../../test-utils/mockCommandContext.js';
 import { MessageType } from '../types.js';
 import { formatDuration } from '../utils/formatters.js';
+import type { Config } from '@google/gemini-cli-core';
 
 describe('statsCommand', () => {
   let mockContext: CommandContext;
@@ -45,6 +46,26 @@ describe('statsCommand', () => {
     );
   });
 
+  it('should fetch and display quota if config is available', async () => {
+    if (!statsCommand.action) throw new Error('Command has no action');
+
+    const mockQuota = { buckets: [] };
+    const mockRefreshUserQuota = vi.fn().mockResolvedValue(mockQuota);
+    mockContext.services.config = {
+      refreshUserQuota: mockRefreshUserQuota,
+    } as unknown as Config;
+
+    await statsCommand.action(mockContext, '');
+
+    expect(mockRefreshUserQuota).toHaveBeenCalled();
+    expect(mockContext.ui.addItem).toHaveBeenCalledWith(
+      expect.objectContaining({
+        quotas: mockQuota,
+      }),
+      expect.any(Number),
+    );
+  });
+
   it('should display model stats when using the "model" subcommand', () => {
     const modelSubCommand = statsCommand.subCommands?.find(
       (sc) => sc.name === 'model',

packages/cli/src/ui/commands/statsCommand.ts

@@ -4,7 +4,6 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { CodeAssistServer, getCodeAssistServer } from '@google/gemini-cli-core';
 import type { HistoryItemStats } from '../types.js';
 import { MessageType } from '../types.js';
 import { formatDuration } from '../utils/formatters.js';
@@ -35,11 +34,8 @@ async function defaultSessionView(context: CommandContext) {
   };
 
   if (context.services.config) {
-    const server = getCodeAssistServer(context.services.config);
-    if (server instanceof CodeAssistServer && server.projectId) {
-      const quota = await server.retrieveUserQuota({
-        project: server.projectId,
-      });
+    const quota = await context.services.config.refreshUserQuota();
+    if (quota) {
       statsItem.quotas = quota;
     }
   }

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

@@ -15,6 +15,7 @@ import {
   DEFAULT_GEMINI_FLASH_MODEL,
   DEFAULT_GEMINI_FLASH_LITE_MODEL,
   PREVIEW_GEMINI_MODEL,
+  PREVIEW_GEMINI_MODEL_AUTO,
 } from '@google/gemini-cli-core';
 import type { Config, ModelSlashCommandEvent } from '@google/gemini-cli-core';
 
@@ -43,23 +44,27 @@ describe('<ModelDialog />', () => {
   const mockGetModel = vi.fn();
   const mockGetPreviewFeatures = vi.fn();
   const mockOnClose = vi.fn();
+  const mockGetHasAccessToPreviewModel = vi.fn();
 
   interface MockConfig extends Partial<Config> {
     setModel: (model: string) => void;
     getModel: () => string;
     getPreviewFeatures: () => boolean;
+    getHasAccessToPreviewModel: () => boolean;
   }
 
   const mockConfig: MockConfig = {
     setModel: mockSetModel,
     getModel: mockGetModel,
     getPreviewFeatures: mockGetPreviewFeatures,
+    getHasAccessToPreviewModel: mockGetHasAccessToPreviewModel,
   };
 
   beforeEach(() => {
     vi.resetAllMocks();
     mockGetModel.mockReturnValue(DEFAULT_GEMINI_MODEL_AUTO);
     mockGetPreviewFeatures.mockReturnValue(false);
+    mockGetHasAccessToPreviewModel.mockReturnValue(false);
 
     // Default implementation for getDisplayString
     mockGetDisplayString.mockImplementation((val: string) => {
@@ -90,6 +95,7 @@ describe('<ModelDialog />', () => {
 
   it('renders "main" view with preview options when preview features are enabled', () => {
     mockGetPreviewFeatures.mockReturnValue(true);
+    mockGetHasAccessToPreviewModel.mockReturnValue(true); // Must have access
     const { lastFrame } = renderComponent();
     expect(lastFrame()).toContain('Auto (Preview)');
   });
@@ -114,13 +120,15 @@ describe('<ModelDialog />', () => {
 
   it('renders "manual" view with preview options when preview features are enabled', async () => {
     mockGetPreviewFeatures.mockReturnValue(true);
-    mockGetModel.mockReturnValue(DEFAULT_GEMINI_MODEL_AUTO);
+    mockGetHasAccessToPreviewModel.mockReturnValue(true); // Must have access
+    mockGetModel.mockReturnValue(PREVIEW_GEMINI_MODEL_AUTO);
     const { lastFrame, stdin } = renderComponent();
 
     // Select "Manual" (index 2 because Preview Auto is first, then Auto (Gemini 2.5))
-    stdin.write('\u001B[B'); // Arrow Down (to Auto (Gemini 2.5))
+    // Press down enough times to ensure we reach the bottom (Manual)
+    stdin.write('\u001B[B'); // Arrow Down
     await waitForUpdate();
-    stdin.write('\u001B[B'); // Arrow Down (to Manual)
+    stdin.write('\u001B[B'); // Arrow Down
     await waitForUpdate();
 
     // Press enter to select Manual
@@ -186,4 +194,50 @@ describe('<ModelDialog />', () => {
     // Should be back to main view (Manual option visible)
     expect(lastFrame()).toContain('Manual');
   });
+
+  describe('Preview Logic', () => {
+    it('should NOT show preview options if user has no access', () => {
+      mockGetHasAccessToPreviewModel.mockReturnValue(false);
+      mockGetPreviewFeatures.mockReturnValue(true); // Even if enabled
+      const { lastFrame } = renderComponent();
+      expect(lastFrame()).not.toContain('Auto (Preview)');
+    });
+
+    it('should NOT show preview options if user has access but preview features are disabled', () => {
+      mockGetHasAccessToPreviewModel.mockReturnValue(true);
+      mockGetPreviewFeatures.mockReturnValue(false);
+      const { lastFrame } = renderComponent();
+      expect(lastFrame()).not.toContain('Auto (Preview)');
+    });
+
+    it('should show preview options if user has access AND preview features are enabled', () => {
+      mockGetHasAccessToPreviewModel.mockReturnValue(true);
+      mockGetPreviewFeatures.mockReturnValue(true);
+      const { lastFrame } = renderComponent();
+      expect(lastFrame()).toContain('Auto (Preview)');
+    });
+
+    it('should show "Gemini 3 is now available" header if user has access but preview features disabled', () => {
+      mockGetHasAccessToPreviewModel.mockReturnValue(true);
+      mockGetPreviewFeatures.mockReturnValue(false);
+      const { lastFrame } = renderComponent();
+      expect(lastFrame()).toContain('Gemini 3 is now available.');
+      expect(lastFrame()).toContain('Enable "Preview features" in /settings');
+    });
+
+    it('should show "Gemini 3 is coming soon" header if user has no access', () => {
+      mockGetHasAccessToPreviewModel.mockReturnValue(false);
+      mockGetPreviewFeatures.mockReturnValue(false);
+      const { lastFrame } = renderComponent();
+      expect(lastFrame()).toContain('Gemini 3 is coming soon.');
+    });
+
+    it('should NOT show header/subheader if preview options are shown', () => {
+      mockGetHasAccessToPreviewModel.mockReturnValue(true);
+      mockGetPreviewFeatures.mockReturnValue(true);
+      const { lastFrame } = renderComponent();
+      expect(lastFrame()).not.toContain('Gemini 3 is now available.');
+      expect(lastFrame()).not.toContain('Gemini 3 is coming soon.');
+    });
+  });
 });

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

@@ -36,6 +36,9 @@ export function ModelDialog({ onClose }: ModelDialogProps): React.JSX.Element {
   // Determine the Preferred Model (read once when the dialog opens).
   const preferredModel = config?.getModel() || DEFAULT_GEMINI_MODEL_AUTO;
 
+  const shouldShowPreviewModels =
+    config?.getPreviewFeatures() && config.getHasAccessToPreviewModel();
+
   const manualModelSelected = useMemo(() => {
     const manualModels = [
       DEFAULT_GEMINI_MODEL,
@@ -82,7 +85,7 @@ export function ModelDialog({ onClose }: ModelDialogProps): React.JSX.Element {
       },
     ];
 
-    if (config?.getPreviewFeatures()) {
+    if (shouldShowPreviewModels) {
       list.unshift({
         value: PREVIEW_GEMINI_MODEL_AUTO,
         title: getDisplayString(PREVIEW_GEMINI_MODEL_AUTO),
@@ -92,7 +95,7 @@ export function ModelDialog({ onClose }: ModelDialogProps): React.JSX.Element {
       });
     }
     return list;
-  }, [config, manualModelSelected]);
+  }, [shouldShowPreviewModels, manualModelSelected]);
 
   const manualOptions = useMemo(() => {
     const list = [
@@ -113,7 +116,7 @@ export function ModelDialog({ onClose }: ModelDialogProps): React.JSX.Element {
       },
     ];
 
-    if (config?.getPreviewFeatures()) {
+    if (shouldShowPreviewModels) {
       list.unshift(
         {
           value: PREVIEW_GEMINI_MODEL,
@@ -128,7 +131,7 @@ export function ModelDialog({ onClose }: ModelDialogProps): React.JSX.Element {
       );
     }
     return list;
-  }, [config]);
+  }, [shouldShowPreviewModels]);
 
   const options = view === 'main' ? mainOptions : manualOptions;
 
@@ -163,13 +166,23 @@ export function ModelDialog({ onClose }: ModelDialogProps): React.JSX.Element {
     [config, onClose],
   );
 
-  const header = config?.getPreviewFeatures()
-    ? 'Gemini 3 is now enabled.'
-    : 'Gemini 3 is now available.';
+  let header;
+  let subheader;
 
-  const subheader = config?.getPreviewFeatures()
-    ? `To disable Gemini 3, disable "Preview features" in /settings.\nLearn more at https://goo.gle/enable-preview-features\n\nWhen you select Auto or Pro, Gemini CLI will attempt to use ${PREVIEW_GEMINI_MODEL} first, before falling back to ${DEFAULT_GEMINI_MODEL}.`
-    : `To use Gemini 3, enable "Preview features" in /settings.\nLearn more at https://goo.gle/enable-preview-features`;
+  // Do not show any header or subheader since it's already showing preview model
+  // options
+  if (shouldShowPreviewModels) {
+    header = undefined;
+    subheader = undefined;
+    // When a user has the access but has not enabled the preview features.
+  } else if (config?.getHasAccessToPreviewModel()) {
+    header = 'Gemini 3 is now available.';
+    subheader =
+      'Enable "Preview features" in /settings.\nLearn more at https://goo.gle/enable-preview-features';
+  } else {
+    header = 'Gemini 3 is coming soon.';
+    subheader = undefined;
+  }
 
   return (
     <Box
@@ -181,11 +194,15 @@ export function ModelDialog({ onClose }: ModelDialogProps): React.JSX.Element {
     >
       <Text bold>Select Model</Text>
 
-      <Box marginTop={1} marginBottom={1} flexDirection="column">
-        <ThemedGradient>
-          <Text>{header}</Text>
-        </ThemedGradient>
-        <Text>{subheader}</Text>
+      <Box flexDirection="column">
+        {header && (
+          <Box marginTop={1}>
+            <ThemedGradient>
+              <Text>{header}</Text>
+            </ThemedGradient>
+          </Box>
+        )}
+        {subheader && <Text>{subheader}</Text>}
       </Box>
 
       <Box marginTop={1}>