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}>

packages/core/src/config/config.test.ts

@@ -34,7 +34,11 @@ import { logRipgrepFallback } from '../telemetry/loggers.js';
 import { RipgrepFallbackEvent } from '../telemetry/types.js';
 import { ToolRegistry } from '../tools/tool-registry.js';
 import { DEFAULT_MODEL_CONFIGS } from './defaultModelConfigs.js';
-import { DEFAULT_GEMINI_MODEL } from './models.js';
+import {
+  DEFAULT_GEMINI_MODEL,
+  DEFAULT_GEMINI_MODEL_AUTO,
+  PREVIEW_GEMINI_MODEL,
+} from './models.js';
 
 vi.mock('fs', async (importOriginal) => {
   const actual = await importOriginal<typeof import('fs')>();
@@ -1736,3 +1740,107 @@ describe('Availability Service Integration', () => {
     expect(spy).toHaveBeenCalled();
   });
 });
+
+describe('Config Quota & Preview Model Access', () => {
+  let config: Config;
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  let mockCodeAssistServer: any;
+
+  const baseParams: ConfigParameters = {
+    cwd: '/tmp',
+    targetDir: '/tmp',
+    debugMode: false,
+    sessionId: 'test-session',
+    model: 'gemini-pro',
+    usageStatisticsEnabled: false,
+    embeddingModel: 'gemini-embedding', // required in type but not in the original file I copied, adding here
+    sandbox: {
+      command: 'docker',
+      image: 'gemini-cli-sandbox',
+    },
+  };
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+    mockCodeAssistServer = {
+      projectId: 'test-project',
+      retrieveUserQuota: vi.fn(),
+    };
+    vi.mocked(getCodeAssistServer).mockReturnValue(mockCodeAssistServer);
+    config = new Config(baseParams);
+  });
+
+  describe('refreshUserQuota', () => {
+    it('should update hasAccessToPreviewModel to true if quota includes preview model', async () => {
+      mockCodeAssistServer.retrieveUserQuota.mockResolvedValue({
+        buckets: [{ modelId: PREVIEW_GEMINI_MODEL }],
+      });
+
+      await config.refreshUserQuota();
+      expect(config.getHasAccessToPreviewModel()).toBe(true);
+    });
+
+    it('should update hasAccessToPreviewModel to false if quota does not include preview model', async () => {
+      mockCodeAssistServer.retrieveUserQuota.mockResolvedValue({
+        buckets: [{ modelId: 'some-other-model' }],
+      });
+
+      await config.refreshUserQuota();
+      expect(config.getHasAccessToPreviewModel()).toBe(false);
+    });
+
+    it('should update hasAccessToPreviewModel to false if buckets are undefined', async () => {
+      mockCodeAssistServer.retrieveUserQuota.mockResolvedValue({});
+
+      await config.refreshUserQuota();
+      expect(config.getHasAccessToPreviewModel()).toBe(false);
+    });
+
+    it('should return undefined and not update if codeAssistServer is missing', async () => {
+      vi.mocked(getCodeAssistServer).mockReturnValue(undefined);
+      const result = await config.refreshUserQuota();
+      expect(result).toBeUndefined();
+      expect(config.getHasAccessToPreviewModel()).toBe(false);
+    });
+
+    it('should return undefined if retrieveUserQuota fails', async () => {
+      mockCodeAssistServer.retrieveUserQuota.mockRejectedValue(
+        new Error('Network error'),
+      );
+      const result = await config.refreshUserQuota();
+      expect(result).toBeUndefined();
+      // Should remain default (false)
+      expect(config.getHasAccessToPreviewModel()).toBe(false);
+    });
+  });
+
+  describe('setPreviewFeatures', () => {
+    it('should reset model to default auto if disabling preview features while using a preview model', () => {
+      config.setPreviewFeatures(true);
+      config.setModel(PREVIEW_GEMINI_MODEL);
+
+      config.setPreviewFeatures(false);
+
+      expect(config.getModel()).toBe(DEFAULT_GEMINI_MODEL_AUTO);
+    });
+
+    it('should NOT reset model if disabling preview features while NOT using a preview model', () => {
+      config.setPreviewFeatures(true);
+      const nonPreviewModel = 'gemini-1.5-pro';
+      config.setModel(nonPreviewModel);
+
+      config.setPreviewFeatures(false);
+
+      expect(config.getModel()).toBe(nonPreviewModel);
+    });
+
+    it('should NOT reset model if enabling preview features', () => {
+      config.setPreviewFeatures(false);
+      config.setModel(PREVIEW_GEMINI_MODEL); // Just pretending it was set somehow
+
+      config.setPreviewFeatures(true);
+
+      expect(config.getModel()).toBe(PREVIEW_GEMINI_MODEL);
+    });
+  });
+});

