feat(billing): implement G1 AI credits overage flow with billing telemetry (#18590)

This commit is contained in:
Gaurav
2026-02-27 10:15:06 -08:00
committed by GitHub
parent fdd844b405
commit b2d6844f9b
55 changed files with 3182 additions and 23 deletions

View File

@@ -47,6 +47,7 @@ import {
type IdeInfo,
type IdeContext,
type UserTierId,
type GeminiUserTier,
type UserFeedbackPayload,
type AgentDefinition,
type ApprovalMode,
@@ -82,6 +83,8 @@ import {
CoreToolCallStatus,
generateSteeringAckMessage,
buildUserSteeringHintPrompt,
logBillingEvent,
ApiKeyUpdatedEvent,
} from '@google/gemini-cli-core';
import { validateAuthMethod } from '../config/auth.js';
import process from 'node:process';
@@ -391,6 +394,9 @@ export const AppContainer = (props: AppContainerProps) => {
? { remaining, limit, resetTime }
: undefined;
});
const [paidTier, setPaidTier] = useState<GeminiUserTier | undefined>(
undefined,
);
const [isConfigInitialized, setConfigInitialized] = useState(false);
@@ -686,10 +692,17 @@ export const AppContainer = (props: AppContainerProps) => {
handleProQuotaChoice,
validationRequest,
handleValidationChoice,
// G1 AI Credits
overageMenuRequest,
handleOverageMenuChoice,
emptyWalletRequest,
handleEmptyWalletChoice,
} = useQuotaAndFallback({
config,
historyManager,
userTier,
paidTier,
settings,
setModelSwitchedFromQuotaError,
onShowAuthSelection: () => setAuthState(AuthState.Updating),
});
@@ -729,6 +742,8 @@ export const AppContainer = (props: AppContainerProps) => {
const handleAuthSelect = useCallback(
async (authType: AuthType | undefined, scope: LoadableSettingScope) => {
if (authType) {
const previousAuthType =
config.getContentGeneratorConfig()?.authType ?? 'unknown';
if (authType === AuthType.LOGIN_WITH_GOOGLE) {
setAuthContext({ requiresRestart: true });
} else {
@@ -741,6 +756,10 @@ export const AppContainer = (props: AppContainerProps) => {
config.setRemoteAdminSettings(undefined);
await config.refreshAuth(authType);
setAuthState(AuthState.Authenticated);
logBillingEvent(
config,
new ApiKeyUpdatedEvent(previousAuthType, authType),
);
} catch (e) {
if (e instanceof ChangeAuthRequestedError) {
return;
@@ -803,6 +822,7 @@ Logging in with Google... Restarting Gemini CLI to continue.
// Only sync when not currently authenticating
if (authState === AuthState.Authenticated) {
setUserTier(config.getUserTier());
setPaidTier(config.getUserPaidTier());
}
}, [config, authState]);
@@ -2006,6 +2026,8 @@ Logging in with Google... Restarting Gemini CLI to continue.
showIdeRestartPrompt ||
!!proQuotaRequest ||
!!validationRequest ||
!!overageMenuRequest ||
!!emptyWalletRequest ||
isSessionBrowserOpen ||
authState === AuthState.AwaitingApiKeyInput ||
!!newAgents;
@@ -2033,6 +2055,8 @@ Logging in with Google... Restarting Gemini CLI to continue.
hasLoopDetectionConfirmationRequest ||
!!proQuotaRequest ||
!!validationRequest ||
!!overageMenuRequest ||
!!emptyWalletRequest ||
!!customDialog;
const allowPlanMode =
@@ -2243,6 +2267,9 @@ Logging in with Google... Restarting Gemini CLI to continue.
stats: quotaStats,
proQuotaRequest,
validationRequest,
// G1 AI Credits dialog state
overageMenuRequest,
emptyWalletRequest,
},
contextFileNames,
errorCount,
@@ -2367,6 +2394,8 @@ Logging in with Google... Restarting Gemini CLI to continue.
quotaStats,
proQuotaRequest,
validationRequest,
overageMenuRequest,
emptyWalletRequest,
contextFileNames,
errorCount,
availableTerminalHeight,
@@ -2448,6 +2477,9 @@ Logging in with Google... Restarting Gemini CLI to continue.
handleClearScreen,
handleProQuotaChoice,
handleValidationChoice,
// G1 AI Credits handlers
handleOverageMenuChoice,
handleEmptyWalletChoice,
openSessionBrowser,
closeSessionBrowser,
handleResumeSession,
@@ -2534,6 +2566,8 @@ Logging in with Google... Restarting Gemini CLI to continue.
handleClearScreen,
handleProQuotaChoice,
handleValidationChoice,
handleOverageMenuChoice,
handleEmptyWalletChoice,
openSessionBrowser,
closeSessionBrowser,
handleResumeSession,

View File

@@ -39,11 +39,18 @@ describe('statsCommand', () => {
mockContext.session.stats.sessionStartTime = startTime;
});
it('should display general session stats when run with no subcommand', () => {
it('should display general session stats when run with no subcommand', async () => {
if (!statsCommand.action) throw new Error('Command has no action');
// eslint-disable-next-line @typescript-eslint/no-floating-promises
statsCommand.action(mockContext, '');
mockContext.services.config = {
refreshUserQuota: vi.fn(),
refreshAvailableCredits: vi.fn(),
getUserTierName: vi.fn(),
getUserPaidTier: vi.fn(),
getModel: vi.fn(),
} as unknown as Config;
await statsCommand.action(mockContext, '');
const expectedDuration = formatDuration(
endTime.getTime() - startTime.getTime(),
@@ -55,6 +62,7 @@ describe('statsCommand', () => {
tier: undefined,
userEmail: 'mock@example.com',
currentModel: undefined,
creditBalance: undefined,
});
});
@@ -78,6 +86,8 @@ describe('statsCommand', () => {
getQuotaRemaining: mockGetQuotaRemaining,
getQuotaLimit: mockGetQuotaLimit,
getQuotaResetTime: mockGetQuotaResetTime,
getUserPaidTier: vi.fn(),
refreshAvailableCredits: vi.fn(),
} as unknown as Config;
await statsCommand.action(mockContext, '');

View File

@@ -11,7 +11,10 @@ import type {
} from '../types.js';
import { MessageType } from '../types.js';
import { formatDuration } from '../utils/formatters.js';
import { UserAccountManager } from '@google/gemini-cli-core';
import {
UserAccountManager,
getG1CreditBalance,
} from '@google/gemini-cli-core';
import {
type CommandContext,
type SlashCommand,
@@ -27,8 +30,10 @@ function getUserIdentity(context: CommandContext) {
const userEmail = cachedAccount ?? undefined;
const tier = context.services.config?.getUserTierName();
const paidTier = context.services.config?.getUserPaidTier();
const creditBalance = getG1CreditBalance(paidTier) ?? undefined;
return { selectedAuthType, userEmail, tier };
return { selectedAuthType, userEmail, tier, creditBalance };
}
async function defaultSessionView(context: CommandContext) {
@@ -43,7 +48,8 @@ async function defaultSessionView(context: CommandContext) {
}
const wallDuration = now.getTime() - sessionStartTime.getTime();
const { selectedAuthType, userEmail, tier } = getUserIdentity(context);
const { selectedAuthType, userEmail, tier, creditBalance } =
getUserIdentity(context);
const currentModel = context.services.config?.getModel();
const statsItem: HistoryItemStats = {
@@ -53,10 +59,14 @@ async function defaultSessionView(context: CommandContext) {
userEmail,
tier,
currentModel,
creditBalance,
};
if (context.services.config) {
const quota = await context.services.config.refreshUserQuota();
const [quota] = await Promise.all([
context.services.config.refreshUserQuota(),
context.services.config.refreshAvailableCredits(),
]);
if (quota) {
statsItem.quotas = quota;
statsItem.pooledRemaining = context.services.config.getQuotaRemaining();

View File

@@ -80,6 +80,8 @@ describe('DialogManager', () => {
stats: undefined,
proQuotaRequest: null,
validationRequest: null,
overageMenuRequest: null,
emptyWalletRequest: null,
},
shouldShowIdePrompt: false,
isFolderTrustDialogOpen: false,
@@ -132,6 +134,8 @@ describe('DialogManager', () => {
resolve: vi.fn(),
},
validationRequest: null,
overageMenuRequest: null,
emptyWalletRequest: null,
},
},
'ProQuotaDialog',

View File

@@ -18,6 +18,8 @@ import { EditorSettingsDialog } from './EditorSettingsDialog.js';
import { PrivacyNotice } from '../privacy/PrivacyNotice.js';
import { ProQuotaDialog } from './ProQuotaDialog.js';
import { ValidationDialog } from './ValidationDialog.js';
import { OverageMenuDialog } from './OverageMenuDialog.js';
import { EmptyWalletDialog } from './EmptyWalletDialog.js';
import { runExitCleanup } from '../../utils/cleanup.js';
import { RELAUNCH_EXIT_CODE } from '../../utils/processUtils.js';
import { SessionBrowser } from './SessionBrowser.js';
@@ -152,6 +154,28 @@ export const DialogManager = ({
/>
);
}
if (uiState.quota.overageMenuRequest) {
return (
<OverageMenuDialog
failedModel={uiState.quota.overageMenuRequest.failedModel}
fallbackModel={uiState.quota.overageMenuRequest.fallbackModel}
resetTime={uiState.quota.overageMenuRequest.resetTime}
creditBalance={uiState.quota.overageMenuRequest.creditBalance}
onChoice={uiActions.handleOverageMenuChoice}
/>
);
}
if (uiState.quota.emptyWalletRequest) {
return (
<EmptyWalletDialog
failedModel={uiState.quota.emptyWalletRequest.failedModel}
fallbackModel={uiState.quota.emptyWalletRequest.fallbackModel}
resetTime={uiState.quota.emptyWalletRequest.resetTime}
onGetCredits={uiState.quota.emptyWalletRequest.onGetCredits}
onChoice={uiActions.handleEmptyWalletChoice}
/>
);
}
if (uiState.shouldShowIdePrompt) {
return (
<IdeIntegrationNudge

View File

@@ -0,0 +1,218 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { renderWithProviders } from '../../test-utils/render.js';
import { waitFor } from '../../test-utils/async.js';
import { act } from 'react';
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
import { EmptyWalletDialog } from './EmptyWalletDialog.js';
const writeKey = (stdin: { write: (data: string) => void }, key: string) => {
act(() => {
stdin.write(key);
});
};
describe('EmptyWalletDialog', () => {
const mockOnChoice = vi.fn();
const mockOnGetCredits = vi.fn();
beforeEach(() => {
vi.clearAllMocks();
});
afterEach(() => {
vi.restoreAllMocks();
});
describe('rendering', () => {
it('should match snapshot with fallback available', async () => {
const { lastFrame, unmount, waitUntilReady } = renderWithProviders(
<EmptyWalletDialog
failedModel="gemini-2.5-pro"
fallbackModel="gemini-3-flash-preview"
resetTime="2:00 PM"
onChoice={mockOnChoice}
/>,
);
await waitUntilReady();
expect(lastFrame()).toMatchSnapshot();
unmount();
});
it('should match snapshot without fallback', async () => {
const { lastFrame, unmount, waitUntilReady } = renderWithProviders(
<EmptyWalletDialog
failedModel="gemini-2.5-pro"
onChoice={mockOnChoice}
/>,
);
await waitUntilReady();
expect(lastFrame()).toMatchSnapshot();
unmount();
});
it('should display the model name and usage limit message', async () => {
const { lastFrame, unmount, waitUntilReady } = renderWithProviders(
<EmptyWalletDialog
failedModel="gemini-2.5-pro"
onChoice={mockOnChoice}
/>,
);
await waitUntilReady();
const output = lastFrame() ?? '';
expect(output).toContain('gemini-2.5-pro');
expect(output).toContain('Usage limit reached');
unmount();
});
it('should display purchase prompt and credits update notice', async () => {
const { lastFrame, unmount, waitUntilReady } = renderWithProviders(
<EmptyWalletDialog
failedModel="gemini-2.5-pro"
onChoice={mockOnChoice}
/>,
);
await waitUntilReady();
const output = lastFrame() ?? '';
expect(output).toContain('purchase more AI Credits');
expect(output).toContain(
'Newly purchased AI credits may take a few minutes to update',
);
unmount();
});
it('should display reset time when provided', async () => {
const { lastFrame, unmount, waitUntilReady } = renderWithProviders(
<EmptyWalletDialog
failedModel="gemini-2.5-pro"
resetTime="3:45 PM"
onChoice={mockOnChoice}
/>,
);
await waitUntilReady();
const output = lastFrame() ?? '';
expect(output).toContain('3:45 PM');
expect(output).toContain('Access resets at');
unmount();
});
it('should not display reset time when not provided', async () => {
const { lastFrame, unmount, waitUntilReady } = renderWithProviders(
<EmptyWalletDialog
failedModel="gemini-2.5-pro"
onChoice={mockOnChoice}
/>,
);
await waitUntilReady();
const output = lastFrame() ?? '';
expect(output).not.toContain('Access resets at');
unmount();
});
it('should display slash command hints', async () => {
const { lastFrame, unmount, waitUntilReady } = renderWithProviders(
<EmptyWalletDialog
failedModel="gemini-2.5-pro"
onChoice={mockOnChoice}
/>,
);
await waitUntilReady();
const output = lastFrame() ?? '';
expect(output).toContain('/stats');
expect(output).toContain('/model');
expect(output).toContain('/auth');
unmount();
});
});
describe('onChoice handling', () => {
it('should call onGetCredits and onChoice when get_credits is selected', async () => {
// get_credits is the first item, so just press Enter
const { unmount, stdin, waitUntilReady } = renderWithProviders(
<EmptyWalletDialog
failedModel="gemini-2.5-pro"
onChoice={mockOnChoice}
onGetCredits={mockOnGetCredits}
/>,
);
await waitUntilReady();
writeKey(stdin, '\r');
await waitFor(() => {
expect(mockOnGetCredits).toHaveBeenCalled();
expect(mockOnChoice).toHaveBeenCalledWith('get_credits');
});
unmount();
});
it('should call onChoice without onGetCredits when onGetCredits is not provided', async () => {
const { unmount, stdin, waitUntilReady } = renderWithProviders(
<EmptyWalletDialog
failedModel="gemini-2.5-pro"
onChoice={mockOnChoice}
/>,
);
await waitUntilReady();
writeKey(stdin, '\r');
await waitFor(() => {
expect(mockOnChoice).toHaveBeenCalledWith('get_credits');
});
unmount();
});
it('should call onChoice with use_fallback when selected', async () => {
// With fallback: items are [get_credits, use_fallback, stop]
// use_fallback is the second item: Down + Enter
const { unmount, stdin, waitUntilReady } = renderWithProviders(
<EmptyWalletDialog
failedModel="gemini-2.5-pro"
fallbackModel="gemini-3-flash-preview"
onChoice={mockOnChoice}
/>,
);
await waitUntilReady();
writeKey(stdin, '\x1b[B'); // Down arrow
writeKey(stdin, '\r');
await waitFor(() => {
expect(mockOnChoice).toHaveBeenCalledWith('use_fallback');
});
unmount();
});
it('should call onChoice with stop when selected', async () => {
// Without fallback: items are [get_credits, stop]
// stop is the second item: Down + Enter
const { unmount, stdin, waitUntilReady } = renderWithProviders(
<EmptyWalletDialog
failedModel="gemini-2.5-pro"
onChoice={mockOnChoice}
/>,
);
await waitUntilReady();
writeKey(stdin, '\x1b[B'); // Down arrow
writeKey(stdin, '\r');
await waitFor(() => {
expect(mockOnChoice).toHaveBeenCalledWith('stop');
});
unmount();
});
});
});

View File

@@ -0,0 +1,110 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type React from 'react';
import { Box, Text } from 'ink';
import { RadioButtonSelect } from './shared/RadioButtonSelect.js';
import { theme } from '../semantic-colors.js';
/** Available choices in the empty wallet dialog */
export type EmptyWalletChoice = 'get_credits' | 'use_fallback' | 'stop';
interface EmptyWalletDialogProps {
/** The model that hit the quota limit */
failedModel: string;
/** The fallback model to offer (omit if none available) */
fallbackModel?: string;
/** Time when access resets (human-readable) */
resetTime?: string;
/** Callback to log click and open the browser for purchasing credits */
onGetCredits?: () => void;
/** Callback when user makes a selection */
onChoice: (choice: EmptyWalletChoice) => void;
}
export function EmptyWalletDialog({
failedModel,
fallbackModel,
resetTime,
onGetCredits,
onChoice,
}: EmptyWalletDialogProps): React.JSX.Element {
const items: Array<{
label: string;
value: EmptyWalletChoice;
key: string;
}> = [
{
label: 'Get AI Credits - Open browser to purchase credits',
value: 'get_credits',
key: 'get_credits',
},
];
if (fallbackModel) {
items.push({
label: `Switch to ${fallbackModel}`,
value: 'use_fallback',
key: 'use_fallback',
});
}
items.push({
label: 'Stop - Abort request',
value: 'stop',
key: 'stop',
});
const handleSelect = (choice: EmptyWalletChoice) => {
if (choice === 'get_credits') {
onGetCredits?.();
}
onChoice(choice);
};
return (
<Box borderStyle="round" flexDirection="column" padding={1}>
<Box marginBottom={1} flexDirection="column">
<Text color={theme.status.warning}>
Usage limit reached for {failedModel}.
</Text>
{resetTime && <Text>Access resets at {resetTime}.</Text>}
<Text>
<Text bold color={theme.text.accent}>
/stats
</Text>{' '}
model for usage details
</Text>
<Text>
<Text bold color={theme.text.accent}>
/model
</Text>{' '}
to switch models.
</Text>
<Text>
<Text bold color={theme.text.accent}>
/auth
</Text>{' '}
to switch to API key.
</Text>
</Box>
<Box marginBottom={1}>
<Text>To continue using this model now, purchase more AI Credits.</Text>
</Box>
<Box marginBottom={1}>
<Text dimColor>
Newly purchased AI credits may take a few minutes to update.
</Text>
</Box>
<Box marginBottom={1}>
<Text>How would you like to proceed?</Text>
</Box>
<Box marginTop={1} marginBottom={1}>
<RadioButtonSelect items={items} onSelect={handleSelect} />
</Box>
</Box>
);
}

View File

@@ -177,6 +177,8 @@ describe('<Footer />', () => {
},
proQuotaRequest: null,
validationRequest: null,
overageMenuRequest: null,
emptyWalletRequest: null,
},
},
},
@@ -203,6 +205,8 @@ describe('<Footer />', () => {
},
proQuotaRequest: null,
validationRequest: null,
overageMenuRequest: null,
emptyWalletRequest: null,
},
},
},
@@ -229,6 +233,8 @@ describe('<Footer />', () => {
},
proQuotaRequest: null,
validationRequest: null,
overageMenuRequest: null,
emptyWalletRequest: null,
},
},
},

View File

@@ -146,6 +146,7 @@ export const HistoryItemDisplay: React.FC<HistoryItemDisplayProps> = ({
}
: undefined
}
creditBalance={itemForDisplay.creditBalance}
/>
)}
{itemForDisplay.type === 'model_stats' && (

View File

@@ -0,0 +1,228 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { renderWithProviders } from '../../test-utils/render.js';
import { waitFor } from '../../test-utils/async.js';
import { act } from 'react';
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
import { OverageMenuDialog } from './OverageMenuDialog.js';
const writeKey = (stdin: { write: (data: string) => void }, key: string) => {
act(() => {
stdin.write(key);
});
};
describe('OverageMenuDialog', () => {
const mockOnChoice = vi.fn();
beforeEach(() => {
vi.clearAllMocks();
});
afterEach(() => {
vi.restoreAllMocks();
});
describe('rendering', () => {
it('should match snapshot with fallback available', async () => {
const { lastFrame, unmount, waitUntilReady } = renderWithProviders(
<OverageMenuDialog
failedModel="gemini-2.5-pro"
fallbackModel="gemini-3-flash-preview"
resetTime="2:00 PM"
creditBalance={500}
onChoice={mockOnChoice}
/>,
);
await waitUntilReady();
expect(lastFrame()).toMatchSnapshot();
unmount();
});
it('should match snapshot without fallback', async () => {
const { lastFrame, unmount, waitUntilReady } = renderWithProviders(
<OverageMenuDialog
failedModel="gemini-2.5-pro"
creditBalance={500}
onChoice={mockOnChoice}
/>,
);
await waitUntilReady();
expect(lastFrame()).toMatchSnapshot();
unmount();
});
it('should display the credit balance', async () => {
const { lastFrame, unmount, waitUntilReady } = renderWithProviders(
<OverageMenuDialog
failedModel="gemini-2.5-pro"
creditBalance={200}
onChoice={mockOnChoice}
/>,
);
await waitUntilReady();
const output = lastFrame() ?? '';
expect(output).toContain('200');
expect(output).toContain('AI Credits available');
unmount();
});
it('should display the model name', async () => {
const { lastFrame, unmount, waitUntilReady } = renderWithProviders(
<OverageMenuDialog
failedModel="gemini-2.5-pro"
creditBalance={100}
onChoice={mockOnChoice}
/>,
);
await waitUntilReady();
const output = lastFrame() ?? '';
expect(output).toContain('gemini-2.5-pro');
expect(output).toContain('Usage limit reached');
unmount();
});
it('should display reset time when provided', async () => {
const { lastFrame, unmount, waitUntilReady } = renderWithProviders(
<OverageMenuDialog
failedModel="gemini-2.5-pro"
resetTime="3:45 PM"
creditBalance={100}
onChoice={mockOnChoice}
/>,
);
await waitUntilReady();
const output = lastFrame() ?? '';
expect(output).toContain('3:45 PM');
expect(output).toContain('Access resets at');
unmount();
});
it('should not display reset time when not provided', async () => {
const { lastFrame, unmount, waitUntilReady } = renderWithProviders(
<OverageMenuDialog
failedModel="gemini-2.5-pro"
creditBalance={100}
onChoice={mockOnChoice}
/>,
);
await waitUntilReady();
const output = lastFrame() ?? '';
expect(output).not.toContain('Access resets at');
unmount();
});
it('should display slash command hints', async () => {
const { lastFrame, unmount, waitUntilReady } = renderWithProviders(
<OverageMenuDialog
failedModel="gemini-2.5-pro"
creditBalance={100}
onChoice={mockOnChoice}
/>,
);
await waitUntilReady();
const output = lastFrame() ?? '';
expect(output).toContain('/stats');
expect(output).toContain('/model');
expect(output).toContain('/auth');
unmount();
});
});
describe('onChoice handling', () => {
it('should call onChoice with use_credits when selected', async () => {
// use_credits is the first item, so just press Enter
const { unmount, stdin, waitUntilReady } = renderWithProviders(
<OverageMenuDialog
failedModel="gemini-2.5-pro"
creditBalance={100}
onChoice={mockOnChoice}
/>,
);
await waitUntilReady();
writeKey(stdin, '\r');
await waitFor(() => {
expect(mockOnChoice).toHaveBeenCalledWith('use_credits');
});
unmount();
});
it('should call onChoice with manage when selected', async () => {
// manage is the second item: Down + Enter
const { unmount, stdin, waitUntilReady } = renderWithProviders(
<OverageMenuDialog
failedModel="gemini-2.5-pro"
creditBalance={100}
onChoice={mockOnChoice}
/>,
);
await waitUntilReady();
writeKey(stdin, '\x1b[B'); // Down arrow
writeKey(stdin, '\r');
await waitFor(() => {
expect(mockOnChoice).toHaveBeenCalledWith('manage');
});
unmount();
});
it('should call onChoice with use_fallback when selected', async () => {
// With fallback: items are [use_credits, manage, use_fallback, stop]
// use_fallback is the third item: Down x2 + Enter
const { unmount, stdin, waitUntilReady } = renderWithProviders(
<OverageMenuDialog
failedModel="gemini-2.5-pro"
fallbackModel="gemini-3-flash-preview"
creditBalance={100}
onChoice={mockOnChoice}
/>,
);
await waitUntilReady();
writeKey(stdin, '\x1b[B'); // Down arrow
writeKey(stdin, '\x1b[B'); // Down arrow
writeKey(stdin, '\r');
await waitFor(() => {
expect(mockOnChoice).toHaveBeenCalledWith('use_fallback');
});
unmount();
});
it('should call onChoice with stop when selected', async () => {
// Without fallback: items are [use_credits, manage, stop]
// stop is the third item: Down x2 + Enter
const { unmount, stdin, waitUntilReady } = renderWithProviders(
<OverageMenuDialog
failedModel="gemini-2.5-pro"
creditBalance={100}
onChoice={mockOnChoice}
/>,
);
await waitUntilReady();
writeKey(stdin, '\x1b[B'); // Down arrow
writeKey(stdin, '\x1b[B'); // Down arrow
writeKey(stdin, '\r');
await waitFor(() => {
expect(mockOnChoice).toHaveBeenCalledWith('stop');
});
unmount();
});
});
});

View File

@@ -0,0 +1,113 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type React from 'react';
import { Box, Text } from 'ink';
import { RadioButtonSelect } from './shared/RadioButtonSelect.js';
import { theme } from '../semantic-colors.js';
/** Available choices in the overage menu dialog */
export type OverageMenuChoice =
| 'use_credits'
| 'use_fallback'
| 'manage'
| 'stop';
interface OverageMenuDialogProps {
/** The model that hit the quota limit */
failedModel: string;
/** The fallback model to offer (omit if none available) */
fallbackModel?: string;
/** Time when access resets (human-readable) */
resetTime?: string;
/** Available G1 AI credit balance */
creditBalance: number;
/** Callback when user makes a selection */
onChoice: (choice: OverageMenuChoice) => void;
}
export function OverageMenuDialog({
failedModel,
fallbackModel,
resetTime,
creditBalance,
onChoice,
}: OverageMenuDialogProps): React.JSX.Element {
const items: Array<{
label: string;
value: OverageMenuChoice;
key: string;
}> = [
{
label: 'Use AI Credits - Continue this request (Overage)',
value: 'use_credits',
key: 'use_credits',
},
{
label: 'Manage - View balance and purchase more credits',
value: 'manage',
key: 'manage',
},
];
if (fallbackModel) {
items.push({
label: `Switch to ${fallbackModel}`,
value: 'use_fallback',
key: 'use_fallback',
});
}
items.push({
label: 'Stop - Abort request',
value: 'stop',
key: 'stop',
});
return (
<Box borderStyle="round" flexDirection="column" padding={1}>
<Box marginBottom={1} flexDirection="column">
<Text color={theme.status.warning}>
Usage limit reached for {failedModel}.
</Text>
{resetTime && <Text>Access resets at {resetTime}.</Text>}
<Text>
<Text bold color={theme.text.accent}>
/stats
</Text>{' '}
model for usage details
</Text>
<Text>
<Text bold color={theme.text.accent}>
/model
</Text>{' '}
to switch models.
</Text>
<Text>
<Text bold color={theme.text.accent}>
/auth
</Text>{' '}
to switch to API key.
</Text>
</Box>
<Box marginBottom={1}>
<Text>
You have{' '}
<Text bold color={theme.status.success}>
{creditBalance}
</Text>{' '}
AI Credits available.
</Text>
</Box>
<Box marginBottom={1}>
<Text>How would you like to proceed?</Text>
</Box>
<Box marginTop={1} marginBottom={1}>
<RadioButtonSelect items={items} onSelect={onChoice} />
</Box>
</Box>
);
}

View File

@@ -395,6 +395,7 @@ interface StatsDisplayProps {
tier?: string;
currentModel?: string;
quotaStats?: QuotaStats;
creditBalance?: number;
}
export const StatsDisplay: React.FC<StatsDisplayProps> = ({
@@ -407,6 +408,7 @@ export const StatsDisplay: React.FC<StatsDisplayProps> = ({
tier,
currentModel,
quotaStats,
creditBalance,
}) => {
const { stats } = useSessionStats();
const { metrics } = stats;
@@ -488,6 +490,17 @@ export const StatsDisplay: React.FC<StatsDisplayProps> = ({
<Text color={theme.text.primary}>{tier}</Text>
</StatRow>
)}
{showUserIdentity && creditBalance != null && creditBalance >= 0 && (
<StatRow title="Google AI Credits:">
<Text
color={
creditBalance > 0 ? theme.text.primary : theme.text.secondary
}
>
{creditBalance.toLocaleString()}
</Text>
</StatRow>
)}
<StatRow title="Tool Calls:">
<Text color={theme.text.primary}>
{tools.totalCalls} ({' '}

View File

@@ -0,0 +1,49 @@
// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
exports[`EmptyWalletDialog > rendering > should match snapshot with fallback available 1`] = `
"╭──────────────────────────────────────────────────────────────────────────────────────────────────╮
│ │
│ Usage limit reached for gemini-2.5-pro. │
│ Access resets at 2:00 PM. │
│ /stats model for usage details │
│ /model to switch models. │
│ /auth to switch to API key. │
│ │
│ To continue using this model now, purchase more AI Credits. │
│ │
│ Newly purchased AI credits may take a few minutes to update. │
│ │
│ How would you like to proceed? │
│ │
│ │
│ ● 1. Get AI Credits - Open browser to purchase credits │
│ 2. Switch to gemini-3-flash-preview │
│ 3. Stop - Abort request │
│ │
│ │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
"
`;
exports[`EmptyWalletDialog > rendering > should match snapshot without fallback 1`] = `
"╭──────────────────────────────────────────────────────────────────────────────────────────────────╮
│ │
│ Usage limit reached for gemini-2.5-pro. │
│ /stats model for usage details │
│ /model to switch models. │
│ /auth to switch to API key. │
│ │
│ To continue using this model now, purchase more AI Credits. │
│ │
│ Newly purchased AI credits may take a few minutes to update. │
│ │
│ How would you like to proceed? │
│ │
│ │
│ ● 1. Get AI Credits - Open browser to purchase credits │
│ 2. Stop - Abort request │
│ │
│ │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
"
`;

View File

@@ -0,0 +1,47 @@
// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
exports[`OverageMenuDialog > rendering > should match snapshot with fallback available 1`] = `
"╭──────────────────────────────────────────────────────────────────────────────────────────────────╮
│ │
│ Usage limit reached for gemini-2.5-pro. │
│ Access resets at 2:00 PM. │
│ /stats model for usage details │
│ /model to switch models. │
│ /auth to switch to API key. │
│ │
│ You have 500 AI Credits available. │
│ │
│ How would you like to proceed? │
│ │
│ │
│ ● 1. Use AI Credits - Continue this request (Overage) │
│ 2. Manage - View balance and purchase more credits │
│ 3. Switch to gemini-3-flash-preview │
│ 4. Stop - Abort request │
│ │
│ │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
"
`;
exports[`OverageMenuDialog > rendering > should match snapshot without fallback 1`] = `
"╭──────────────────────────────────────────────────────────────────────────────────────────────────╮
│ │
│ Usage limit reached for gemini-2.5-pro. │
│ /stats model for usage details │
│ /model to switch models. │
│ /auth to switch to API key. │
│ │
│ You have 500 AI Credits available. │
│ │
│ How would you like to proceed? │
│ │
│ │
│ ● 1. Use AI Credits - Continue this request (Overage) │
│ 2. Manage - View balance and purchase more credits │
│ 3. Stop - Abort request │
│ │
│ │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
"
`;

View File

@@ -18,6 +18,7 @@ import type { AuthState } from '../types.js';
import { type PermissionsDialogProps } from '../components/PermissionsModifyTrustDialog.js';
import type { SessionInfo } from '../../utils/sessionUtils.js';
import { type NewAgentsChoice } from '../components/NewAgentsNotification.js';
import type { OverageMenuIntent, EmptyWalletIntent } from './UIStateContext.js';
export interface UIActions {
handleThemeSelect: (
@@ -62,6 +63,8 @@ export interface UIActions {
choice: 'retry_later' | 'retry_once' | 'retry_always' | 'upgrade',
) => void;
handleValidationChoice: (choice: 'verify' | 'change_auth' | 'cancel') => void;
handleOverageMenuChoice: (choice: OverageMenuIntent) => void;
handleEmptyWalletChoice: (choice: EmptyWalletIntent) => void;
openSessionBrowser: () => void;
closeSessionBrowser: () => void;
handleResumeSession: (session: SessionInfo) => Promise<void>;

View File

@@ -54,6 +54,34 @@ export interface ValidationDialogRequest {
resolve: (intent: ValidationIntent) => void;
}
/** Intent for overage menu dialog */
export type OverageMenuIntent =
| 'use_credits'
| 'use_fallback'
| 'manage'
| 'stop';
export interface OverageMenuDialogRequest {
failedModel: string;
fallbackModel?: string;
resetTime?: string;
creditBalance: number;
userEmail?: string;
resolve: (intent: OverageMenuIntent) => void;
}
/** Intent for empty wallet dialog */
export type EmptyWalletIntent = 'get_credits' | 'use_fallback' | 'stop';
export interface EmptyWalletDialogRequest {
failedModel: string;
fallbackModel?: string;
resetTime?: string;
userEmail?: string;
onGetCredits: () => void;
resolve: (intent: EmptyWalletIntent) => void;
}
import { type UseHistoryManagerReturn } from '../hooks/useHistoryManager.js';
import { type RestartReason } from '../hooks/useIdeTrustListener.js';
import type { TerminalBackgroundColor } from '../utils/terminalCapabilityManager.js';
@@ -64,6 +92,9 @@ export interface QuotaState {
stats: QuotaStats | undefined;
proQuotaRequest: ProQuotaDialogRequest | null;
validationRequest: ValidationDialogRequest | null;
// G1 AI Credits overage flow
overageMenuRequest: OverageMenuDialogRequest | null;
emptyWalletRequest: EmptyWalletDialogRequest | null;
}
export interface UIState {

View File

@@ -0,0 +1,240 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { vi, describe, it, expect, beforeEach, afterEach } from 'vitest';
import { handleCreditsFlow } from './creditsFlowHandler.js';
import type { UseHistoryManagerReturn } from './useHistoryManager.js';
import {
type Config,
type GeminiUserTier,
makeFakeConfig,
getG1CreditBalance,
shouldAutoUseCredits,
shouldShowOverageMenu,
shouldShowEmptyWalletMenu,
logBillingEvent,
G1_CREDIT_TYPE,
UserTierId,
} from '@google/gemini-cli-core';
import { MessageType } from '../types.js';
vi.mock('@google/gemini-cli-core', async (importOriginal) => {
const actual =
await importOriginal<typeof import('@google/gemini-cli-core')>();
return {
...actual,
getG1CreditBalance: vi.fn(),
shouldAutoUseCredits: vi.fn(),
shouldShowOverageMenu: vi.fn(),
shouldShowEmptyWalletMenu: vi.fn(),
logBillingEvent: vi.fn(),
openBrowserSecurely: vi.fn(),
};
});
describe('handleCreditsFlow', () => {
let mockConfig: Config;
let mockHistoryManager: UseHistoryManagerReturn;
let isDialogPending: React.MutableRefObject<boolean>;
let mockSetOverageMenuRequest: ReturnType<typeof vi.fn>;
let mockSetEmptyWalletRequest: ReturnType<typeof vi.fn>;
let mockSetModelSwitchedFromQuotaError: ReturnType<typeof vi.fn>;
const mockPaidTier: GeminiUserTier = {
id: UserTierId.STANDARD,
availableCredits: [{ creditType: G1_CREDIT_TYPE, creditAmount: '100' }],
};
beforeEach(() => {
mockConfig = makeFakeConfig();
mockHistoryManager = {
addItem: vi.fn(),
history: [],
updateItem: vi.fn(),
clearItems: vi.fn(),
loadHistory: vi.fn(),
};
isDialogPending = { current: false };
mockSetOverageMenuRequest = vi.fn();
mockSetEmptyWalletRequest = vi.fn();
mockSetModelSwitchedFromQuotaError = vi.fn();
vi.spyOn(mockConfig, 'setQuotaErrorOccurred');
vi.spyOn(mockConfig, 'setOverageStrategy');
vi.mocked(getG1CreditBalance).mockReturnValue(100);
vi.mocked(shouldAutoUseCredits).mockReturnValue(false);
vi.mocked(shouldShowOverageMenu).mockReturnValue(false);
vi.mocked(shouldShowEmptyWalletMenu).mockReturnValue(false);
});
afterEach(() => {
vi.restoreAllMocks();
});
function makeArgs(
overrides?: Partial<Parameters<typeof handleCreditsFlow>[0]>,
) {
return {
config: mockConfig,
paidTier: mockPaidTier,
overageStrategy: 'ask' as const,
failedModel: 'gemini-3-pro-preview',
fallbackModel: 'gemini-3-flash-preview',
usageLimitReachedModel: 'all Pro models',
resetTime: '3:45 PM',
historyManager: mockHistoryManager,
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
isDialogPending,
setOverageMenuRequest: mockSetOverageMenuRequest,
setEmptyWalletRequest: mockSetEmptyWalletRequest,
...overrides,
};
}
it('should return null if credit balance is null (non-G1 user)', async () => {
vi.mocked(getG1CreditBalance).mockReturnValue(null);
const result = await handleCreditsFlow(makeArgs());
expect(result).toBeNull();
});
it('should return null if credits are already auto-used (strategy=always)', async () => {
vi.mocked(shouldAutoUseCredits).mockReturnValue(true);
const result = await handleCreditsFlow(makeArgs());
expect(result).toBeNull();
});
it('should show overage menu and return retry_with_credits when use_credits selected', async () => {
vi.mocked(shouldShowOverageMenu).mockReturnValue(true);
const flowPromise = handleCreditsFlow(makeArgs());
// Extract the resolve callback from the setOverageMenuRequest call
expect(mockSetOverageMenuRequest).toHaveBeenCalledOnce();
const request = mockSetOverageMenuRequest.mock.calls[0][0];
expect(request.failedModel).toBe('all Pro models');
expect(request.creditBalance).toBe(100);
// Simulate user choosing 'use_credits'
request.resolve('use_credits');
const result = await flowPromise;
expect(result).toBe('retry_with_credits');
expect(mockConfig.setOverageStrategy).toHaveBeenCalledWith('always');
expect(logBillingEvent).toHaveBeenCalled();
});
it('should show overage menu and return retry_always when use_fallback selected', async () => {
vi.mocked(shouldShowOverageMenu).mockReturnValue(true);
const flowPromise = handleCreditsFlow(makeArgs());
const request = mockSetOverageMenuRequest.mock.calls[0][0];
request.resolve('use_fallback');
const result = await flowPromise;
expect(result).toBe('retry_always');
});
it('should show overage menu and return stop when stop selected', async () => {
vi.mocked(shouldShowOverageMenu).mockReturnValue(true);
const flowPromise = handleCreditsFlow(makeArgs());
const request = mockSetOverageMenuRequest.mock.calls[0][0];
request.resolve('stop');
const result = await flowPromise;
expect(result).toBe('stop');
});
it('should return stop immediately if dialog is already pending (overage)', async () => {
vi.mocked(shouldShowOverageMenu).mockReturnValue(true);
isDialogPending.current = true;
const result = await handleCreditsFlow(makeArgs());
expect(result).toBe('stop');
expect(mockSetOverageMenuRequest).not.toHaveBeenCalled();
});
it('should show empty wallet menu and return stop when get_credits selected', async () => {
vi.mocked(shouldShowEmptyWalletMenu).mockReturnValue(true);
const flowPromise = handleCreditsFlow(makeArgs());
expect(mockSetEmptyWalletRequest).toHaveBeenCalledOnce();
const request = mockSetEmptyWalletRequest.mock.calls[0][0];
expect(request.failedModel).toBe('all Pro models');
request.resolve('get_credits');
const result = await flowPromise;
expect(result).toBe('stop');
expect(mockHistoryManager.addItem).toHaveBeenCalledWith(
expect.objectContaining({
type: MessageType.INFO,
text: expect.stringContaining('few minutes'),
}),
expect.any(Number),
);
});
it('should show empty wallet menu and return retry_always when use_fallback selected', async () => {
vi.mocked(shouldShowEmptyWalletMenu).mockReturnValue(true);
const flowPromise = handleCreditsFlow(makeArgs());
const request = mockSetEmptyWalletRequest.mock.calls[0][0];
request.resolve('use_fallback');
const result = await flowPromise;
expect(result).toBe('retry_always');
});
it('should return stop immediately if dialog is already pending (empty wallet)', async () => {
vi.mocked(shouldShowEmptyWalletMenu).mockReturnValue(true);
isDialogPending.current = true;
const result = await handleCreditsFlow(makeArgs());
expect(result).toBe('stop');
expect(mockSetEmptyWalletRequest).not.toHaveBeenCalled();
});
it('should return null if no flow conditions are met', async () => {
vi.mocked(getG1CreditBalance).mockReturnValue(100);
vi.mocked(shouldAutoUseCredits).mockReturnValue(false);
vi.mocked(shouldShowOverageMenu).mockReturnValue(false);
vi.mocked(shouldShowEmptyWalletMenu).mockReturnValue(false);
const result = await handleCreditsFlow(makeArgs());
expect(result).toBeNull();
});
it('should clear dialog state after overage menu resolves', async () => {
vi.mocked(shouldShowOverageMenu).mockReturnValue(true);
const flowPromise = handleCreditsFlow(makeArgs());
expect(isDialogPending.current).toBe(true);
const request = mockSetOverageMenuRequest.mock.calls[0][0];
request.resolve('stop');
await flowPromise;
expect(isDialogPending.current).toBe(false);
// Verify null was set to clear the request
expect(mockSetOverageMenuRequest).toHaveBeenCalledWith(null);
});
it('should clear dialog state after empty wallet menu resolves', async () => {
vi.mocked(shouldShowEmptyWalletMenu).mockReturnValue(true);
const flowPromise = handleCreditsFlow(makeArgs());
expect(isDialogPending.current).toBe(true);
const request = mockSetEmptyWalletRequest.mock.calls[0][0];
request.resolve('stop');
await flowPromise;
expect(isDialogPending.current).toBe(false);
expect(mockSetEmptyWalletRequest).toHaveBeenCalledWith(null);
});
});

View File

@@ -0,0 +1,290 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import {
type Config,
type FallbackIntent,
type GeminiUserTier,
type OverageOption,
getG1CreditBalance,
shouldAutoUseCredits,
shouldShowOverageMenu,
shouldShowEmptyWalletMenu,
openBrowserSecurely,
logBillingEvent,
OverageMenuShownEvent,
OverageOptionSelectedEvent,
EmptyWalletMenuShownEvent,
CreditPurchaseClickEvent,
buildG1Url,
G1_UTM_CAMPAIGNS,
UserAccountManager,
recordOverageOptionSelected,
recordCreditPurchaseClick,
} from '@google/gemini-cli-core';
import { MessageType } from '../types.js';
import type { UseHistoryManagerReturn } from './useHistoryManager.js';
import type {
OverageMenuIntent,
EmptyWalletIntent,
EmptyWalletDialogRequest,
} from '../contexts/UIStateContext.js';
interface CreditsFlowArgs {
config: Config;
paidTier: GeminiUserTier;
overageStrategy: 'ask' | 'always' | 'never';
failedModel: string;
fallbackModel: string;
usageLimitReachedModel: string;
resetTime: string | undefined;
historyManager: UseHistoryManagerReturn;
setModelSwitchedFromQuotaError: (value: boolean) => void;
isDialogPending: React.MutableRefObject<boolean>;
setOverageMenuRequest: (
req: {
failedModel: string;
fallbackModel: string;
resetTime: string | undefined;
creditBalance: number;
resolve: (intent: OverageMenuIntent) => void;
} | null,
) => void;
setEmptyWalletRequest: (req: EmptyWalletDialogRequest | null) => void;
}
/**
* Handles the G1 AI Credits flow when a quota error occurs.
* Returns a FallbackIntent if the credits flow handled the error,
* or null to fall through to the default ProQuotaDialog.
*/
export async function handleCreditsFlow(
args: CreditsFlowArgs,
): Promise<FallbackIntent | null> {
const creditBalance = getG1CreditBalance(args.paidTier);
// creditBalance is null when user is not eligible for G1 credits.
if (creditBalance == null) {
return null;
}
const { overageStrategy } = args;
// If credits are already auto-enabled (strategy='always'), the request
// that just failed already included enabledCreditTypes — credits didn't
// help. Fall through to ProQuotaDialog which offers the Flash downgrade.
if (shouldAutoUseCredits(overageStrategy, creditBalance)) {
return null;
}
// Show overage menu when strategy is 'ask' and credits > 0
if (shouldShowOverageMenu(overageStrategy, creditBalance)) {
return handleOverageMenu(args, creditBalance);
}
// Show empty wallet when credits === 0 and strategy isn't 'never'
if (shouldShowEmptyWalletMenu(overageStrategy, creditBalance)) {
return handleEmptyWalletMenu(args);
}
return null;
}
// ---------------------------------------------------------------------------
// Overage menu flow
// ---------------------------------------------------------------------------
async function handleOverageMenu(
args: CreditsFlowArgs,
creditBalance: number,
): Promise<FallbackIntent> {
const {
config,
fallbackModel,
usageLimitReachedModel,
overageStrategy,
resetTime,
isDialogPending,
setOverageMenuRequest,
setModelSwitchedFromQuotaError,
historyManager,
} = args;
logBillingEvent(
config,
new OverageMenuShownEvent(
usageLimitReachedModel,
creditBalance,
overageStrategy,
),
);
if (isDialogPending.current) {
return 'stop';
}
isDialogPending.current = true;
setModelSwitchedFromQuotaError(true);
config.setQuotaErrorOccurred(true);
const overageIntent = await new Promise<OverageMenuIntent>((resolve) => {
setOverageMenuRequest({
failedModel: usageLimitReachedModel,
fallbackModel,
resetTime,
creditBalance,
resolve,
});
});
setOverageMenuRequest(null);
isDialogPending.current = false;
logOverageOptionSelected(
config,
usageLimitReachedModel,
overageIntent,
creditBalance,
);
switch (overageIntent) {
case 'use_credits':
setModelSwitchedFromQuotaError(false);
config.setQuotaErrorOccurred(false);
config.setOverageStrategy('always');
historyManager.addItem(
{
type: MessageType.INFO,
text: `Using AI Credits for this request.`,
},
Date.now(),
);
return 'retry_with_credits';
case 'use_fallback':
return 'retry_always';
case 'manage':
logCreditPurchaseClick(config, 'manage', usageLimitReachedModel);
await openG1Url('activity', G1_UTM_CAMPAIGNS.MANAGE_ACTIVITY);
return 'stop';
case 'stop':
default:
return 'stop';
}
}
// ---------------------------------------------------------------------------
// Empty wallet flow
// ---------------------------------------------------------------------------
async function handleEmptyWalletMenu(
args: CreditsFlowArgs,
): Promise<FallbackIntent> {
const {
config,
fallbackModel,
usageLimitReachedModel,
resetTime,
isDialogPending,
setEmptyWalletRequest,
setModelSwitchedFromQuotaError,
} = args;
logBillingEvent(
config,
new EmptyWalletMenuShownEvent(usageLimitReachedModel),
);
if (isDialogPending.current) {
return 'stop';
}
isDialogPending.current = true;
setModelSwitchedFromQuotaError(true);
config.setQuotaErrorOccurred(true);
const emptyWalletIntent = await new Promise<EmptyWalletIntent>((resolve) => {
setEmptyWalletRequest({
failedModel: usageLimitReachedModel,
fallbackModel,
resetTime,
onGetCredits: () => {
logCreditPurchaseClick(
config,
'empty_wallet_menu',
usageLimitReachedModel,
);
void openG1Url('credits', G1_UTM_CAMPAIGNS.EMPTY_WALLET_ADD_CREDITS);
},
resolve,
});
});
setEmptyWalletRequest(null);
isDialogPending.current = false;
switch (emptyWalletIntent) {
case 'get_credits':
args.historyManager.addItem(
{
type: MessageType.INFO,
text: 'Newly purchased AI credits may take a few minutes to update. Run /stats to check your balance.',
},
Date.now(),
);
return 'stop';
case 'use_fallback':
return 'retry_always';
case 'stop':
default:
return 'stop';
}
}
// ---------------------------------------------------------------------------
// Telemetry helpers
// ---------------------------------------------------------------------------
function logOverageOptionSelected(
config: Config,
model: string,
option: OverageOption,
creditBalance: number,
): void {
logBillingEvent(
config,
new OverageOptionSelectedEvent(model, option, creditBalance),
);
recordOverageOptionSelected(config, {
selected_option: option,
model,
});
}
function logCreditPurchaseClick(
config: Config,
source: 'overage_menu' | 'empty_wallet_menu' | 'manage',
model: string,
): void {
logBillingEvent(config, new CreditPurchaseClickEvent(source, model));
recordCreditPurchaseClick(config, { source, model });
}
async function openG1Url(
path: 'activity' | 'credits',
campaign: string,
): Promise<void> {
try {
const userEmail = new UserAccountManager().getCachedGoogleAccount() ?? '';
await openBrowserSecurely(buildG1Url(path, userEmail, campaign));
} catch {
// Ignore browser open errors
}
}

View File

@@ -14,7 +14,8 @@ import {
type Mock,
} from 'vitest';
import { act } from 'react';
import { renderHook } from '../../test-utils/render.js';
import { renderHook, mockSettings } from '../../test-utils/render.js';
import { waitFor } from '../../test-utils/async.js';
import {
type Config,
type FallbackModelHandler,
@@ -29,6 +30,12 @@ import {
ModelNotFoundError,
DEFAULT_GEMINI_MODEL,
DEFAULT_GEMINI_FLASH_MODEL,
getG1CreditBalance,
shouldAutoUseCredits,
shouldShowOverageMenu,
shouldShowEmptyWalletMenu,
logBillingEvent,
G1_CREDIT_TYPE,
} from '@google/gemini-cli-core';
import { useQuotaAndFallback } from './useQuotaAndFallback.js';
import type { UseHistoryManagerReturn } from './useHistoryManager.js';
@@ -37,6 +44,19 @@ import { MessageType } from '../types.js';
// Use a type alias for SpyInstance as it's not directly exported
type SpyInstance = ReturnType<typeof vi.spyOn>;
vi.mock('@google/gemini-cli-core', async (importOriginal) => {
const actual =
await importOriginal<typeof import('@google/gemini-cli-core')>();
return {
...actual,
getG1CreditBalance: vi.fn(),
shouldAutoUseCredits: vi.fn(),
shouldShowOverageMenu: vi.fn(),
shouldShowEmptyWalletMenu: vi.fn(),
logBillingEvent: vi.fn(),
};
});
describe('useQuotaAndFallback', () => {
let mockConfig: Config;
let mockHistoryManager: UseHistoryManagerReturn;
@@ -74,10 +94,16 @@ describe('useQuotaAndFallback', () => {
vi.spyOn(mockConfig, 'setModel');
vi.spyOn(mockConfig, 'setActiveModel');
vi.spyOn(mockConfig, 'activateFallbackMode');
// Mock billing utility functions
vi.mocked(getG1CreditBalance).mockReturnValue(0);
vi.mocked(shouldAutoUseCredits).mockReturnValue(false);
vi.mocked(shouldShowOverageMenu).mockReturnValue(false);
vi.mocked(shouldShowEmptyWalletMenu).mockReturnValue(false);
});
afterEach(() => {
vi.clearAllMocks();
vi.restoreAllMocks();
});
it('should register a fallback handler on initialization', () => {
@@ -88,6 +114,8 @@ describe('useQuotaAndFallback', () => {
userTier: UserTierId.FREE,
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
onShowAuthSelection: mockOnShowAuthSelection,
paidTier: null,
settings: mockSettings,
}),
);
@@ -109,6 +137,8 @@ describe('useQuotaAndFallback', () => {
userTier: UserTierId.FREE,
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
onShowAuthSelection: mockOnShowAuthSelection,
paidTier: null,
settings: mockSettings,
}),
);
@@ -140,6 +170,8 @@ describe('useQuotaAndFallback', () => {
userTier: UserTierId.FREE,
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
onShowAuthSelection: mockOnShowAuthSelection,
paidTier: null,
settings: mockSettings,
}),
);
@@ -193,6 +225,8 @@ describe('useQuotaAndFallback', () => {
userTier: UserTierId.FREE,
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
onShowAuthSelection: mockOnShowAuthSelection,
paidTier: null,
settings: mockSettings,
}),
);
@@ -232,6 +266,8 @@ describe('useQuotaAndFallback', () => {
userTier: UserTierId.FREE,
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
onShowAuthSelection: mockOnShowAuthSelection,
paidTier: null,
settings: mockSettings,
}),
);
@@ -264,6 +300,8 @@ describe('useQuotaAndFallback', () => {
userTier: UserTierId.FREE,
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
onShowAuthSelection: mockOnShowAuthSelection,
paidTier: null,
settings: mockSettings,
}),
);
@@ -330,6 +368,8 @@ describe('useQuotaAndFallback', () => {
setModelSwitchedFromQuotaError:
mockSetModelSwitchedFromQuotaError,
onShowAuthSelection: mockOnShowAuthSelection,
paidTier: null,
settings: mockSettings,
}),
);
@@ -385,6 +425,8 @@ describe('useQuotaAndFallback', () => {
userTier: UserTierId.FREE,
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
onShowAuthSelection: mockOnShowAuthSelection,
paidTier: null,
settings: mockSettings,
}),
);
@@ -430,6 +472,8 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
userTier: UserTierId.FREE,
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
onShowAuthSelection: mockOnShowAuthSelection,
paidTier: null,
settings: mockSettings,
}),
);
@@ -464,6 +508,243 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
});
});
describe('G1 AI Credits Flow', () => {
const mockPaidTier = {
id: UserTierId.STANDARD,
userTier: UserTierId.STANDARD,
availableCredits: [
{
creditType: G1_CREDIT_TYPE,
creditAmount: '100',
},
],
};
beforeEach(() => {
// Default to having credits
vi.mocked(getG1CreditBalance).mockReturnValue(100);
});
it('should fall through to ProQuotaDialog if credits are already active (strategy=always)', async () => {
// If shouldAutoUseCredits is true, credits were already active on the
// failed request — they didn't help. Fall through to ProQuotaDialog
// so the user can downgrade to Flash instead of retrying infinitely.
vi.mocked(shouldAutoUseCredits).mockReturnValue(true);
const { result } = renderHook(() =>
useQuotaAndFallback({
config: mockConfig,
historyManager: mockHistoryManager,
userTier: UserTierId.STANDARD,
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
onShowAuthSelection: mockOnShowAuthSelection,
paidTier: mockPaidTier,
settings: mockSettings,
}),
);
const handler = setFallbackHandlerSpy.mock
.calls[0][0] as FallbackModelHandler;
const error = new TerminalQuotaError(
'pro quota',
mockGoogleApiError,
1000 * 60 * 5,
);
const intentPromise = handler(
PREVIEW_GEMINI_MODEL,
'gemini-flash',
error,
);
// Since credits didn't help, the ProQuotaDialog should be shown
await waitFor(() => {
expect(result.current.proQuotaRequest).not.toBeNull();
});
// Resolve it to verify the flow completes
act(() => {
result.current.handleProQuotaChoice('stop');
});
const intent = await intentPromise;
expect(intent).toBe('stop');
});
it('should show overage menu if balance > 0 and not auto-using', async () => {
vi.mocked(shouldAutoUseCredits).mockReturnValue(false);
vi.mocked(shouldShowOverageMenu).mockReturnValue(true);
const { result } = renderHook(() =>
useQuotaAndFallback({
config: mockConfig,
historyManager: mockHistoryManager,
userTier: UserTierId.STANDARD,
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
onShowAuthSelection: mockOnShowAuthSelection,
paidTier: mockPaidTier,
settings: mockSettings,
}),
);
const handler = setFallbackHandlerSpy.mock
.calls[0][0] as FallbackModelHandler;
let promise: Promise<FallbackIntent | null>;
act(() => {
promise = handler(
PREVIEW_GEMINI_MODEL,
'gemini-flash',
new TerminalQuotaError('pro quota', mockGoogleApiError),
);
});
expect(result.current.overageMenuRequest).not.toBeNull();
expect(result.current.overageMenuRequest?.creditBalance).toBe(100);
expect(logBillingEvent).toHaveBeenCalled();
// Simulate choosing "Use Credits"
await act(async () => {
result.current.handleOverageMenuChoice('use_credits');
await promise!;
});
const intent = await promise!;
expect(intent).toBe('retry_with_credits');
});
it('should handle use_fallback from overage menu', async () => {
vi.mocked(shouldAutoUseCredits).mockReturnValue(false);
vi.mocked(shouldShowOverageMenu).mockReturnValue(true);
const { result } = renderHook(() =>
useQuotaAndFallback({
config: mockConfig,
historyManager: mockHistoryManager,
userTier: UserTierId.STANDARD,
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
onShowAuthSelection: mockOnShowAuthSelection,
paidTier: mockPaidTier,
settings: mockSettings,
}),
);
const handler = setFallbackHandlerSpy.mock
.calls[0][0] as FallbackModelHandler;
let promise: Promise<FallbackIntent | null>;
act(() => {
promise = handler(
PREVIEW_GEMINI_MODEL,
'gemini-flash',
new TerminalQuotaError('pro quota', mockGoogleApiError),
);
});
// Simulate choosing "Switch to fallback"
await act(async () => {
result.current.handleOverageMenuChoice('use_fallback');
await promise!;
});
const intent = await promise!;
expect(intent).toBe('retry_always');
});
it('should show empty wallet menu if balance is 0', async () => {
vi.mocked(getG1CreditBalance).mockReturnValue(0);
vi.mocked(shouldAutoUseCredits).mockReturnValue(false);
vi.mocked(shouldShowOverageMenu).mockReturnValue(false);
vi.mocked(shouldShowEmptyWalletMenu).mockReturnValue(true);
const { result } = renderHook(() =>
useQuotaAndFallback({
config: mockConfig,
historyManager: mockHistoryManager,
userTier: UserTierId.STANDARD,
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
onShowAuthSelection: mockOnShowAuthSelection,
paidTier: { ...mockPaidTier, availableCredits: [] },
settings: mockSettings,
}),
);
const handler = setFallbackHandlerSpy.mock
.calls[0][0] as FallbackModelHandler;
let promise: Promise<FallbackIntent | null>;
act(() => {
promise = handler(
PREVIEW_GEMINI_MODEL,
'gemini-flash',
new TerminalQuotaError('pro quota', mockGoogleApiError),
);
});
expect(result.current.emptyWalletRequest).not.toBeNull();
expect(logBillingEvent).toHaveBeenCalled();
// Simulate choosing "Stop"
await act(async () => {
result.current.handleEmptyWalletChoice('stop');
await promise!;
});
const intent = await promise!;
expect(intent).toBe('stop');
});
it('should add info message to history when get_credits is selected', async () => {
vi.mocked(getG1CreditBalance).mockReturnValue(0);
vi.mocked(shouldAutoUseCredits).mockReturnValue(false);
vi.mocked(shouldShowOverageMenu).mockReturnValue(false);
vi.mocked(shouldShowEmptyWalletMenu).mockReturnValue(true);
const { result } = renderHook(() =>
useQuotaAndFallback({
config: mockConfig,
historyManager: mockHistoryManager,
userTier: UserTierId.STANDARD,
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
onShowAuthSelection: mockOnShowAuthSelection,
paidTier: { ...mockPaidTier, availableCredits: [] },
settings: mockSettings,
}),
);
const handler = setFallbackHandlerSpy.mock
.calls[0][0] as FallbackModelHandler;
let promise: Promise<FallbackIntent | null>;
act(() => {
promise = handler(
PREVIEW_GEMINI_MODEL,
'gemini-flash',
new TerminalQuotaError('pro quota', mockGoogleApiError),
);
});
expect(result.current.emptyWalletRequest).not.toBeNull();
// Simulate choosing "Get AI Credits"
await act(async () => {
result.current.handleEmptyWalletChoice('get_credits');
await promise!;
});
const intent = await promise!;
expect(intent).toBe('stop');
expect(mockHistoryManager.addItem).toHaveBeenCalledWith(
expect.objectContaining({
type: MessageType.INFO,
text: expect.stringContaining('few minutes'),
}),
expect.any(Number),
);
});
});
describe('handleProQuotaChoice', () => {
it('should do nothing if there is no pending pro quota request', () => {
const { result } = renderHook(() =>
@@ -473,6 +754,8 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
userTier: UserTierId.FREE,
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
onShowAuthSelection: mockOnShowAuthSelection,
paidTier: null,
settings: mockSettings,
}),
);
@@ -491,6 +774,8 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
userTier: UserTierId.FREE,
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
onShowAuthSelection: mockOnShowAuthSelection,
paidTier: null,
settings: mockSettings,
}),
);
@@ -522,6 +807,8 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
userTier: UserTierId.FREE,
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
onShowAuthSelection: mockOnShowAuthSelection,
paidTier: null,
settings: mockSettings,
}),
);
@@ -566,6 +853,8 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
userTier: UserTierId.FREE,
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
onShowAuthSelection: mockOnShowAuthSelection,
paidTier: null,
settings: mockSettings,
}),
);
@@ -602,6 +891,8 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
userTier: UserTierId.FREE,
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
onShowAuthSelection: mockOnShowAuthSelection,
paidTier: null,
settings: mockSettings,
}),
);
@@ -646,6 +937,8 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
userTier: UserTierId.FREE,
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
onShowAuthSelection: mockOnShowAuthSelection,
paidTier: null,
settings: mockSettings,
}),
);
@@ -661,6 +954,8 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
userTier: UserTierId.FREE,
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
onShowAuthSelection: mockOnShowAuthSelection,
paidTier: null,
settings: mockSettings,
}),
);
@@ -703,6 +998,8 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
userTier: UserTierId.FREE,
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
onShowAuthSelection: mockOnShowAuthSelection,
paidTier: null,
settings: mockSettings,
}),
);
@@ -745,6 +1042,8 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
userTier: UserTierId.FREE,
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
onShowAuthSelection: mockOnShowAuthSelection,
paidTier: null,
settings: mockSettings,
}),
);
@@ -775,6 +1074,8 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
userTier: UserTierId.FREE,
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
onShowAuthSelection: mockOnShowAuthSelection,
paidTier: null,
settings: mockSettings,
}),
);
@@ -805,6 +1106,8 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
userTier: UserTierId.FREE,
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
onShowAuthSelection: mockOnShowAuthSelection,
paidTier: null,
settings: mockSettings,
}),
);

