MediaMetz/gemini-cli
1
0
Fork 0
mirror of https://github.com/google-gemini/gemini-cli.git synced 2026-03-11 22:51:00 -07:00
Code Issues Packages Projects Releases Wiki Activity

fix(cli): support quota error fallbacks for all authentication types. (#20338)

Browse Source 
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.
This commit is contained in:
Sehoon Shon
2026-02-25 15:16:21 -05:00
parent ac454bcfbc
commit efb6c3aa89
7 changed files with 89 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
packages/cli/src/ui/components/DialogManager.tsx
View File

@@ -135,6 +135,7 @@ export const DialogManager = ({
isModelNotFoundError={
!!uiState.quota.proQuotaRequest.isModelNotFoundError
}
authType={uiState.quota.proQuotaRequest.authType}
onChoice={uiActions.handleProQuotaChoice}
/>
);

51
packages/cli/src/ui/components/ProQuotaDialog.test.tsx
View File

@@ -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(
<ProQuotaDialog
failedModel="gemini-2.5-pro"
@@ -70,6 +71,7 @@ describe('ProQuotaDialog', () => {
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(
<ProQuotaDialog
failedModel="gemini-2.5-pro"
fallbackModel="gemini-2.5-flash"
message="paid tier quota error"
isTerminalQuotaError={true}
isModelNotFoundError={false}
authType={AuthType.USE_GEMINI}
onChoice={mockOnChoice}
/>,
);
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(
<ProQuotaDialog
@@ -130,7 +165,7 @@ describe('ProQuotaDialog', () => {
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(
<ProQuotaDialog
failedModel="gemini-2.5-pro"
@@ -138,6 +173,7 @@ describe('ProQuotaDialog', () => {
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(
<ProQuotaDialog
failedModel="gemini-3-pro-preview"
@@ -212,6 +248,7 @@ describe('ProQuotaDialog', () => {
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(
<ProQuotaDialog
failedModel="gemini-3-pro-preview"
@@ -249,6 +286,7 @@ describe('ProQuotaDialog', () => {
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',

17
packages/cli/src/ui/components/ProQuotaDialog.tsx
View File

@@ -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,

2
packages/cli/src/ui/contexts/UIStateContext.tsx
View File

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

36
packages/cli/src/ui/hooks/useQuotaAndFallback.test.ts
View File

@@ -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', () => {

12
packages/cli/src/ui/hooks/useQuotaAndFallback.ts
View File

@@ -55,14 +55,7 @@ export function useQuotaAndFallback({
fallbackModel,
error,
): Promise<FallbackIntent | null> => {
// 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,
});
},
);

6
packages/core/src/core/contentGenerator.ts
View File

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