diff --git a/packages/cli/src/ui/components/DialogManager.tsx b/packages/cli/src/ui/components/DialogManager.tsx index 3d56c68e5b..c90194052a 100644 --- a/packages/cli/src/ui/components/DialogManager.tsx +++ b/packages/cli/src/ui/components/DialogManager.tsx @@ -135,6 +135,7 @@ export const DialogManager = ({ isModelNotFoundError={ !!uiState.quota.proQuotaRequest.isModelNotFoundError } + authType={uiState.quota.proQuotaRequest.authType} onChoice={uiActions.handleProQuotaChoice} /> ); diff --git a/packages/cli/src/ui/components/ProQuotaDialog.test.tsx b/packages/cli/src/ui/components/ProQuotaDialog.test.tsx index f74f5fa447..d97d53314e 100644 --- a/packages/cli/src/ui/components/ProQuotaDialog.test.tsx +++ b/packages/cli/src/ui/components/ProQuotaDialog.test.tsx @@ -13,6 +13,7 @@ import { RadioButtonSelect } from './shared/RadioButtonSelect.js'; import { PREVIEW_GEMINI_MODEL, DEFAULT_GEMINI_FLASH_MODEL, + AuthType, } from '@google/gemini-cli-core'; // Mock the child component to make it easier to test the parent @@ -62,7 +63,7 @@ describe('ProQuotaDialog', () => { describe('for non-flash model failures', () => { describe('when it is a terminal quota error', () => { - it('should render switch, upgrade, and stop options for paid tiers', () => { + it('should render switch, upgrade, and stop options for LOGIN_WITH_GOOGLE', () => { const { unmount } = render( { message="paid tier quota error" isTerminalQuotaError={true} isModelNotFoundError={false} + authType={AuthType.LOGIN_WITH_GOOGLE} onChoice={mockOnChoice} />, ); @@ -99,6 +101,39 @@ describe('ProQuotaDialog', () => { unmount(); }); + it('should NOT render upgrade option for USE_GEMINI', () => { + const { unmount } = render( + , + ); + + expect(RadioButtonSelect).toHaveBeenCalledWith( + expect.objectContaining({ + items: [ + { + label: 'Switch to gemini-2.5-flash', + value: 'retry_always', + key: 'retry_always', + }, + { + label: 'Stop', + value: 'retry_later', + key: 'retry_later', + }, + ], + }), + undefined, + ); + unmount(); + }); + it('should render "Keep trying" and "Stop" options when failed model and fallback model are the same', () => { const { unmount } = render( { unmount(); }); - it('should render switch, upgrade, and stop options for free tier', () => { + it('should render switch, upgrade, and stop options for LOGIN_WITH_GOOGLE (free tier)', () => { const { unmount } = render( { message="free tier quota error" isTerminalQuotaError={true} isModelNotFoundError={false} + authType={AuthType.LOGIN_WITH_GOOGLE} onChoice={mockOnChoice} />, ); @@ -204,7 +240,7 @@ describe('ProQuotaDialog', () => { }); describe('when it is a model not found error', () => { - it('should render switch and stop options regardless of tier', () => { + it('should render switch, upgrade, and stop options for LOGIN_WITH_GOOGLE', () => { const { unmount } = render( { message="You don't have access to gemini-3-pro-preview yet." isTerminalQuotaError={false} isModelNotFoundError={true} + authType={AuthType.LOGIN_WITH_GOOGLE} onChoice={mockOnChoice} />, ); @@ -241,7 +278,7 @@ describe('ProQuotaDialog', () => { unmount(); }); - it('should render switch and stop options for paid tier as well', () => { + it('should NOT render upgrade option for USE_GEMINI', () => { const { unmount } = render( { message="You don't have access to gemini-3-pro-preview yet." isTerminalQuotaError={false} isModelNotFoundError={true} + authType={AuthType.USE_GEMINI} onChoice={mockOnChoice} />, ); @@ -261,11 +299,6 @@ describe('ProQuotaDialog', () => { value: 'retry_always', key: 'retry_always', }, - { - label: 'Upgrade for higher limits', - value: 'upgrade', - key: 'upgrade', - }, { label: 'Stop', value: 'retry_later', diff --git a/packages/cli/src/ui/components/ProQuotaDialog.tsx b/packages/cli/src/ui/components/ProQuotaDialog.tsx index ccc20b3e75..82a679db8c 100644 --- a/packages/cli/src/ui/components/ProQuotaDialog.tsx +++ b/packages/cli/src/ui/components/ProQuotaDialog.tsx @@ -8,6 +8,7 @@ import type React from 'react'; import { Box, Text } from 'ink'; import { RadioButtonSelect } from './shared/RadioButtonSelect.js'; import { theme } from '../semantic-colors.js'; +import { AuthType } from '@google/gemini-cli-core'; interface ProQuotaDialogProps { failedModel: string; @@ -15,6 +16,7 @@ interface ProQuotaDialogProps { message: string; isTerminalQuotaError: boolean; isModelNotFoundError?: boolean; + authType?: AuthType; onChoice: ( choice: 'retry_later' | 'retry_once' | 'retry_always' | 'upgrade', ) => void; @@ -26,6 +28,7 @@ export function ProQuotaDialog({ message, isTerminalQuotaError, isModelNotFoundError, + authType, onChoice, }: ProQuotaDialogProps): React.JSX.Element { let items; @@ -51,11 +54,15 @@ export function ProQuotaDialog({ value: 'retry_always' as const, key: 'retry_always', }, - { - label: 'Upgrade for higher limits', - value: 'upgrade' as const, - key: 'upgrade', - }, + ...(authType === AuthType.LOGIN_WITH_GOOGLE + ? [ + { + label: 'Upgrade for higher limits', + value: 'upgrade' as const, + key: 'upgrade', + }, + ] + : []), { label: `Stop`, value: 'retry_later' as const, diff --git a/packages/cli/src/ui/contexts/UIStateContext.tsx b/packages/cli/src/ui/contexts/UIStateContext.tsx index 9fb2852361..79464271b8 100644 --- a/packages/cli/src/ui/contexts/UIStateContext.tsx +++ b/packages/cli/src/ui/contexts/UIStateContext.tsx @@ -24,6 +24,7 @@ import type { ApprovalMode, UserTierId, IdeInfo, + AuthType, FallbackIntent, ValidationIntent, AgentDefinition, @@ -42,6 +43,7 @@ export interface ProQuotaDialogRequest { message: string; isTerminalQuotaError: boolean; isModelNotFoundError?: boolean; + authType?: AuthType; resolve: (intent: FallbackIntent) => void; } diff --git a/packages/cli/src/ui/hooks/useQuotaAndFallback.test.ts b/packages/cli/src/ui/hooks/useQuotaAndFallback.test.ts index 5d6db5abfa..2272de5bf9 100644 --- a/packages/cli/src/ui/hooks/useQuotaAndFallback.test.ts +++ b/packages/cli/src/ui/hooks/useQuotaAndFallback.test.ts @@ -96,9 +96,13 @@ describe('useQuotaAndFallback', () => { }); describe('Fallback Handler Logic', () => { - // Helper function to render the hook and extract the registered handler - const getRegisteredHandler = (): FallbackModelHandler => { - renderHook(() => + it('should show fallback dialog but omit switch to API key message if authType is not LOGIN_WITH_GOOGLE', async () => { + // Override the default mock from beforeEach for this specific test + vi.spyOn(mockConfig, 'getContentGeneratorConfig').mockReturnValue({ + authType: AuthType.USE_GEMINI, + }); + + const { result } = renderHook(() => useQuotaAndFallback({ config: mockConfig, historyManager: mockHistoryManager, @@ -107,20 +111,24 @@ describe('useQuotaAndFallback', () => { onShowAuthSelection: mockOnShowAuthSelection, }), ); - return setFallbackHandlerSpy.mock.calls[0][0] as FallbackModelHandler; - }; - it('should return null and take no action if authType is not LOGIN_WITH_GOOGLE', async () => { - // Override the default mock from beforeEach for this specific test - vi.spyOn(mockConfig, 'getContentGeneratorConfig').mockReturnValue({ - authType: AuthType.USE_GEMINI, + const handler = setFallbackHandlerSpy.mock + .calls[0][0] as FallbackModelHandler; + + const error = new TerminalQuotaError( + 'pro quota', + mockGoogleApiError, + 1000 * 60 * 5, + ); + + act(() => { + void handler('gemini-pro', 'gemini-flash', error); }); - const handler = getRegisteredHandler(); - const result = await handler('gemini-pro', 'gemini-flash', new Error()); - - expect(result).toBeNull(); - expect(mockHistoryManager.addItem).not.toHaveBeenCalled(); + expect(result.current.proQuotaRequest).not.toBeNull(); + expect(result.current.proQuotaRequest?.message).not.toContain( + '/auth to switch to API key.', + ); }); describe('Interactive Fallback', () => { diff --git a/packages/cli/src/ui/hooks/useQuotaAndFallback.ts b/packages/cli/src/ui/hooks/useQuotaAndFallback.ts index 1ba03f2a47..a9e2b0c867 100644 --- a/packages/cli/src/ui/hooks/useQuotaAndFallback.ts +++ b/packages/cli/src/ui/hooks/useQuotaAndFallback.ts @@ -55,14 +55,7 @@ export function useQuotaAndFallback({ fallbackModel, error, ): Promise => { - // Fallbacks are currently only handled for OAuth users. const contentGeneratorConfig = config.getContentGeneratorConfig(); - if ( - !contentGeneratorConfig || - contentGeneratorConfig.authType !== AuthType.LOGIN_WITH_GOOGLE - ) { - return null; - } let message: string; let isTerminalQuotaError = false; @@ -78,7 +71,9 @@ export function useQuotaAndFallback({ error.retryDelayMs ? getResetTimeMessage(error.retryDelayMs) : null, `/stats model for usage details`, `/model to switch models.`, - `/auth to switch to API key.`, + contentGeneratorConfig?.authType === AuthType.LOGIN_WITH_GOOGLE + ? `/auth to switch to API key.` + : null, ].filter(Boolean); message = messageLines.join('\n'); } else if (error instanceof ModelNotFoundError) { @@ -122,6 +117,7 @@ export function useQuotaAndFallback({ message, isTerminalQuotaError, isModelNotFoundError, + authType: contentGeneratorConfig?.authType, }); }, ); diff --git a/packages/core/src/core/contentGenerator.ts b/packages/core/src/core/contentGenerator.ts index 98d8d50020..29f2ff03df 100644 --- a/packages/core/src/core/contentGenerator.ts +++ b/packages/core/src/core/contentGenerator.ts @@ -77,6 +77,12 @@ export function getAuthTypeFromEnv(): AuthType | undefined { if (process.env['GEMINI_API_KEY']) { return AuthType.USE_GEMINI; } + if ( + process.env['CLOUD_SHELL'] === 'true' || + process.env['GEMINI_CLI_USE_COMPUTE_ADC'] === 'true' + ) { + return AuthType.COMPUTE_ADC; + } return undefined; }