diff --git a/packages/cli/src/ui/commands/modelCommand.test.ts b/packages/cli/src/ui/commands/modelCommand.test.ts index 8bab5962b9..ed2da93a1c 100644 --- a/packages/cli/src/ui/commands/modelCommand.test.ts +++ b/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( diff --git a/packages/cli/src/ui/commands/modelCommand.ts b/packages/cli/src/ui/commands/modelCommand.ts index d355d0521a..fd89223a7c 100644 --- a/packages/cli/src/ui/commands/modelCommand.ts +++ b/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', + }; + }, }; diff --git a/packages/cli/src/ui/commands/statsCommand.test.ts b/packages/cli/src/ui/commands/statsCommand.test.ts index 13d007ac2e..2a054ecc4d 100644 --- a/packages/cli/src/ui/commands/statsCommand.test.ts +++ b/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', diff --git a/packages/cli/src/ui/commands/statsCommand.ts b/packages/cli/src/ui/commands/statsCommand.ts index 5657d826b1..718da86f69 100644 --- a/packages/cli/src/ui/commands/statsCommand.ts +++ b/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; } } diff --git a/packages/cli/src/ui/components/ModelDialog.test.tsx b/packages/cli/src/ui/components/ModelDialog.test.tsx index 25b472dee9..390be01b74 100644 --- a/packages/cli/src/ui/components/ModelDialog.test.tsx +++ b/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('', () => { const mockGetModel = vi.fn(); const mockGetPreviewFeatures = vi.fn(); const mockOnClose = vi.fn(); + const mockGetHasAccessToPreviewModel = vi.fn(); interface MockConfig extends Partial { 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('', () => { 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('', () => { 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('', () => { // 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.'); + }); + }); }); diff --git a/packages/cli/src/ui/components/ModelDialog.tsx b/packages/cli/src/ui/components/ModelDialog.tsx index ff92e96939..5be22beda8 100644 --- a/packages/cli/src/ui/components/ModelDialog.tsx +++ b/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 ( Select Model - - - {header} - - {subheader} + + {header && ( + + + {header} + + + )} + {subheader && {subheader}} diff --git a/packages/core/src/config/config.test.ts b/packages/core/src/config/config.test.ts index b1db06cfad..f0c61ed454 100644 --- a/packages/core/src/config/config.test.ts +++ b/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(); @@ -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); + }); + }); +}); diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index bfb91dbf81..845f174529 100644 --- a/packages/core/src/config/config.ts +++ b/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 { @@ -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 { + 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; }