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

feat: implement G1 AI credits overage flow with billing telemetry

Browse Source 
Adds end-to-end support for Google One AI credits in quota exhaustion flows:

- New billing module (packages/core/src/billing/) with credit balance
  checking, overage strategy management, and G1 URL construction
- OverageMenuDialog and EmptyWalletDialog UI components for quota
  exhaustion with credit purchase options
- Credits flow handler extracted to creditsFlowHandler.ts with overage
  menu, empty wallet, and auto-use-credits logic
- Server-side credit tracking: enabledCreditTypes on requests,
  consumed/remaining credits from streaming responses
- Billing telemetry events (overage menu shown, option selected, credits
  used, credit purchase click, API key updated)
- OpenTelemetry metrics for overage option and credit purchase counters
- Credit balance display in /stats command with refresh support
- Settings: general.overageStrategy (ask/always/never) for credit usage
- Error handling: INSUFFICIENT_G1_CREDITS_BALANCE as terminal error
  regardless of domain field presence
- Persistent info message after
This commit is contained in:
Gaurav Ghosh
2026-02-25 03:21:10 -08:00
parent 29e8f2abf4
commit 6836f0e1b2
53 changed files with 3159 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/cli/settings.md
+6
View File
@@ -80,6 +80,12 @@ they appear in the UI.
| -------- | ------------- | ---------------------------- | ------- |
| IDE Mode | `ide.enabled` | Enable IDE integration mode. | `false` |
### Billing
| UI Label | Setting | Description | Default |
| ---------------- | ------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------- | ------- |
| Overage Strategy | `billing.overageStrategy` | How to handle quota exhaustion when AI credits are available. 'ask' prompts each time, 'always' automatically uses credits, 'never' disables credit usage. | `"ask"` |
### Model
| UI Label | Setting | Description | Default |
docs/reference/configuration.md
+9
View File
@@ -357,6 +357,15 @@ their corresponding top-level category object in your `settings.json` file.
- **Default:** `true`
- **Requires restart:** Yes
#### `billing`
- **`billing.overageStrategy`** (enum):
- **Description:** How to handle quota exhaustion when AI credits are
available. 'ask' prompts each time, 'always' automatically uses credits,
'never' disables credit usage.
- **Default:** `"ask"`
- **Values:** `"ask"`, `"always"`, `"never"`
#### `model`
- **`model.name`** (string):
packages/cli/src/config/settingsSchema.ts
+30
View File
@@ -828,6 +828,36 @@ const SETTINGS_SCHEMA = {
ref: 'TelemetrySettings',
},
billing: {
type: 'object',
label: 'Billing',
category: 'Advanced',
requiresRestart: false,
default: {},
description: 'Billing and AI credits settings.',
showInDialog: true,
properties: {
overageStrategy: {
type: 'enum',
label: 'Overage Strategy',
category: 'Advanced',
requiresRestart: false,
default: 'ask',
description: oneLine`
How to handle quota exhaustion when AI credits are available.
'ask' prompts each time, 'always' automatically uses credits,
'never' disables credit usage.
`,
showInDialog: true,
options: [
{ value: 'ask', label: 'Ask each time' },
{ value: 'always', label: 'Always use credits' },
{ value: 'never', label: 'Never use credits' },
],
},
},
},
model: {
type: 'object',
label: 'Model',
packages/cli/src/test-utils/render.tsx
+2
View File
@@ -553,6 +553,8 @@ const mockUIActions: UIActions = {
handleClearScreen: vi.fn(),
handleProQuotaChoice: vi.fn(),
handleValidationChoice: vi.fn(),
handleOverageMenuChoice: vi.fn(),
handleEmptyWalletChoice: vi.fn(),
setQueueErrorMessage: vi.fn(),
popAllMessages: vi.fn(),
handleApiKeySubmit: vi.fn(),
packages/cli/src/ui/AppContainer.tsx
+34
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';
@@ -414,6 +417,9 @@ export const AppContainer = (props: AppContainerProps) => {
? { remaining, limit, resetTime }
: undefined;
});
const [paidTier, setPaidTier] = useState<GeminiUserTier | undefined>(
undefined,
);
const [isConfigInitialized, setConfigInitialized] = useState(false);
@@ -709,10 +715,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),
});
@@ -752,6 +765,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 {
@@ -764,6 +779,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;
@@ -826,6 +845,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]);
@@ -2056,6 +2076,8 @@ Logging in with Google... Restarting Gemini CLI to continue.
showIdeRestartPrompt ||
!!proQuotaRequest ||
!!validationRequest ||
!!overageMenuRequest ||
!!emptyWalletRequest ||
isSessionBrowserOpen ||
authState === AuthState.AwaitingApiKeyInput ||
!!newAgents;
@@ -2083,6 +2105,8 @@ Logging in with Google... Restarting Gemini CLI to continue.
hasLoopDetectionConfirmationRequest ||
!!proQuotaRequest ||
!!validationRequest ||
!!overageMenuRequest ||
!!emptyWalletRequest ||
!!customDialog;
const allowPlanMode =
@@ -2293,6 +2317,9 @@ Logging in with Google... Restarting Gemini CLI to continue.
stats: quotaStats,
proQuotaRequest,
validationRequest,
// G1 AI Credits dialog state
overageMenuRequest,
emptyWalletRequest,
},
contextFileNames,
errorCount,
@@ -2417,6 +2444,8 @@ Logging in with Google... Restarting Gemini CLI to continue.
quotaStats,
proQuotaRequest,
validationRequest,
overageMenuRequest,
emptyWalletRequest,
contextFileNames,
errorCount,
availableTerminalHeight,
@@ -2498,6 +2527,9 @@ Logging in with Google... Restarting Gemini CLI to continue.
handleClearScreen,
handleProQuotaChoice,
handleValidationChoice,
// G1 AI Credits handlers
handleOverageMenuChoice,
handleEmptyWalletChoice,
openSessionBrowser,
closeSessionBrowser,
handleResumeSession,
@@ -2583,6 +2615,8 @@ Logging in with Google... Restarting Gemini CLI to continue.
handleClearScreen,
handleProQuotaChoice,
handleValidationChoice,
handleOverageMenuChoice,
handleEmptyWalletChoice,
openSessionBrowser,
closeSessionBrowser,
handleResumeSession,
packages/cli/src/ui/commands/statsCommand.test.ts
+13 -3
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: null,
});
});
@@ -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, '');
packages/cli/src/ui/commands/statsCommand.ts
+14 -4
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);
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();
packages/cli/src/ui/components/DialogManager.test.tsx
+4
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',
packages/cli/src/ui/components/DialogManager.tsx
+24
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';
@@ -151,6 +153,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
packages/cli/src/ui/components/EmptyWalletDialog.test.tsx
+261
View File
@@ -0,0 +1,261 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { render } from '../../test-utils/render.js';
import { act } from 'react';
import {
describe,
it,
expect,
vi,
beforeEach,
afterEach,
type Mock,
} from 'vitest';
import { EmptyWalletDialog } from './EmptyWalletDialog.js';
import { RadioButtonSelect } from './shared/RadioButtonSelect.js';
// Mock the child component to make it easier to test the parent
vi.mock('./shared/RadioButtonSelect.js', () => ({
RadioButtonSelect: vi.fn(),
}));
describe('EmptyWalletDialog', () => {
const mockOnChoice = vi.fn();
const mockOnGetCredits = vi.fn();
beforeEach(() => {
vi.clearAllMocks();
});
afterEach(() => {
vi.restoreAllMocks();
});
describe('rendering', () => {
it('should render with correct menu options when fallback is available', async () => {
const { unmount, waitUntilReady } = render(
<EmptyWalletDialog
failedModel="gemini-2.5-pro"
fallbackModel="gemini-3-flash-preview"
resetTime="2:00 PM"
onChoice={mockOnChoice}
/>,
);
await waitUntilReady();
expect(RadioButtonSelect).toHaveBeenCalledWith(
expect.objectContaining({
items: [
{
label: 'Get AI Credits - Open browser to purchase credits',
value: 'get_credits',
key: 'get_credits',
},
{
label: 'Switch to gemini-3-flash-preview',
value: 'use_fallback',
key: 'use_fallback',
},
{
label: 'Stop - Abort request',
value: 'stop',
key: 'stop',
},
],
}),
undefined,
);
unmount();
});
it('should omit fallback option when fallbackModel is not provided', async () => {
const { unmount, waitUntilReady } = render(
<EmptyWalletDialog
failedModel="gemini-2.5-pro"
onChoice={mockOnChoice}
/>,
);
await waitUntilReady();
expect(RadioButtonSelect).toHaveBeenCalledWith(
expect.objectContaining({
items: [
{
label: 'Get AI Credits - Open browser to purchase credits',
value: 'get_credits',
key: 'get_credits',
},
{
label: 'Stop - Abort request',
value: 'stop',
key: 'stop',
},
],
}),
undefined,
);
unmount();
});
it('should display the model name and usage limit message', async () => {
const { lastFrame, unmount, waitUntilReady } = render(
<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 } = render(
<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 } = render(
<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 } = render(
<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 } = render(
<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 () => {
const { unmount, waitUntilReady } = render(
<EmptyWalletDialog
failedModel="gemini-2.5-pro"
onChoice={mockOnChoice}
onGetCredits={mockOnGetCredits}
/>,
);
await waitUntilReady();
const onSelect = (RadioButtonSelect as Mock).mock.calls[0][0].onSelect;
await act(async () => {
onSelect('get_credits');
});
expect(mockOnGetCredits).toHaveBeenCalled();
expect(mockOnChoice).toHaveBeenCalledWith('get_credits');
unmount();
});
it('should call onChoice without onGetCredits when onGetCredits is not provided', async () => {
const { unmount, waitUntilReady } = render(
<EmptyWalletDialog
failedModel="gemini-2.5-pro"
onChoice={mockOnChoice}
/>,
);
await waitUntilReady();
const onSelect = (RadioButtonSelect as Mock).mock.calls[0][0].onSelect;
await act(async () => {
onSelect('get_credits');
});
expect(mockOnChoice).toHaveBeenCalledWith('get_credits');
unmount();
});
it('should call onChoice with use_fallback when selected', async () => {
const { unmount, waitUntilReady } = render(
<EmptyWalletDialog
failedModel="gemini-2.5-pro"
fallbackModel="gemini-3-flash-preview"
onChoice={mockOnChoice}
/>,
);
await waitUntilReady();
const onSelect = (RadioButtonSelect as Mock).mock.calls[0][0].onSelect;
await act(async () => {
onSelect('use_fallback');
});
expect(mockOnChoice).toHaveBeenCalledWith('use_fallback');
unmount();
});
it('should call onChoice with stop when selected', async () => {
const { unmount, waitUntilReady } = render(
<EmptyWalletDialog
failedModel="gemini-2.5-pro"
onChoice={mockOnChoice}
/>,
);
await waitUntilReady();
const onSelect = (RadioButtonSelect as Mock).mock.calls[0][0].onSelect;
await act(async () => {
onSelect('stop');
});
expect(mockOnChoice).toHaveBeenCalledWith('stop');
unmount();
});
});
});
packages/cli/src/ui/components/EmptyWalletDialog.tsx
+110
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>
);
}
packages/cli/src/ui/components/Footer.test.tsx
+6
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,
},
},
},
packages/cli/src/ui/components/HistoryItemDisplay.tsx
+1
View File
@@ -146,6 +146,7 @@ export const HistoryItemDisplay: React.FC<HistoryItemDisplayProps> = ({
}
: undefined
}
creditBalance={itemForDisplay.creditBalance}
/>
)}
{itemForDisplay.type === 'model_stats' && (
packages/cli/src/ui/components/OverageMenuDialog.test.tsx
+263
View File
@@ -0,0 +1,263 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { render } from '../../test-utils/render.js';
import { act } from 'react';
import { describe, it, expect, vi, beforeEach, type Mock } from 'vitest';
import { OverageMenuDialog } from './OverageMenuDialog.js';
import { RadioButtonSelect } from './shared/RadioButtonSelect.js';
// Mock the child component to make it easier to test the parent
vi.mock('./shared/RadioButtonSelect.js', () => ({
RadioButtonSelect: vi.fn(),
}));
describe('OverageMenuDialog', () => {
const mockOnChoice = vi.fn();
beforeEach(() => {
vi.clearAllMocks();
});
describe('rendering', () => {
it('should render with correct menu options when fallback is available', async () => {
const { unmount, waitUntilReady } = render(
<OverageMenuDialog
failedModel="gemini-2.5-pro"
fallbackModel="gemini-3-flash-preview"
resetTime="2:00 PM"
creditBalance={500}
onChoice={mockOnChoice}
/>,
);
await waitUntilReady();
expect(RadioButtonSelect).toHaveBeenCalledWith(
expect.objectContaining({
items: [
{
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',
},
{
label: 'Switch to gemini-3-flash-preview',
value: 'use_fallback',
key: 'use_fallback',
},
{
label: 'Stop - Abort request',
value: 'stop',
key: 'stop',
},
],
}),
undefined,
);
unmount();
});
it('should omit fallback option when fallbackModel is not provided', async () => {
const { unmount, waitUntilReady } = render(
<OverageMenuDialog
failedModel="gemini-2.5-pro"
creditBalance={500}
onChoice={mockOnChoice}
/>,
);
await waitUntilReady();
expect(RadioButtonSelect).toHaveBeenCalledWith(
expect.objectContaining({
items: [
{
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',
},
{
label: 'Stop - Abort request',
value: 'stop',
key: 'stop',
},
],
}),
undefined,
);
unmount();
});
it('should display the credit balance', async () => {
const { lastFrame, unmount, waitUntilReady } = render(
<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 } = render(
<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 } = render(
<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 } = render(
<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 } = render(
<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 () => {
const { unmount, waitUntilReady } = render(
<OverageMenuDialog
failedModel="gemini-2.5-pro"
creditBalance={100}
onChoice={mockOnChoice}
/>,
);
await waitUntilReady();
const onSelect = (RadioButtonSelect as Mock).mock.calls[0][0].onSelect;
await act(async () => {
onSelect('use_credits');
});
expect(mockOnChoice).toHaveBeenCalledWith('use_credits');
unmount();
});
it('should call onChoice with manage when selected', async () => {
const { unmount, waitUntilReady } = render(
<OverageMenuDialog
failedModel="gemini-2.5-pro"
creditBalance={100}
onChoice={mockOnChoice}
/>,
);
await waitUntilReady();
const onSelect = (RadioButtonSelect as Mock).mock.calls[0][0].onSelect;
await act(async () => {
onSelect('manage');
});
expect(mockOnChoice).toHaveBeenCalledWith('manage');
unmount();
});
it('should call onChoice with use_fallback when selected', async () => {
const { unmount, waitUntilReady } = render(
<OverageMenuDialog
failedModel="gemini-2.5-pro"
fallbackModel="gemini-3-flash-preview"
creditBalance={100}
onChoice={mockOnChoice}
/>,
);
await waitUntilReady();
const onSelect = (RadioButtonSelect as Mock).mock.calls[0][0].onSelect;
await act(async () => {
onSelect('use_fallback');
});
expect(mockOnChoice).toHaveBeenCalledWith('use_fallback');
unmount();
});
it('should call onChoice with stop when selected', async () => {
const { unmount, waitUntilReady } = render(
<OverageMenuDialog
failedModel="gemini-2.5-pro"
creditBalance={100}
onChoice={mockOnChoice}
/>,
);
await waitUntilReady();
const onSelect = (RadioButtonSelect as Mock).mock.calls[0][0].onSelect;
await act(async () => {
onSelect('stop');
});
expect(mockOnChoice).toHaveBeenCalledWith('stop');
unmount();
});
});
});
packages/cli/src/ui/components/OverageMenuDialog.tsx
+113
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>
);
}
packages/cli/src/ui/components/StatsDisplay.tsx
+14
View File
@@ -395,6 +395,8 @@ interface StatsDisplayProps {
tier?: string;
currentModel?: string;
quotaStats?: QuotaStats;
/** G1 AI Credits balance, null if not eligible */
creditBalance?: number | null;
}
export const StatsDisplay: React.FC<StatsDisplayProps> = ({
@@ -407,6 +409,7 @@ export const StatsDisplay: React.FC<StatsDisplayProps> = ({
tier,
currentModel,
quotaStats,
creditBalance,
}) => {
const { stats } = useSessionStats();
const { metrics } = stats;
@@ -488,6 +491,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} ({' '}
packages/cli/src/ui/contexts/UIActionsContext.tsx
+3
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>;
packages/cli/src/ui/contexts/UIStateContext.tsx
+31
View File
@@ -52,6 +52,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';
@@ -62,6 +90,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 {
packages/cli/src/ui/hooks/creditsFlowHandler.test.ts
+240
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);
});
});
packages/cli/src/ui/hooks/creditsFlowHandler.ts
+290
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 credits — they 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
}
}
packages/cli/src/ui/hooks/useQuotaAndFallback.test.ts
+304 -1
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,6 +94,12 @@ 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(() => {
@@ -88,6 +114,8 @@ describe('useQuotaAndFallback', () => {
userTier: UserTierId.FREE,
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
onShowAuthSelection: mockOnShowAuthSelection,
paidTier: null,
settings: mockSettings,
}),
);
@@ -105,6 +133,8 @@ describe('useQuotaAndFallback', () => {
userTier: UserTierId.FREE,
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
onShowAuthSelection: mockOnShowAuthSelection,
paidTier: null,
settings: mockSettings,
}),
);
return setFallbackHandlerSpy.mock.calls[0][0] as FallbackModelHandler;
@@ -132,6 +162,8 @@ describe('useQuotaAndFallback', () => {
userTier: UserTierId.FREE,
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
onShowAuthSelection: mockOnShowAuthSelection,
paidTier: null,
settings: mockSettings,
}),
);
@@ -185,6 +217,8 @@ describe('useQuotaAndFallback', () => {
userTier: UserTierId.FREE,
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
onShowAuthSelection: mockOnShowAuthSelection,
paidTier: null,
settings: mockSettings,
}),
);
@@ -224,6 +258,8 @@ describe('useQuotaAndFallback', () => {
userTier: UserTierId.FREE,
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
onShowAuthSelection: mockOnShowAuthSelection,
paidTier: null,
settings: mockSettings,
}),
);
@@ -256,6 +292,8 @@ describe('useQuotaAndFallback', () => {
userTier: UserTierId.FREE,
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
onShowAuthSelection: mockOnShowAuthSelection,
paidTier: null,
settings: mockSettings,
}),
);
@@ -322,6 +360,8 @@ describe('useQuotaAndFallback', () => {
setModelSwitchedFromQuotaError:
mockSetModelSwitchedFromQuotaError,
onShowAuthSelection: mockOnShowAuthSelection,
paidTier: null,
settings: mockSettings,
}),
);
@@ -377,6 +417,8 @@ describe('useQuotaAndFallback', () => {
userTier: UserTierId.FREE,
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
onShowAuthSelection: mockOnShowAuthSelection,
paidTier: null,
settings: mockSettings,
}),
);
@@ -422,6 +464,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,
}),
);
@@ -456,6 +500,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(() =>
@@ -465,6 +746,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,
}),
);
@@ -483,6 +766,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,
}),
);
@@ -514,6 +799,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,
}),
);
@@ -558,6 +845,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,
}),
);
@@ -594,6 +883,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,
}),
);
@@ -638,6 +929,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,
}),
);
@@ -653,6 +946,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,
}),
);
@@ -695,6 +990,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,
}),
);
@@ -737,6 +1034,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,
}),
);
@@ -767,6 +1066,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,
}),
);
@@ -797,6 +1098,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,
}),
);
packages/cli/src/ui/hooks/useQuotaAndFallback.ts
+105 -4
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 (
@@ -70,12 +95,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.`,
`/auth to switch to API key.`,
@@ -130,7 +195,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(() => {
@@ -208,11 +282,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,
};
}
@@ -225,5 +326,5 @@ function getResetTimeMessage(delayMs: number): string {
timeZoneName: 'short',
});
return `Access resets at ${timeFormatter.format(resetDate)}.`;
return timeFormatter.format(resetDate);
}
packages/cli/src/ui/types.ts
+2
View File
@@ -201,6 +201,8 @@ export type HistoryItemStats = HistoryItemQuotaBase & {
type: 'stats';
duration: string;
quotas?: RetrieveUserQuotaResponse;
/** G1 AI Credits balance, null if not eligible */
creditBalance?: number | null;
};
export type HistoryItemModelStats = HistoryItemQuotaBase & {
packages/core/src/billing/billing.test.ts
+254
View File
@@ -0,0 +1,254 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { describe, it, expect } from 'vitest';
import type { GeminiUserTier } from '../code_assist/types.js';
import {
buildG1Url,
getG1CreditBalance,
G1_CREDIT_TYPE,
G1_UTM_CAMPAIGNS,
isOverageEligibleModel,
shouldAutoUseCredits,
shouldShowEmptyWalletMenu,
shouldShowOverageMenu,
wrapInAccountChooser,
} from './billing.js';
describe('billing', () => {
describe('wrapInAccountChooser', () => {
it('should wrap URL with AccountChooser redirect', () => {
const result = wrapInAccountChooser(
'user@gmail.com',
'https://one.google.com/ai/activity',
);
expect(result).toBe(
'https://accounts.google.com/AccountChooser?Email=user%40gmail.com&continue=https%3A%2F%2Fone.google.com%2Fai%2Factivity',
);
});
it('should handle special characters in email', () => {
const result = wrapInAccountChooser(
'user+test@example.com',
'https://example.com',
);
expect(result).toContain('Email=user%2Btest%40example.com');
});
});
describe('buildG1Url', () => {
it('should build activity URL with UTM params wrapped in AccountChooser', () => {
const result = buildG1Url(
'activity',
'user@gmail.com',
G1_UTM_CAMPAIGNS.MANAGE_ACTIVITY,
);
// Should contain AccountChooser prefix
expect(result).toContain('https://accounts.google.com/AccountChooser');
expect(result).toContain('Email=user%40gmail.com');
// The continue URL should contain the G1 activity path and UTM params
expect(result).toContain('one.google.com%2Fai%2Factivity');
expect(result).toContain('utm_source%3Dgemini_cli');
expect(result).toContain(
'utm_campaign%3Dhydrogen_cli_settings_ai_credits_activity_page',
);
});
it('should build credits URL with UTM params wrapped in AccountChooser', () => {
const result = buildG1Url(
'credits',
'test@example.com',
G1_UTM_CAMPAIGNS.EMPTY_WALLET_ADD_CREDITS,
);
expect(result).toContain('https://accounts.google.com/AccountChooser');
expect(result).toContain('one.google.com%2Fai%2Fcredits');
expect(result).toContain(
'utm_campaign%3Dhydrogen_cli_insufficient_credits_add_credits',
);
});
});
describe('getG1CreditBalance', () => {
it('should return null for null tier', () => {
expect(getG1CreditBalance(null)).toBeNull();
});
it('should return null for undefined tier', () => {
expect(getG1CreditBalance(undefined)).toBeNull();
});
it('should return null for tier without availableCredits', () => {
const tier: GeminiUserTier = { id: 'PERSONAL' };
expect(getG1CreditBalance(tier)).toBeNull();
});
it('should return null for empty availableCredits array', () => {
const tier: GeminiUserTier = { id: 'PERSONAL', availableCredits: [] };
expect(getG1CreditBalance(tier)).toBeNull();
});
it('should return null when no G1 credit type found', () => {
const tier: GeminiUserTier = {
id: 'PERSONAL',
availableCredits: [
{ creditType: 'CREDIT_TYPE_UNSPECIFIED', creditAmount: '100' },
],
};
expect(getG1CreditBalance(tier)).toBeNull();
});
it('should return G1 credit balance when present', () => {
const tier: GeminiUserTier = {
id: 'PERSONAL',
availableCredits: [{ creditType: G1_CREDIT_TYPE, creditAmount: '500' }],
};
expect(getG1CreditBalance(tier)).toBe(500);
});
it('should return G1 credit balance when multiple credit types present', () => {
const tier: GeminiUserTier = {
id: 'PERSONAL',
availableCredits: [
{ creditType: 'CREDIT_TYPE_UNSPECIFIED', creditAmount: '100' },
{ creditType: G1_CREDIT_TYPE, creditAmount: '750' },
],
};
expect(getG1CreditBalance(tier)).toBe(750);
});
it('should return 0 for invalid credit amount', () => {
const tier: GeminiUserTier = {
id: 'PERSONAL',
availableCredits: [
{ creditType: G1_CREDIT_TYPE, creditAmount: 'invalid' },
],
};
expect(getG1CreditBalance(tier)).toBe(0);
});
it('should handle large credit amounts (int64 as string)', () => {
const tier: GeminiUserTier = {
id: 'PERSONAL',
availableCredits: [
{ creditType: G1_CREDIT_TYPE, creditAmount: '9999999999' },
],
};
expect(getG1CreditBalance(tier)).toBe(9999999999);
});
it('should sum multiple credits of the same G1 type', () => {
const tier: GeminiUserTier = {
id: 'PERSONAL',
availableCredits: [
{ creditType: G1_CREDIT_TYPE, creditAmount: '1000' },
{ creditType: G1_CREDIT_TYPE, creditAmount: '8' },
],
};
expect(getG1CreditBalance(tier)).toBe(1008);
});
});
describe('shouldAutoUseCredits', () => {
it('should return true when strategy is always and balance > 0', () => {
expect(shouldAutoUseCredits('always', 100)).toBe(true);
});
it('should return false when strategy is always but balance is 0', () => {
expect(shouldAutoUseCredits('always', 0)).toBe(false);
});
it('should return false when strategy is ask', () => {
expect(shouldAutoUseCredits('ask', 100)).toBe(false);
});
it('should return false when strategy is never', () => {
expect(shouldAutoUseCredits('never', 100)).toBe(false);
});
it('should return false when creditBalance is null (ineligible)', () => {
expect(shouldAutoUseCredits('always', null)).toBe(false);
});
});
describe('shouldShowOverageMenu', () => {
it('should return true when strategy is ask and balance > 0', () => {
expect(shouldShowOverageMenu('ask', 100)).toBe(true);
});
it('should return false when strategy is ask but balance is 0', () => {
expect(shouldShowOverageMenu('ask', 0)).toBe(false);
});
it('should return false when strategy is always', () => {
expect(shouldShowOverageMenu('always', 100)).toBe(false);
});
it('should return false when strategy is never', () => {
expect(shouldShowOverageMenu('never', 100)).toBe(false);
});
it('should return false when creditBalance is null (ineligible)', () => {
expect(shouldShowOverageMenu('ask', null)).toBe(false);
});
});
describe('shouldShowEmptyWalletMenu', () => {
it('should return true when strategy is ask and balance is 0', () => {
expect(shouldShowEmptyWalletMenu('ask', 0)).toBe(true);
});
it('should return true when strategy is always and balance is 0', () => {
expect(shouldShowEmptyWalletMenu('always', 0)).toBe(true);
});
it('should return false when strategy is never', () => {
expect(shouldShowEmptyWalletMenu('never', 0)).toBe(false);
});
it('should return false when balance > 0', () => {
expect(shouldShowEmptyWalletMenu('ask', 100)).toBe(false);
});
it('should return false when creditBalance is null (ineligible)', () => {
expect(shouldShowEmptyWalletMenu('ask', null)).toBe(false);
});
});
describe('isOverageEligibleModel', () => {
it('should return true for gemini-3-pro-preview', () => {
expect(isOverageEligibleModel('gemini-3-pro-preview')).toBe(true);
});
it('should return true for gemini-3.1-pro-preview', () => {
expect(isOverageEligibleModel('gemini-3.1-pro-preview')).toBe(false);
});
it('should return true for gemini-3.1-pro-preview-customtools', () => {
expect(isOverageEligibleModel('gemini-3.1-pro-preview-customtools')).toBe(
false,
);
});
it('should return false for gemini-3-flash-preview', () => {
expect(isOverageEligibleModel('gemini-3-flash-preview')).toBe(false);
});
it('should return false for gemini-2.5-pro', () => {
expect(isOverageEligibleModel('gemini-2.5-pro')).toBe(false);
});
it('should return false for gemini-2.5-flash', () => {
expect(isOverageEligibleModel('gemini-2.5-flash')).toBe(false);
});
it('should return false for custom model names', () => {
expect(isOverageEligibleModel('my-custom-model')).toBe(false);
});
});
});
packages/core/src/billing/billing.ts
+180
View File
@@ -0,0 +1,180 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type {
AvailableCredits,
CreditType,
GeminiUserTier,
} from '../code_assist/types.js';
import { PREVIEW_GEMINI_MODEL } from '../config/models.js';
/**
* Strategy for handling quota exhaustion when AI credits are available.
* - 'ask': Prompt the user each time
* - 'always': Automatically use credits
* - 'never': Never use credits, show standard fallback
*/
export type OverageStrategy = 'ask' | 'always' | 'never';
/** Credit type for Google One AI credits */
export const G1_CREDIT_TYPE: CreditType = 'GOOGLE_ONE_AI';
/**
* The set of models that support AI credits overage billing.
* Only these models are eligible for the credits-based retry flow.
*/
export const OVERAGE_ELIGIBLE_MODELS = new Set([PREVIEW_GEMINI_MODEL]);
/**
* Checks if a model is eligible for AI credits overage billing.
* @param model The model name to check.
* @returns true if the model supports credits overage, false otherwise.
*/
export function isOverageEligibleModel(model: string): boolean {
return OVERAGE_ELIGIBLE_MODELS.has(model);
}
/** Base URL for Google One AI page */
const G1_AI_BASE_URL = 'https://one.google.com/ai';
/** AccountChooser URL for redirecting with email context */
const ACCOUNT_CHOOSER_URL = 'https://accounts.google.com/AccountChooser';
/** UTM parameters for CLI tracking */
const UTM_SOURCE = 'gemini_cli';
// TODO: change to 'desktop' when G1 service fix is rolled out
const UTM_MEDIUM = 'web';
/**
* Wraps a URL in the AccountChooser redirect to maintain user context.
* @param email User's email address for account selection
* @param continueUrl The destination URL after account selection
* @returns The full AccountChooser redirect URL
*/
export function wrapInAccountChooser(
email: string,
continueUrl: string,
): string {
const params = new URLSearchParams({
Email: email,
continue: continueUrl,
});
return `${ACCOUNT_CHOOSER_URL}?${params.toString()}`;
}
/**
* UTM campaign identifiers per the design doc.
*/
export const G1_UTM_CAMPAIGNS = {
/** From Interception Flow "Manage" link (user has credits) */
MANAGE_ACTIVITY: 'hydrogen_cli_settings_ai_credits_activity_page',
/** From "Manage" to add more credits */
MANAGE_ADD_CREDITS: 'hydrogen_cli_settings_add_credits',
/** From Empty Wallet Flow "Get AI Credits" link */
EMPTY_WALLET_ADD_CREDITS: 'hydrogen_cli_insufficient_credits_add_credits',
} as const;
/**
* Builds a G1 AI URL with UTM tracking parameters.
* @param path The path segment (e.g., 'activity' or 'credits')
* @param email User's email for AccountChooser wrapper
* @param campaign The UTM campaign identifier
* @returns The complete URL wrapped in AccountChooser
*/
export function buildG1Url(
path: 'activity' | 'credits',
email: string,
campaign: string,
): string {
const baseUrl = `${G1_AI_BASE_URL}/${path}`;
const params = new URLSearchParams({
utm_source: UTM_SOURCE,
utm_medium: UTM_MEDIUM,
utm_campaign: campaign,
});
const urlWithUtm = `${baseUrl}?${params.toString()}`;
return wrapInAccountChooser(email, urlWithUtm);
}
/**
* Extracts the G1 AI credit balance from a tier's available credits.
* @param tier The user tier to check
* @returns The credit amount as a number, 0 if eligible but empty, or null if not eligible
*/
export function getG1CreditBalance(
tier: GeminiUserTier | null | undefined,
): number | null {
if (!tier?.availableCredits) {
return null;
}
const g1Credits = tier.availableCredits.filter(
(credit: AvailableCredits) => credit.creditType === G1_CREDIT_TYPE,
);
if (g1Credits.length === 0) {
return null;
}
// creditAmount is an int64 represented as string; sum all matching entries
return g1Credits.reduce((sum, credit) => {
const amount = parseInt(credit.creditAmount ?? '0', 10);
return sum + (isNaN(amount) ? 0 : amount);
}, 0);
}
export const MIN_CREDIT_BALANCE = 50;
/**
* Determines if credits should be automatically used based on the overage strategy.
* @param strategy The configured overage strategy
* @param creditBalance The available credit balance
* @returns true if credits should be auto-used, false otherwise
*/
export function shouldAutoUseCredits(
strategy: OverageStrategy,
creditBalance: number | null,
): boolean {
return (
strategy === 'always' &&
creditBalance != null &&
creditBalance >= MIN_CREDIT_BALANCE
);
}
/**
* Determines if the overage menu should be shown based on the strategy.
* @param strategy The configured overage strategy
* @param creditBalance The available credit balance
* @returns true if the menu should be shown
*/
export function shouldShowOverageMenu(
strategy: OverageStrategy,
creditBalance: number | null,
): boolean {
return (
strategy === 'ask' &&
creditBalance != null &&
creditBalance >= MIN_CREDIT_BALANCE
);
}
/**
* Determines if the empty wallet menu should be shown.
* @param strategy The configured overage strategy
* @param creditBalance The available credit balance
* @returns true if the empty wallet menu should be shown
*/
export function shouldShowEmptyWalletMenu(
strategy: OverageStrategy,
creditBalance: number | null,
): boolean {
return (
strategy !== 'never' &&
creditBalance != null &&
creditBalance < MIN_CREDIT_BALANCE
);
}
packages/core/src/billing/index.ts
+7
View File
@@ -0,0 +1,7 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
export * from './billing.js';
packages/core/src/code_assist/codeAssist.test.ts
+4
View File
@@ -73,6 +73,8 @@ describe('codeAssist', () => {
'session-123',
'free-tier',
'free-tier-name',
undefined,
mockConfig,
);
expect(generator).toBeInstanceOf(MockedCodeAssistServer);
});
@@ -103,6 +105,8 @@ describe('codeAssist', () => {
undefined, // No session ID
'free-tier',
'free-tier-name',
undefined,
mockConfig,
);
expect(generator).toBeInstanceOf(MockedCodeAssistServer);
});
packages/core/src/code_assist/codeAssist.ts
+2
View File
@@ -36,6 +36,8 @@ export async function createCodeAssistContentGenerator(
sessionId,
userData.userTier,
userData.userTierName,
userData.paidTier,
config,
);
}
packages/core/src/code_assist/converter.ts
+19
View File
@@ -27,12 +27,14 @@ import type {
ToolConfig,
} from '@google/genai';
import { GenerateContentResponse } from '@google/genai';
import type { Credits } from './types.js';
export interface CAGenerateContentRequest {
model: string;
project?: string;
user_prompt_id?: string;
request: VertexGenerateContentRequest;
enabled_credit_types?: string[];
}
interface VertexGenerateContentRequest {
@@ -74,6 +76,8 @@ interface VertexGenerationConfig {
export interface CaGenerateContentResponse {
response: VertexGenerateContentResponse;
traceId?: string;
consumedCredits?: Credits[];
remainingCredits?: Credits[];
}
interface VertexGenerateContentResponse {
@@ -121,12 +125,14 @@ export function toGenerateContentRequest(
userPromptId: string,
project?: string,
sessionId?: string,
enabledCreditTypes?: string[],
): CAGenerateContentRequest {
return {
model: req.model,
project,
user_prompt_id: userPromptId,
request: toVertexGenerateContentRequest(req, sessionId),
enabled_credit_types: enabledCreditTypes,
};
}
@@ -299,3 +305,16 @@ function toVertexGenerationConfig(
thinkingConfig: config.thinkingConfig,
};
}
export function fromGenerateContentResponseUsage(
metadata?: GenerateContentResponseUsageMetadata,
): GenerateContentResponseUsageMetadata | undefined {
if (!metadata) {
return undefined;
}
return {
promptTokenCount: metadata.promptTokenCount,
candidatesTokenCount: metadata.candidatesTokenCount,
totalTokenCount: metadata.totalTokenCount,
};
}
packages/core/src/code_assist/server.ts
+101 -1
View File
@@ -21,7 +21,10 @@ import type {
ConversationInteraction,
StreamingLatency,
RecordCodeAssistMetricsRequest,
GeminiUserTier,
Credits,
} from './types.js';
import { UserTierId } from './types.js';
import type {
ListExperimentsRequest,
ListExperimentsResponse,
@@ -37,7 +40,15 @@ import type {
import * as readline from 'node:readline';
import { Readable } from 'node:stream';
import type { ContentGenerator } from '../core/contentGenerator.js';
import { UserTierId } from './types.js';
import type { Config } from '../config/config.js';
import {
G1_CREDIT_TYPE,
getG1CreditBalance,
isOverageEligibleModel,
shouldAutoUseCredits,
} from '../billing/billing.js';
import { logBillingEvent } from '../telemetry/loggers.js';
import { CreditsUsedEvent } from '../telemetry/billingEvents.js';
import type {
CaCountTokenResponse,
CaGenerateContentResponse,
@@ -71,6 +82,8 @@ export class CodeAssistServer implements ContentGenerator {
readonly sessionId?: string,
readonly userTier?: UserTierId,
readonly userTierName?: string,
readonly paidTier?: GeminiUserTier,
readonly config?: Config,
) {}
async generateContentStream(
@@ -79,6 +92,19 @@ export class CodeAssistServer implements ContentGenerator {
// eslint-disable-next-line @typescript-eslint/no-unused-vars
role: LlmRole,
): Promise<AsyncGenerator<GenerateContentResponse>> {
const autoUse = this.config
? shouldAutoUseCredits(
this.config.getBillingSettings().overageStrategy,
getG1CreditBalance(this.paidTier),
)
: false;
const modelIsEligible = isOverageEligibleModel(req.model);
const shouldEnableCredits = modelIsEligible && autoUse;
const enabledCreditTypes = shouldEnableCredits
? ([G1_CREDIT_TYPE] as string[])
: undefined;
const responses =
await this.requestStreamingPost<CaGenerateContentResponse>(
'streamGenerateContent',
@@ -87,6 +113,7 @@ export class CodeAssistServer implements ContentGenerator {
userPromptId,
this.projectId,
this.sessionId,
enabledCreditTypes,
),
req.config?.abortSignal,
);
@@ -98,6 +125,9 @@ export class CodeAssistServer implements ContentGenerator {
return (async function* (
server: CodeAssistServer,
): AsyncGenerator<GenerateContentResponse> {
let totalConsumed = 0;
let lastRemaining = 0;
for await (const response of responses) {
if (isFirst) {
streamingLatency.firstMessageLatency = formatProtoJsonDuration(
@@ -120,8 +150,38 @@ export class CodeAssistServer implements ContentGenerator {
req.config?.abortSignal,
);
if (response.consumedCredits) {
for (const credit of response.consumedCredits) {
if (credit.creditType === G1_CREDIT_TYPE && credit.creditAmount) {
totalConsumed += parseInt(credit.creditAmount, 10) || 0;
}
}
}
if (response.remainingCredits) {
// Sum all G1 credit entries for consistency with getG1CreditBalance
lastRemaining = response.remainingCredits.reduce((sum, credit) => {
if (credit.creditType === G1_CREDIT_TYPE && credit.creditAmount) {
return sum + (parseInt(credit.creditAmount, 10) || 0);
}
return sum;
}, 0);
server.updateCredits(response.remainingCredits);
}
yield translatedResponse;
}
// Emit credits used telemetry after the stream completes
if (totalConsumed > 0 && server.config) {
logBillingEvent(
server.config,
new CreditsUsedEvent(
req.model ?? 'unknown',
totalConsumed,
lastRemaining,
),
);
}
})(this);
}
@@ -139,6 +199,7 @@ export class CodeAssistServer implements ContentGenerator {
userPromptId,
this.projectId,
this.sessionId,
undefined,
),
req.config?.abortSignal,
);
@@ -158,9 +219,29 @@ export class CodeAssistServer implements ContentGenerator {
req.config?.abortSignal,
);
if (response.remainingCredits) {
this.updateCredits(response.remainingCredits);
}
return translatedResponse;
}
private updateCredits(remainingCredits: Credits[]): void {
if (!this.paidTier) {
return;
}
// Replace the G1 credits entries with the latest remaining amounts.
// Non-G1 credits are preserved as-is.
const nonG1Credits = (this.paidTier.availableCredits ?? []).filter(
(c) => c.creditType !== G1_CREDIT_TYPE,
);
const updatedG1Credits = remainingCredits.filter(
(c) => c.creditType === G1_CREDIT_TYPE,
);
this.paidTier.availableCredits = [...nonG1Credits, ...updatedG1Credits];
}
async onboardUser(
req: OnboardUserRequest,
): Promise<LongRunningOperationResponse> {
@@ -190,6 +271,25 @@ export class CodeAssistServer implements ContentGenerator {
}
}
async refreshAvailableCredits(): Promise<void> {
if (!this.paidTier) {
return;
}
const res = await this.loadCodeAssist({
cloudaicompanionProject: this.projectId,
metadata: {
ideType: 'IDE_UNSPECIFIED',
platform: 'PLATFORM_UNSPECIFIED',
pluginType: 'GEMINI',
duetProject: this.projectId,
},
mode: 'HEALTH_CHECK',
});
if (res.paidTier?.availableCredits) {
this.paidTier.availableCredits = res.paidTier.availableCredits;
}
}
async fetchAdminControls(
req: FetchAdminControlsRequest,
): Promise<FetchAdminControlsResponse> {
packages/core/src/code_assist/setup.ts
+3
View File
@@ -51,6 +51,7 @@ export interface UserData {
projectId: string;
userTier: UserTierId;
userTierName?: string;
paidTier?: GeminiUserTier;
}
/**
@@ -136,6 +137,7 @@ export async function setupUser(
projectId,
userTier: loadRes.paidTier?.id ?? loadRes.currentTier.id,
userTierName: loadRes.paidTier?.name ?? loadRes.currentTier.name,
paidTier: loadRes.paidTier ?? undefined,
};
}
@@ -146,6 +148,7 @@ export async function setupUser(
projectId: loadRes.cloudaicompanionProject,
userTier: loadRes.paidTier?.id ?? loadRes.currentTier.id,
userTierName: loadRes.paidTier?.name ?? loadRes.currentTier.name,
paidTier: loadRes.paidTier ?? undefined,
};
}
packages/core/src/code_assist/types.ts
+32
View File
@@ -39,11 +39,41 @@ export type ClientMetadataPluginType =
| 'AIPLUGIN_INTELLIJ'
| 'AIPLUGIN_STUDIO';
/**
* Credit types that can be used for API consumption.
*/
export type CreditType = 'CREDIT_TYPE_UNSPECIFIED' | 'GOOGLE_ONE_AI';
/**
* Represents a credit amount for a specific credit type.
* Used in LoadCodeAssistResponse for available credits and
* in GenerateContentResponse for consumed/remaining credits.
*/
export interface Credits {
creditType: CreditType;
creditAmount: string; // int64 represented as string in JSON
}
/** Alias for Credits used in available_credits context */
export type AvailableCredits = Credits;
/** Alias for Credits used in consumedCredits context */
export type ConsumedCredits = Credits;
/** Alias for Credits used in remainingCredits context */
export type RemainingCredits = Credits;
export interface LoadCodeAssistRequest {
cloudaicompanionProject?: string;
metadata: ClientMetadata;
mode?: LoadCodeAssistMode;
}
export type LoadCodeAssistMode =
| 'MODE_UNSPECIFIED'
| 'FULL_ELIGIBILITY_CHECK'
| 'HEALTH_CHECK';
/**
* Represents LoadCodeAssistResponse proto json field
* http://google3/google/internal/cloud/code/v1internal/cloudcode.proto;l=224
@@ -69,6 +99,8 @@ export interface GeminiUserTier {
privacyNotice?: PrivacyNotice;
hasAcceptedTos?: boolean;
hasOnboardedPreviously?: boolean;
/** Available AI credits for this tier (e.g., Google One AI credits) */
availableCredits?: AvailableCredits[];
}
/**
packages/core/src/config/config.ts
+43
View File
@@ -13,6 +13,7 @@ import type {
ContentGenerator,
ContentGeneratorConfig,
} from '../core/contentGenerator.js';
import type { OverageStrategy } from '../billing/billing.js';
import {
AuthType,
createContentGenerator,
@@ -110,6 +111,7 @@ import { ApprovalMode, type PolicyEngineConfig } from '../policy/types.js';
import { HookSystem } from '../hooks/index.js';
import type {
UserTierId,
GeminiUserTier,
RetrieveUserQuotaResponse,
AdminControlsSettings,
} from '../code_assist/types.js';
@@ -549,6 +551,9 @@ export interface ConfigParameters {
agents?: AgentSettings;
}>;
enableConseca?: boolean;
billing?: {
overageStrategy?: OverageStrategy;
};
}
export class Config {
@@ -725,6 +730,10 @@ export class Config {
}>)
| undefined;
private readonly billing: {
overageStrategy: OverageStrategy;
};
private readonly enableAgents: boolean;
private agents: AgentSettings;
private readonly enableEventDrivenScheduler: boolean;
@@ -962,6 +971,10 @@ export class Config {
this.onModelChange = params.onModelChange;
this.onReload = params.onReload;
this.billing = {
overageStrategy: params.billing?.overageStrategy ?? 'ask',
};
if (params.contextFileName) {
setGeminiMdFilename(params.contextFileName);
}
@@ -1224,6 +1237,10 @@ export class Config {
return this.contentGenerator?.userTierName;
}
getUserPaidTier(): GeminiUserTier | undefined {
return this.contentGenerator?.paidTier;
}
/**
* Provides access to the BaseLlmClient for stateless LLM operations.
*/
@@ -1531,6 +1548,19 @@ export class Config {
this.hasAccessToPreviewModel = hasAccess;
}
async refreshAvailableCredits(): Promise<void> {
const codeAssistServer = getCodeAssistServer(this);
if (!codeAssistServer) {
return;
}
try {
await codeAssistServer.refreshAvailableCredits();
} catch {
// Non-fatal: proceed even if refresh fails.
// The actual credit balance will be verified server-side.
}
}
async refreshUserQuota(): Promise<RetrieveUserQuotaResponse | undefined> {
const codeAssistServer = getCodeAssistServer(this);
if (!codeAssistServer || !codeAssistServer.projectId) {
@@ -2005,6 +2035,19 @@ export class Config {
return this.telemetrySettings.outfile;
}
getBillingSettings(): { overageStrategy: OverageStrategy } {
return this.billing;
}
/**
* Updates the overage strategy at runtime.
* Used to switch from 'ask' to 'always' after the user accepts credits
* via the overage dialog, so subsequent API calls auto-include credits.
*/
setOverageStrategy(strategy: OverageStrategy): void {
this.billing.overageStrategy = strategy;
}
getTelemetryUseCollector(): boolean {
return this.telemetrySettings.useCollector ?? false;
}
packages/core/src/core/contentGenerator.ts
+3 -1
View File
@@ -17,7 +17,7 @@ import { createCodeAssistContentGenerator } from '../code_assist/codeAssist.js';
import type { Config } from '../config/config.js';
import { loadApiKey } from './apiKeyCredentialStorage.js';
import type { UserTierId } from '../code_assist/types.js';
import type { UserTierId, GeminiUserTier } from '../code_assist/types.js';
import { LoggingContentGenerator } from './loggingContentGenerator.js';
import { InstallationManager } from '../utils/installationManager.js';
import { FakeContentGenerator } from './fakeContentGenerator.js';
@@ -49,6 +49,8 @@ export interface ContentGenerator {
userTier?: UserTierId;
userTierName?: string;
paidTier?: GeminiUserTier;
}
export enum AuthType {
packages/core/src/core/fakeContentGenerator.ts
+2 -1
View File
@@ -14,7 +14,7 @@ import {
} from '@google/genai';
import { promises } from 'node:fs';
import type { ContentGenerator } from './contentGenerator.js';
import type { UserTierId } from '../code_assist/types.js';
import type { UserTierId, GeminiUserTier } from '../code_assist/types.js';
import { safeJsonStringify } from '../utils/safeJsonStringify.js';
import type { LlmRole } from '../telemetry/types.js';
@@ -44,6 +44,7 @@ export class FakeContentGenerator implements ContentGenerator {
private callCounter = 0;
userTier?: UserTierId;
userTierName?: string;
paidTier?: GeminiUserTier;
constructor(private readonly responses: FakeResponse[]) {}
packages/core/src/core/loggingContentGenerator.ts
+5 -1
View File
@@ -24,7 +24,7 @@ import {
} from '../telemetry/types.js';
import type { LlmRole } from '../telemetry/llmRole.js';
import type { Config } from '../config/config.js';
import type { UserTierId } from '../code_assist/types.js';
import type { UserTierId, GeminiUserTier } from '../code_assist/types.js';
import {
logApiError,
logApiRequest,
@@ -153,6 +153,10 @@ export class LoggingContentGenerator implements ContentGenerator {
return this.wrapped.userTierName;
}
get paidTier(): GeminiUserTier | undefined {
return this.wrapped.paidTier;
}
private logApiRequest(
contents: Content[],
model: string,
packages/core/src/fallback/handler.ts
+3
View File
@@ -140,6 +140,9 @@ async function processIntent(
// based on the availability service state (which is updated before this).
return true;
case 'retry_with_credits':
return true;
case 'stop':
// Do not switch model on stop. User wants to stay on current model (and stop).
return false;
packages/core/src/fallback/types.ts
+1
View File
@@ -17,6 +17,7 @@ import type {
export type FallbackIntent =
| 'retry_always' // Retry with fallback model and stick to it for future requests.
| 'retry_once' // Retry with fallback model for this request only.
| 'retry_with_credits' // Retry the current request using Google One AI credits (and potentially future ones if strategy is 'always').
| 'stop' // Switch to fallback for future requests, but stop the current request.
| 'retry_later' // Stop the current request and do not fallback. Intend to try again later with the same model.
| 'upgrade'; // Give user an option to upgrade the tier.
packages/core/src/index.ts
+5
View File
@@ -18,6 +18,7 @@ export * from './policy/policy-engine.js';
export * from './policy/toml-loader.js';
export * from './policy/config.js';
export * from './policy/integrity.js';
export * from './billing/index.js';
export * from './confirmation-bus/types.js';
export * from './confirmation-bus/message-bus.js';
@@ -76,6 +77,7 @@ export * from './utils/quotaErrorDetection.js';
export * from './utils/userAccountManager.js';
export * from './utils/authConsent.js';
export * from './utils/googleQuotaErrors.js';
export * from './utils/googleErrors.js';
export * from './utils/fileUtils.js';
export * from './utils/planUtils.js';
export * from './utils/approvalModeUtils.js';
@@ -98,6 +100,7 @@ export * from './utils/ignorePatterns.js';
export * from './utils/partUtils.js';
export * from './utils/promptIdContext.js';
export * from './utils/thoughtUtils.js';
export * from './utils/secure-browser-launcher.js';
export * from './utils/debugLogger.js';
export * from './utils/events.js';
export * from './utils/extensionLoader.js';
@@ -183,6 +186,8 @@ export { OAuthUtils } from './mcp/oauth-utils.js';
// Export telemetry functions
export * from './telemetry/index.js';
export * from './telemetry/billingEvents.js';
export { logBillingEvent } from './telemetry/loggers.js';
export { sessionId, createSessionId } from './utils/session.js';
export * from './utils/compatibility.js';
export * from './utils/browser.js';
packages/core/src/telemetry/billingEvents.test.ts
+206
View File
@@ -0,0 +1,206 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
import { makeFakeConfig } from '../test-utils/config.js';
import {
OverageMenuShownEvent,
OverageOptionSelectedEvent,
EmptyWalletMenuShownEvent,
CreditPurchaseClickEvent,
CreditsUsedEvent,
ApiKeyUpdatedEvent,
EVENT_OVERAGE_MENU_SHOWN,
EVENT_OVERAGE_OPTION_SELECTED,
EVENT_EMPTY_WALLET_MENU_SHOWN,
EVENT_CREDIT_PURCHASE_CLICK,
EVENT_CREDITS_USED,
EVENT_API_KEY_UPDATED,
} from './billingEvents.js';
describe('billingEvents', () => {
const fakeConfig = makeFakeConfig();
beforeEach(() => {
vi.useFakeTimers();
vi.setSystemTime(new Date('2026-01-15T10:30:00.000Z'));
});
afterEach(() => {
vi.useRealTimers();
});
describe('OverageMenuShownEvent', () => {
it('should construct with correct properties', () => {
const event = new OverageMenuShownEvent(
'gemini-3-pro-preview',
500,
'ask',
);
expect(event['event.name']).toBe('overage_menu_shown');
expect(event.model).toBe('gemini-3-pro-preview');
expect(event.credit_balance).toBe(500);
expect(event.overage_strategy).toBe('ask');
});
it('should produce correct OpenTelemetry attributes', () => {
const event = new OverageMenuShownEvent(
'gemini-3-pro-preview',
500,
'ask',
);
const attrs = event.toOpenTelemetryAttributes(fakeConfig);
expect(attrs['event.name']).toBe(EVENT_OVERAGE_MENU_SHOWN);
expect(attrs['model']).toBe('gemini-3-pro-preview');
expect(attrs['credit_balance']).toBe(500);
expect(attrs['overage_strategy']).toBe('ask');
});
it('should produce a human-readable log body', () => {
const event = new OverageMenuShownEvent(
'gemini-3-pro-preview',
500,
'ask',
);
expect(event.toLogBody()).toContain('gemini-3-pro-preview');
expect(event.toLogBody()).toContain('500');
});
});
describe('OverageOptionSelectedEvent', () => {
it('should construct with correct properties', () => {
const event = new OverageOptionSelectedEvent(
'gemini-3-pro-preview',
'use_credits',
100,
);
expect(event['event.name']).toBe('overage_option_selected');
expect(event.selected_option).toBe('use_credits');
expect(event.credit_balance).toBe(100);
});
it('should produce correct OpenTelemetry attributes', () => {
const event = new OverageOptionSelectedEvent(
'gemini-3-pro-preview',
'use_fallback',
200,
);
const attrs = event.toOpenTelemetryAttributes(fakeConfig);
expect(attrs['event.name']).toBe(EVENT_OVERAGE_OPTION_SELECTED);
expect(attrs['selected_option']).toBe('use_fallback');
});
it('should produce a human-readable log body', () => {
const event = new OverageOptionSelectedEvent(
'gemini-3-pro-preview',
'manage',
100,
);
expect(event.toLogBody()).toContain('manage');
expect(event.toLogBody()).toContain('gemini-3-pro-preview');
});
});
describe('EmptyWalletMenuShownEvent', () => {
it('should construct with correct properties', () => {
const event = new EmptyWalletMenuShownEvent('gemini-3-pro-preview');
expect(event['event.name']).toBe('empty_wallet_menu_shown');
expect(event.model).toBe('gemini-3-pro-preview');
});
it('should produce correct OpenTelemetry attributes', () => {
const event = new EmptyWalletMenuShownEvent('gemini-3-pro-preview');
const attrs = event.toOpenTelemetryAttributes(fakeConfig);
expect(attrs['event.name']).toBe(EVENT_EMPTY_WALLET_MENU_SHOWN);
expect(attrs['model']).toBe('gemini-3-pro-preview');
});
it('should produce a human-readable log body', () => {
const event = new EmptyWalletMenuShownEvent('gemini-3-pro-preview');
expect(event.toLogBody()).toContain('gemini-3-pro-preview');
});
});
describe('CreditPurchaseClickEvent', () => {
it('should construct with correct properties', () => {
const event = new CreditPurchaseClickEvent(
'empty_wallet_menu',
'gemini-3-pro-preview',
);
expect(event['event.name']).toBe('credit_purchase_click');
expect(event.source).toBe('empty_wallet_menu');
expect(event.model).toBe('gemini-3-pro-preview');
});
it('should produce correct OpenTelemetry attributes', () => {
const event = new CreditPurchaseClickEvent(
'overage_menu',
'gemini-3-pro-preview',
);
const attrs = event.toOpenTelemetryAttributes(fakeConfig);
expect(attrs['event.name']).toBe(EVENT_CREDIT_PURCHASE_CLICK);
expect(attrs['source']).toBe('overage_menu');
});
it('should produce a human-readable log body', () => {
const event = new CreditPurchaseClickEvent(
'manage',
'gemini-3-pro-preview',
);
expect(event.toLogBody()).toContain('manage');
expect(event.toLogBody()).toContain('gemini-3-pro-preview');
});
});
describe('CreditsUsedEvent', () => {
it('should construct with correct properties', () => {
const event = new CreditsUsedEvent('gemini-3-pro-preview', 10, 490);
expect(event['event.name']).toBe('credits_used');
expect(event.credits_consumed).toBe(10);
expect(event.credits_remaining).toBe(490);
});
it('should produce correct OpenTelemetry attributes', () => {
const event = new CreditsUsedEvent('gemini-3-pro-preview', 10, 490);
const attrs = event.toOpenTelemetryAttributes(fakeConfig);
expect(attrs['event.name']).toBe(EVENT_CREDITS_USED);
expect(attrs['credits_consumed']).toBe(10);
expect(attrs['credits_remaining']).toBe(490);
});
it('should produce a human-readable log body', () => {
const event = new CreditsUsedEvent('gemini-3-pro-preview', 10, 490);
const body = event.toLogBody();
expect(body).toContain('10');
expect(body).toContain('490');
expect(body).toContain('gemini-3-pro-preview');
});
});
describe('ApiKeyUpdatedEvent', () => {
it('should construct with correct properties', () => {
const event = new ApiKeyUpdatedEvent('google_login', 'api_key');
expect(event['event.name']).toBe('api_key_updated');
expect(event.previous_auth_type).toBe('google_login');
expect(event.new_auth_type).toBe('api_key');
});
it('should produce correct OpenTelemetry attributes', () => {
const event = new ApiKeyUpdatedEvent('google_login', 'api_key');
const attrs = event.toOpenTelemetryAttributes(fakeConfig);
expect(attrs['event.name']).toBe(EVENT_API_KEY_UPDATED);
expect(attrs['previous_auth_type']).toBe('google_login');
expect(attrs['new_auth_type']).toBe('api_key');
});
it('should produce a human-readable log body', () => {
const event = new ApiKeyUpdatedEvent('google_login', 'api_key');
const body = event.toLogBody();
expect(body).toContain('google_login');
expect(body).toContain('api_key');
});
});
});
packages/core/src/telemetry/billingEvents.ts
+255
View File
@@ -0,0 +1,255 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type { Config } from '../config/config.js';
import type { LogAttributes } from '@opentelemetry/api-logs';
import type { BaseTelemetryEvent } from './types.js';
import { getCommonAttributes } from './telemetryAttributes.js';
import type { OverageStrategy } from '../billing/billing.js';
/** Overage menu option that can be selected by the user */
export type OverageOption =
| 'use_credits'
| 'use_fallback'
| 'manage'
| 'stop'
| 'get_credits';
// ============================================================================
// Event: Overage Menu Shown
// ============================================================================
export const EVENT_OVERAGE_MENU_SHOWN = 'gemini_cli.overage_menu_shown';
export class OverageMenuShownEvent implements BaseTelemetryEvent {
'event.name': 'overage_menu_shown';
'event.timestamp': string;
model: string;
credit_balance: number;
overage_strategy: OverageStrategy;
constructor(
model: string,
creditBalance: number,
overageStrategy: OverageStrategy,
) {
this['event.name'] = 'overage_menu_shown';
this['event.timestamp'] = new Date().toISOString();
this.model = model;
this.credit_balance = creditBalance;
this.overage_strategy = overageStrategy;
}
toOpenTelemetryAttributes(config: Config): LogAttributes {
return {
...getCommonAttributes(config),
'event.name': EVENT_OVERAGE_MENU_SHOWN,
'event.timestamp': this['event.timestamp'],
model: this.model,
credit_balance: this.credit_balance,
overage_strategy: this.overage_strategy,
};
}
toLogBody(): string {
return `Overage menu shown for model ${this.model} with ${this.credit_balance} credits available.`;
}
}
// ============================================================================
// Event: Overage Option Selected
// ============================================================================
export const EVENT_OVERAGE_OPTION_SELECTED =
'gemini_cli.overage_option_selected';
export class OverageOptionSelectedEvent implements BaseTelemetryEvent {
'event.name': 'overage_option_selected';
'event.timestamp': string;
model: string;
selected_option: OverageOption;
credit_balance: number;
constructor(
model: string,
selectedOption: OverageOption,
creditBalance: number,
) {
this['event.name'] = 'overage_option_selected';
this['event.timestamp'] = new Date().toISOString();
this.model = model;
this.selected_option = selectedOption;
this.credit_balance = creditBalance;
}
toOpenTelemetryAttributes(config: Config): LogAttributes {
return {
...getCommonAttributes(config),
'event.name': EVENT_OVERAGE_OPTION_SELECTED,
'event.timestamp': this['event.timestamp'],
model: this.model,
selected_option: this.selected_option,
credit_balance: this.credit_balance,
};
}
toLogBody(): string {
return `Overage option '${this.selected_option}' selected for model ${this.model}.`;
}
}
// ============================================================================
// Event: Empty Wallet Menu Shown
// ============================================================================
export const EVENT_EMPTY_WALLET_MENU_SHOWN =
'gemini_cli.empty_wallet_menu_shown';
export class EmptyWalletMenuShownEvent implements BaseTelemetryEvent {
'event.name': 'empty_wallet_menu_shown';
'event.timestamp': string;
model: string;
constructor(model: string) {
this['event.name'] = 'empty_wallet_menu_shown';
this['event.timestamp'] = new Date().toISOString();
this.model = model;
}
toOpenTelemetryAttributes(config: Config): LogAttributes {
return {
...getCommonAttributes(config),
'event.name': EVENT_EMPTY_WALLET_MENU_SHOWN,
'event.timestamp': this['event.timestamp'],
model: this.model,
};
}
toLogBody(): string {
return `Empty wallet menu shown for model ${this.model}.`;
}
}
// ============================================================================
// Event: Credit Purchase Click
// ============================================================================
export const EVENT_CREDIT_PURCHASE_CLICK = 'gemini_cli.credit_purchase_click';
export class CreditPurchaseClickEvent implements BaseTelemetryEvent {
'event.name': 'credit_purchase_click';
'event.timestamp': string;
source: 'overage_menu' | 'empty_wallet_menu' | 'manage';
model: string;
constructor(
source: 'overage_menu' | 'empty_wallet_menu' | 'manage',
model: string,
) {
this['event.name'] = 'credit_purchase_click';
this['event.timestamp'] = new Date().toISOString();
this.source = source;
this.model = model;
}
toOpenTelemetryAttributes(config: Config): LogAttributes {
return {
...getCommonAttributes(config),
'event.name': EVENT_CREDIT_PURCHASE_CLICK,
'event.timestamp': this['event.timestamp'],
source: this.source,
model: this.model,
};
}
toLogBody(): string {
return `Credit purchase clicked from ${this.source} for model ${this.model}.`;
}
}
// ============================================================================
// Event: Credits Used
// ============================================================================
export const EVENT_CREDITS_USED = 'gemini_cli.credits_used';
export class CreditsUsedEvent implements BaseTelemetryEvent {
'event.name': 'credits_used';
'event.timestamp': string;
model: string;
credits_consumed: number;
credits_remaining: number;
constructor(
model: string,
creditsConsumed: number,
creditsRemaining: number,
) {
this['event.name'] = 'credits_used';
this['event.timestamp'] = new Date().toISOString();
this.model = model;
this.credits_consumed = creditsConsumed;
this.credits_remaining = creditsRemaining;
}
toOpenTelemetryAttributes(config: Config): LogAttributes {
return {
...getCommonAttributes(config),
'event.name': EVENT_CREDITS_USED,
'event.timestamp': this['event.timestamp'],
model: this.model,
credits_consumed: this.credits_consumed,
credits_remaining: this.credits_remaining,
};
}
toLogBody(): string {
return `${this.credits_consumed} credits consumed for model ${this.model}. ${this.credits_remaining} remaining.`;
}
}
// ============================================================================
// Event: API Key Updated (Auth Type Changed)
// ============================================================================
export const EVENT_API_KEY_UPDATED = 'gemini_cli.api_key_updated';
export class ApiKeyUpdatedEvent implements BaseTelemetryEvent {
'event.name': 'api_key_updated';
'event.timestamp': string;
previous_auth_type: string;
new_auth_type: string;
constructor(previousAuthType: string, newAuthType: string) {
this['event.name'] = 'api_key_updated';
this['event.timestamp'] = new Date().toISOString();
this.previous_auth_type = previousAuthType;
this.new_auth_type = newAuthType;
}
toOpenTelemetryAttributes(config: Config): LogAttributes {
return {
...getCommonAttributes(config),
'event.name': EVENT_API_KEY_UPDATED,
'event.timestamp': this['event.timestamp'],
previous_auth_type: this.previous_auth_type,
new_auth_type: this.new_auth_type,
};
}
toLogBody(): string {
return `Auth type changed from ${this.previous_auth_type} to ${this.new_auth_type}.`;
}
}
/** Union type of all billing-related telemetry events */
export type BillingTelemetryEvent =
| OverageMenuShownEvent
| OverageOptionSelectedEvent
| EmptyWalletMenuShownEvent
| CreditPurchaseClickEvent
| CreditsUsedEvent
| ApiKeyUpdatedEvent;
packages/core/src/telemetry/conseca-logger.test.ts
+1
View File
@@ -39,6 +39,7 @@ describe('conseca-logger', () => {
getTelemetryLogPromptsEnabled: vi.fn().mockReturnValue(true),
isInteractive: vi.fn().mockReturnValue(true),
getExperiments: vi.fn().mockReturnValue({ experimentIds: [] }),
getContentGeneratorConfig: vi.fn().mockReturnValue({ authType: 'oauth' }),
} as unknown as Config;
mockLogger = {
packages/core/src/telemetry/index.ts
+4
View File
@@ -77,6 +77,7 @@ export type { TelemetryEvent } from './types.js';
export { SpanStatusCode, ValueType } from '@opentelemetry/api';
export { SemanticAttributes } from '@opentelemetry/semantic-conventions';
export * from './uiTelemetry.js';
export * from './billingEvents.js';
export {
MemoryMonitor,
initializeMemoryMonitor,
@@ -145,6 +146,9 @@ export {
GenAiOperationName,
GenAiProviderName,
GenAiTokenType,
// Billing metrics functions
recordOverageOptionSelected,
recordCreditPurchaseClick,
} from './metrics.js';
export { runInDevTraceSpan, type SpanMetadata } from './trace.js';
export { startupProfiler, StartupProfiler } from './startupProfiler.js';
packages/core/src/telemetry/loggers.test.ts
+20 -4
View File
@@ -280,6 +280,7 @@ describe('loggers', () => {
isInteractive: () => false,
getExperiments: () => undefined,
getExperimentsAsync: async () => undefined,
getContentGeneratorConfig: () => undefined,
} as unknown as Config;
it('should log a user prompt', () => {
@@ -319,6 +320,7 @@ describe('loggers', () => {
isInteractive: () => false,
getExperiments: () => undefined,
getExperimentsAsync: async () => undefined,
getContentGeneratorConfig: () => undefined,
} as unknown as Config;
const event = new UserPromptEvent(
11,
@@ -356,7 +358,8 @@ describe('loggers', () => {
isInteractive: () => false,
getExperiments: () => undefined,
getExperimentsAsync: async () => undefined,
} as Config;
getContentGeneratorConfig: () => undefined,
} as unknown as Config;
const mockMetrics = {
recordApiResponseMetrics: vi.fn(),
@@ -558,7 +561,8 @@ describe('loggers', () => {
isInteractive: () => false,
getExperiments: () => undefined,
getExperimentsAsync: async () => undefined,
} as Config;
getContentGeneratorConfig: () => undefined,
} as unknown as Config;
const mockMetrics = {
recordApiResponseMetrics: vi.fn(),
@@ -996,6 +1000,7 @@ describe('loggers', () => {
isInteractive: () => false,
getExperiments: () => undefined,
getExperimentsAsync: async () => undefined,
getContentGeneratorConfig: () => undefined,
} as unknown as Config;
it('should log flash fallback event', () => {
@@ -1025,6 +1030,7 @@ describe('loggers', () => {
isInteractive: () => false,
getExperiments: () => undefined,
getExperimentsAsync: async () => undefined,
getContentGeneratorConfig: () => undefined,
} as unknown as Config;
beforeEach(() => {
@@ -1121,7 +1127,8 @@ describe('loggers', () => {
isInteractive: () => false,
getExperiments: () => undefined,
getExperimentsAsync: async () => undefined,
} as Config;
getContentGeneratorConfig: () => undefined,
} as unknown as Config;
const mockMetrics = {
recordToolCallMetrics: vi.fn(),
@@ -1741,7 +1748,8 @@ describe('loggers', () => {
isInteractive: () => false,
getExperiments: () => undefined,
getExperimentsAsync: async () => undefined,
} as Config;
getContentGeneratorConfig: () => undefined,
} as unknown as Config;
const mockMetrics = {
recordFileOperationMetric: vi.fn(),
@@ -1803,6 +1811,7 @@ describe('loggers', () => {
isInteractive: () => false,
getExperiments: () => undefined,
getExperimentsAsync: async () => undefined,
getContentGeneratorConfig: () => undefined,
} as unknown as Config;
it('should log a tool output truncated event', () => {
@@ -1842,6 +1851,7 @@ describe('loggers', () => {
isInteractive: () => false,
getExperiments: () => undefined,
getExperimentsAsync: async () => undefined,
getContentGeneratorConfig: () => undefined,
} as unknown as Config;
beforeEach(() => {
@@ -2099,6 +2109,7 @@ describe('loggers', () => {
isInteractive: () => false,
getExperiments: () => undefined,
getExperimentsAsync: async () => undefined,
getContentGeneratorConfig: () => undefined,
} as unknown as Config;
beforeEach(() => {
@@ -2146,6 +2157,7 @@ describe('loggers', () => {
isInteractive: () => false,
getExperiments: () => undefined,
getExperimentsAsync: async () => undefined,
getContentGeneratorConfig: () => undefined,
} as unknown as Config;
beforeEach(() => {
@@ -2193,6 +2205,7 @@ describe('loggers', () => {
isInteractive: () => false,
getExperiments: () => undefined,
getExperimentsAsync: async () => undefined,
getContentGeneratorConfig: () => undefined,
} as unknown as Config;
beforeEach(() => {
@@ -2231,6 +2244,7 @@ describe('loggers', () => {
isInteractive: () => false,
getExperiments: () => undefined,
getExperimentsAsync: async () => undefined,
getContentGeneratorConfig: () => undefined,
} as unknown as Config;
beforeEach(() => {
@@ -2284,6 +2298,7 @@ describe('loggers', () => {
isInteractive: () => false,
getExperiments: () => undefined,
getExperimentsAsync: async () => undefined,
getContentGeneratorConfig: () => undefined,
} as unknown as Config;
beforeEach(() => {
@@ -2322,6 +2337,7 @@ describe('loggers', () => {
getExperiments: () => undefined,
getExperimentsAsync: async () => undefined,
getTelemetryLogPromptsEnabled: () => false,
getContentGeneratorConfig: () => undefined,
} as unknown as Config;
beforeEach(() => {
packages/core/src/telemetry/loggers.ts
+15
View File
@@ -84,6 +84,7 @@ import type { UiEvent } from './uiTelemetry.js';
import { uiTelemetryService } from './uiTelemetry.js';
import { ClearcutLogger } from './clearcut-logger/clearcut-logger.js';
import { debugLogger } from '../utils/debugLogger.js';
import type { BillingTelemetryEvent } from './billingEvents.js';
export function logCliConfiguration(
config: Config,
@@ -827,3 +828,17 @@ export function logTokenStorageInitialization(
recordTokenStorageInitialization(config, event);
});
}
export function logBillingEvent(
config: Config,
event: BillingTelemetryEvent,
): void {
bufferTelemetryEvent(() => {
const logger = logs.getLogger(SERVICE_NAME);
const logRecord: LogRecord = {
body: event.toLogBody(),
attributes: event.toOpenTelemetryAttributes(config),
};
logger.emit(logRecord);
});
}
packages/core/src/telemetry/metrics.ts
+52
View File
@@ -41,6 +41,8 @@ const EVENT_HOOK_CALL_COUNT = 'gemini_cli.hook_call.count';
const EVENT_HOOK_CALL_LATENCY = 'gemini_cli.hook_call.latency';
const KEYCHAIN_AVAILABILITY_COUNT = 'gemini_cli.keychain.availability.count';
const TOKEN_STORAGE_TYPE_COUNT = 'gemini_cli.token_storage.type.count';
const OVERAGE_OPTION_COUNT = 'gemini_cli.overage_option.count';
const CREDIT_PURCHASE_COUNT = 'gemini_cli.credit_purchase.count';
// Agent Metrics
const AGENT_RUN_COUNT = 'gemini_cli.agent.run.count';
@@ -259,6 +261,26 @@ const COUNTER_DEFINITIONS = {
forced: boolean;
},
},
[OVERAGE_OPTION_COUNT]: {
description: 'Counts overage option selections.',
valueType: ValueType.INT,
assign: (c: Counter) => (overageOptionCounter = c),
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
attributes: {} as {
selected_option: string;
model: string;
},
},
[CREDIT_PURCHASE_COUNT]: {
description: 'Counts credit purchase link clicks.',
valueType: ValueType.INT,
assign: (c: Counter) => (creditPurchaseCounter = c),
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
attributes: {} as {
source: string;
model: string;
},
},
} as const;
const HISTOGRAM_DEFINITIONS = {
@@ -597,6 +619,8 @@ let hookCallCounter: Counter | undefined;
let hookCallLatencyHistogram: Histogram | undefined;
let keychainAvailabilityCounter: Counter | undefined;
let tokenStorageTypeCounter: Counter | undefined;
let overageOptionCounter: Counter | undefined;
let creditPurchaseCounter: Counter | undefined;
// OpenTelemetry GenAI Semantic Convention Metrics
let genAiClientTokenUsageHistogram: Histogram | undefined;
@@ -1334,3 +1358,31 @@ export function recordTokenStorageInitialization(
forced: event.forced,
});
}
/**
* Records a metric for an overage option selection.
*/
export function recordOverageOptionSelected(
config: Config,
attributes: MetricDefinitions[typeof OVERAGE_OPTION_COUNT]['attributes'],
): void {
if (!overageOptionCounter || !isMetricsInitialized) return;
overageOptionCounter.add(1, {
...baseMetricDefinition.getCommonAttributes(config),
...attributes,
});
}
/**
* Records a metric for a credit purchase link click.
*/
export function recordCreditPurchaseClick(
config: Config,
attributes: MetricDefinitions[typeof CREDIT_PURCHASE_COUNT]['attributes'],
): void {
if (!creditPurchaseCounter || !isMetricsInitialized) return;
creditPurchaseCounter.add(1, {
...baseMetricDefinition.getCommonAttributes(config),
...attributes,
});
}
packages/core/src/telemetry/sanitize.test.ts
+1
View File
@@ -32,6 +32,7 @@ function createMockConfig(logPromptsEnabled: boolean): Config {
getModel: () => 'gemini-1.5-flash',
isInteractive: () => true,
getUserEmail: () => undefined,
getContentGeneratorConfig: () => undefined,
} as unknown as Config;
}
packages/core/src/telemetry/sdk.test.ts
+1
View File
@@ -77,6 +77,7 @@ describe('Telemetry SDK', () => {
isInteractive: () => false,
getExperiments: () => undefined,
getExperimentsAsync: async () => undefined,
getContentGeneratorConfig: () => undefined,
} as unknown as Config;
});
packages/core/src/telemetry/telemetryAttributes.ts
+2
View File
@@ -15,11 +15,13 @@ const installationManager = new InstallationManager();
export function getCommonAttributes(config: Config): Attributes {
const email = userAccountManager.getCachedGoogleAccount();
const experiments = config.getExperiments();
const authType = config.getContentGeneratorConfig()?.authType;
return {
'session.id': config.getSessionId(),
'installation.id': installationManager.getInstallationId(),
interactive: config.isInteractive(),
...(email && { 'user.email': email }),
...(authType && { auth_type: authType }),
...(experiments &&
experiments.experimentIds.length > 0 && {
'experiments.ids': experiments.experimentIds,
packages/core/src/utils/googleErrors.ts
+2 -2
View File
@@ -16,8 +16,8 @@
export interface ErrorInfo {
'@type': 'type.googleapis.com/google.rpc.ErrorInfo';
reason: string;
domain: string;
metadata: { [key: string]: string };
domain?: string;
metadata?: { [key: string]: string };
}
export interface RetryInfo {
packages/core/src/utils/googleQuotaErrors.test.ts
+20
View File
@@ -297,6 +297,26 @@ describe('classifyGoogleError', () => {
expect(result).toBeInstanceOf(TerminalQuotaError);
});
it('should return TerminalQuotaError for INSUFFICIENT_G1_CREDITS_BALANCE without domain', () => {
const apiError: GoogleApiError = {
code: 429,
message: 'Resource has been exhausted (e.g. check quota).',
details: [
{
'@type': 'type.googleapis.com/google.rpc.ErrorInfo',
reason: 'INSUFFICIENT_G1_CREDITS_BALANCE',
},
],
};
vi.spyOn(errorParser, 'parseGoogleApiError').mockReturnValue(apiError);
const result = classifyGoogleError(new Error());
expect(result).toBeInstanceOf(TerminalQuotaError);
expect((result as TerminalQuotaError).isInsufficientCredits).toBe(true);
expect((result as TerminalQuotaError).reason).toBe(
'INSUFFICIENT_G1_CREDITS_BALANCE',
);
});
it('should prioritize daily limit over retry info', () => {
const apiError: GoogleApiError = {
code: 429,
packages/core/src/utils/googleQuotaErrors.ts
+19
View File
@@ -19,17 +19,24 @@ import { getErrorStatus, ModelNotFoundError } from './httpErrors.js';
*/
export class TerminalQuotaError extends Error {
retryDelayMs?: number;
reason?: string;
constructor(
message: string,
override readonly cause: GoogleApiError,
retryDelaySeconds?: number,
reason?: string,
) {
super(message);
this.name = 'TerminalQuotaError';
this.retryDelayMs = retryDelaySeconds
? retryDelaySeconds * 1000
: undefined;
this.reason = reason;
}
get isInsufficientCredits(): boolean {
return this.reason === 'INSUFFICIENT_G1_CREDITS_BALANCE';
}
}
@@ -121,6 +128,7 @@ function classifyValidationRequiredError(
}
if (
!errorInfo.domain ||
!CLOUDCODE_DOMAINS.includes(errorInfo.domain) ||
errorInfo.reason !== 'VALIDATION_REQUIRED'
) {
@@ -293,6 +301,16 @@ export function classifyGoogleError(error: unknown): unknown {
}
if (errorInfo) {
// INSUFFICIENT_G1_CREDITS_BALANCE is always terminal, regardless of domain
if (errorInfo.reason === 'INSUFFICIENT_G1_CREDITS_BALANCE') {
return new TerminalQuotaError(
`${googleApiError.message}`,
googleApiError,
delaySeconds,
errorInfo.reason,
);
}
// New Cloud Code API quota handling
if (errorInfo.domain) {
const validDomains = [
@@ -313,6 +331,7 @@ export function classifyGoogleError(error: unknown): unknown {
`${googleApiError.message}`,
googleApiError,
delaySeconds,
errorInfo.reason,
);
}
}
schemas/settings.schema.json
+18
View File
@@ -520,6 +520,24 @@
"markdownDescription": "Telemetry configuration.\n\n- Category: `Advanced`\n- Requires restart: `yes`",
"$ref": "#/$defs/TelemetrySettings"
},
"billing": {
"title": "Billing",
"description": "Billing and AI credits settings.",
"markdownDescription": "Billing and AI credits settings.\n\n- Category: `Advanced`\n- Requires restart: `no`\n- Default: `{}`",
"default": {},
"type": "object",
"properties": {
"overageStrategy": {
"title": "Overage Strategy",
"description": "How to handle quota exhaustion when AI credits are available. 'ask' prompts each time, 'always' automatically uses credits, 'never' disables credit usage.",
"markdownDescription": "How to handle quota exhaustion when AI credits are available. 'ask' prompts each time, 'always' automatically uses credits, 'never' disables credit usage.\n\n- Category: `Advanced`\n- Requires restart: `no`\n- Default: `ask`",
"default": "ask",
"type": "string",
"enum": ["ask", "always", "never"]
}
},
"additionalProperties": false
},
"model": {
"title": "Model",
"description": "Settings related to the generative model.",