packages/core/src/config/config.ts

@@ -48,8 +48,10 @@ import { tokenLimit } from '../core/tokenLimits.js';
 import {
   DEFAULT_GEMINI_EMBEDDING_MODEL,
   DEFAULT_GEMINI_FLASH_MODEL,
+  DEFAULT_GEMINI_MODEL_AUTO,
   DEFAULT_THINKING_MODE,
   isPreviewModel,
+  PREVIEW_GEMINI_MODEL,
 } from './models.js';
 import { shouldAttemptBrowserLaunch } from '../utils/browser.js';
 import type { MCPOAuthConfig } from '../mcp/oauth-provider.js';
@@ -81,6 +83,7 @@ import { PolicyEngine } from '../policy/policy-engine.js';
 import type { PolicyEngineConfig } from '../policy/types.js';
 import { HookSystem } from '../hooks/index.js';
 import type { UserTierId } from '../code_assist/types.js';
+import type { RetrieveUserQuotaResponse } from '../code_assist/types.js';
 import { getCodeAssistServer } from '../code_assist/codeAssist.js';
 import type { Experiments } from '../code_assist/experiments/experiments.js';
 import { AgentRegistry } from '../agents/registry.js';
@@ -379,6 +382,7 @@ export class Config {
   private readonly bugCommand: BugCommandSettings | undefined;
   private model: string;
   private previewFeatures: boolean | undefined;
+  private hasAccessToPreviewModel: boolean = false;
   private readonly noBrowser: boolean;
   private readonly folderTrust: boolean;
   private ideMode: boolean;
@@ -732,6 +736,10 @@ export class Config {
 
     const codeAssistServer = getCodeAssistServer(this);
     if (codeAssistServer) {
+      if (codeAssistServer.projectId) {
+        await this.refreshUserQuota();
+      }
+
       this.experimentsPromise = getExperiments(codeAssistServer)
         .then((experiments) => {
           this.setExperiments(experiments);
@@ -753,8 +761,21 @@ export class Config {
       this.experimentsPromise = undefined;
     }
 
+    const authType = this.contentGeneratorConfig.authType;
+    if (
+      authType === AuthType.USE_GEMINI ||
+      authType === AuthType.USE_VERTEX_AI
+    ) {
+      this.setHasAccessToPreviewModel(true);
+    }
+
     // Reset the session flag since we're explicitly changing auth and using default model
     this.inFallbackMode = false;
+
+    // Update model if user no longer has access to the preview model
+    if (!this.hasAccessToPreviewModel && isPreviewModel(this.model)) {
+      this.setModel(DEFAULT_GEMINI_MODEL_AUTO);
+    }
   }
 
   async getExperimentsAsync(): Promise<Experiments | undefined> {
@@ -948,9 +969,43 @@ export class Config {
   }
 
   setPreviewFeatures(previewFeatures: boolean) {
+    // If it's using a preview model and it's turning off previewFeatures,
+    // switch the model to the default auto mode.
+    if (this.previewFeatures && !previewFeatures) {
+      if (isPreviewModel(this.getModel())) {
+        this.setModel(DEFAULT_GEMINI_MODEL_AUTO);
+      }
+    }
     this.previewFeatures = previewFeatures;
   }
 
+  getHasAccessToPreviewModel(): boolean {
+    return this.hasAccessToPreviewModel;
+  }
+
+  setHasAccessToPreviewModel(hasAccess: boolean): void {
+    this.hasAccessToPreviewModel = hasAccess;
+  }
+
+  async refreshUserQuota(): Promise<RetrieveUserQuotaResponse | undefined> {
+    const codeAssistServer = getCodeAssistServer(this);
+    if (!codeAssistServer || !codeAssistServer.projectId) {
+      return undefined;
+    }
+    try {
+      const quota = await codeAssistServer.retrieveUserQuota({
+        project: codeAssistServer.projectId,
+      });
+      const hasAccess =
+        quota.buckets?.some((b) => b.modelId === PREVIEW_GEMINI_MODEL) ?? false;
+      this.setHasAccessToPreviewModel(hasAccess);
+      return quota;
+    } catch (e) {
+      debugLogger.debug('Failed to retrieve user quota', e);
+      return undefined;
+    }
+  }
+
   getCoreTools(): string[] | undefined {
     return this.coreTools;
   }