View File

@@ -16,7 +16,9 @@ import {
type UserTierId,
VALID_GEMINI_MODELS,
isProModel,
isOverageEligibleModel,
getDisplayString,
type GeminiUserTier,
} from '@google/gemini-cli-core';
import { useCallback, useEffect, useRef, useState } from 'react';
import { type UseHistoryManagerReturn } from './useHistoryManager.js';
@@ -24,12 +26,20 @@ import { MessageType } from '../types.js';
import {
type ProQuotaDialogRequest,
type ValidationDialogRequest,
type OverageMenuDialogRequest,
type OverageMenuIntent,
type EmptyWalletDialogRequest,
type EmptyWalletIntent,
} from '../contexts/UIStateContext.js';
import type { LoadedSettings } from '../../config/settings.js';
import { handleCreditsFlow } from './creditsFlowHandler.js';
interface UseQuotaAndFallbackArgs {
config: Config;
historyManager: UseHistoryManagerReturn;
userTier: UserTierId | undefined;
paidTier: GeminiUserTier | null | undefined;
settings: LoadedSettings;
setModelSwitchedFromQuotaError: (value: boolean) => void;
onShowAuthSelection: () => void;
}
@@ -38,6 +48,8 @@ export function useQuotaAndFallback({
config,
historyManager,
userTier,
paidTier,
settings,
setModelSwitchedFromQuotaError,
onShowAuthSelection,
}: UseQuotaAndFallbackArgs) {
@@ -45,9 +57,22 @@ export function useQuotaAndFallback({
useState<ProQuotaDialogRequest | null>(null);
const [validationRequest, setValidationRequest] =
useState<ValidationDialogRequest | null>(null);
// G1 AI Credits dialog states
const [overageMenuRequest, setOverageMenuRequest] =
useState<OverageMenuDialogRequest | null>(null);
const [emptyWalletRequest, setEmptyWalletRequest] =
useState<EmptyWalletDialogRequest | null>(null);
const isDialogPending = useRef(false);
const isValidationPending = useRef(false);
// Initial overage strategy from settings; runtime value read from config at call time.
const initialOverageStrategy =
(settings.merged.billing?.overageStrategy as
| 'ask'
| 'always'
| 'never'
| undefined) ?? 'ask';
// Set up Flash fallback handler
useEffect(() => {
const fallbackHandler: FallbackModelHandler = async (
@@ -63,12 +88,52 @@ export function useQuotaAndFallback({
const usageLimitReachedModel = isProModel(failedModel)
? 'all Pro models'
: failedModel;
if (error instanceof TerminalQuotaError) {
isTerminalQuotaError = true;
// Common part of the message for both tiers
const isInsufficientCredits = error.isInsufficientCredits;
// G1 Credits Flow: Only apply if user has a tier that supports credits
// (paidTier?.availableCredits indicates the user is a G1 subscriber)
// Skip if the error explicitly says they have insufficient credits (e.g. they
// just exhausted them or zero balance cache is delayed).
if (
!isInsufficientCredits &&
paidTier?.availableCredits &&
isOverageEligibleModel(failedModel)
) {
const resetTime = error.retryDelayMs
? getResetTimeMessage(error.retryDelayMs)
: undefined;
const overageStrategy =
config.getBillingSettings().overageStrategy ??
initialOverageStrategy;
const creditsResult = await handleCreditsFlow({
config,
paidTier,
overageStrategy,
failedModel,
fallbackModel,
usageLimitReachedModel,
resetTime,
historyManager,
setModelSwitchedFromQuotaError,
isDialogPending,
setOverageMenuRequest,
setEmptyWalletRequest,
});
if (creditsResult) return creditsResult;
}
// Default: Show existing ProQuotaDialog (for overageStrategy: 'never' or non-G1 users)
const messageLines = [
`Usage limit reached for ${usageLimitReachedModel}.`,
error.retryDelayMs ? getResetTimeMessage(error.retryDelayMs) : null,
error.retryDelayMs
? `Access resets at ${getResetTimeMessage(error.retryDelayMs)}.`
: null,
`/stats model for usage details`,
`/model to switch models.`,
contentGeneratorConfig?.authType === AuthType.LOGIN_WITH_GOOGLE
@@ -126,7 +191,16 @@ export function useQuotaAndFallback({
};
config.setFallbackModelHandler(fallbackHandler);
}, [config, historyManager, userTier, setModelSwitchedFromQuotaError]);
}, [
config,
historyManager,
userTier,
paidTier,
settings,
initialOverageStrategy,
setModelSwitchedFromQuotaError,
onShowAuthSelection,
]);
// Set up validation handler for 403 VALIDATION_REQUIRED errors
useEffect(() => {
@@ -204,11 +278,38 @@ export function useQuotaAndFallback({
[validationRequest, onShowAuthSelection],
);
// Handler for overage menu dialog (G1 AI Credits flow)
const handleOverageMenuChoice = useCallback(
(choice: OverageMenuIntent) => {
if (!overageMenuRequest) return;
overageMenuRequest.resolve(choice);
// State will be cleared by the effect callback after the promise resolves
},
[overageMenuRequest],
);
// Handler for empty wallet dialog (G1 AI Credits flow)
const handleEmptyWalletChoice = useCallback(
(choice: EmptyWalletIntent) => {
if (!emptyWalletRequest) return;
emptyWalletRequest.resolve(choice);
// State will be cleared by the effect callback after the promise resolves
},
[emptyWalletRequest],
);
return {
proQuotaRequest,
handleProQuotaChoice,
validationRequest,
handleValidationChoice,
// G1 AI Credits
overageMenuRequest,
handleOverageMenuChoice,
emptyWalletRequest,
handleEmptyWalletChoice,
};
}
@@ -221,5 +322,5 @@ function getResetTimeMessage(delayMs: number): string {
timeZoneName: 'short',
});
return `Access resets at ${timeFormatter.format(resetDate)}.`;
return timeFormatter.format(resetDate);
}

View File

@@ -201,6 +201,7 @@ export type HistoryItemStats = HistoryItemQuotaBase & {
type: 'stats';
duration: string;
quotas?: RetrieveUserQuotaResponse;
creditBalance?: number;
};
export type HistoryItemModelStats = HistoryItemQuotaBase & {