From efb6c3aa89cb5767e4422e3e4aa7d734f1b993e5 Mon Sep 17 00:00:00 2001 From: Sehoon Shon Date: Wed, 25 Feb 2026 15:16:21 -0500 Subject: [PATCH] fix(cli): support quota error fallbacks for all authentication types. (#20338) Previously, model fallbacks on persistent quota errors were restricted to users authenticated with LOGIN_WITH_GOOGLE. This commit enables fallback prompts for all authentication methods, conditionally rendering the 'switch to API key' message exclusively for OAuth users. Additionally, this updates the environment parser to properly respect CLOUD_SHELL and COMPUTE_ADC variables. --- .../cli/src/ui/components/DialogManager.tsx | 1 + .../src/ui/components/ProQuotaDialog.test.tsx | 51 +++++++++++++++---- .../cli/src/ui/components/ProQuotaDialog.tsx | 17 +++++-- .../cli/src/ui/contexts/UIStateContext.tsx | 2 + .../src/ui/hooks/useQuotaAndFallback.test.ts | 36 ++++++++----- .../cli/src/ui/hooks/useQuotaAndFallback.ts | 12 ++--- packages/core/src/core/contentGenerator.ts | 6 +++ 7 files changed, 89 insertions(+), 36 deletions(-) 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; }