mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-07-22 15:51:18 -07:00
feat(billing): implement G1 AI credits overage flow with billing telemetry (#18590)
This commit is contained in:
@@ -80,6 +80,12 @@ they appear in the UI.
|
|||||||
| -------- | ------------- | ---------------------------- | ------- |
|
| -------- | ------------- | ---------------------------- | ------- |
|
||||||
| IDE Mode | `ide.enabled` | Enable IDE integration mode. | `false` |
|
| 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
|
### Model
|
||||||
|
|
||||||
| UI Label | Setting | Description | Default |
|
| UI Label | Setting | Description | Default |
|
||||||
|
|||||||
@@ -357,6 +357,15 @@ their corresponding top-level category object in your `settings.json` file.
|
|||||||
- **Default:** `true`
|
- **Default:** `true`
|
||||||
- **Requires restart:** Yes
|
- **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`
|
||||||
|
|
||||||
- **`model.name`** (string):
|
- **`model.name`** (string):
|
||||||
|
|||||||
@@ -828,6 +828,36 @@ const SETTINGS_SCHEMA = {
|
|||||||
ref: 'TelemetrySettings',
|
ref: 'TelemetrySettings',
|
||||||
},
|
},
|
||||||
|
|
||||||
|
billing: {
|
||||||
|
type: 'object',
|
||||||
|
label: 'Billing',
|
||||||
|
category: 'Advanced',
|
||||||
|
requiresRestart: false,
|
||||||
|
default: {},
|
||||||
|
description: 'Billing and AI credits settings.',
|
||||||
|
showInDialog: false,
|
||||||
|
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: {
|
model: {
|
||||||
type: 'object',
|
type: 'object',
|
||||||
label: 'Model',
|
label: 'Model',
|
||||||
|
|||||||
@@ -591,6 +591,8 @@ const mockUIActions: UIActions = {
|
|||||||
handleClearScreen: vi.fn(),
|
handleClearScreen: vi.fn(),
|
||||||
handleProQuotaChoice: vi.fn(),
|
handleProQuotaChoice: vi.fn(),
|
||||||
handleValidationChoice: vi.fn(),
|
handleValidationChoice: vi.fn(),
|
||||||
|
handleOverageMenuChoice: vi.fn(),
|
||||||
|
handleEmptyWalletChoice: vi.fn(),
|
||||||
setQueueErrorMessage: vi.fn(),
|
setQueueErrorMessage: vi.fn(),
|
||||||
popAllMessages: vi.fn(),
|
popAllMessages: vi.fn(),
|
||||||
handleApiKeySubmit: vi.fn(),
|
handleApiKeySubmit: vi.fn(),
|
||||||
|
|||||||
@@ -47,6 +47,7 @@ import {
|
|||||||
type IdeInfo,
|
type IdeInfo,
|
||||||
type IdeContext,
|
type IdeContext,
|
||||||
type UserTierId,
|
type UserTierId,
|
||||||
|
type GeminiUserTier,
|
||||||
type UserFeedbackPayload,
|
type UserFeedbackPayload,
|
||||||
type AgentDefinition,
|
type AgentDefinition,
|
||||||
type ApprovalMode,
|
type ApprovalMode,
|
||||||
@@ -82,6 +83,8 @@ import {
|
|||||||
CoreToolCallStatus,
|
CoreToolCallStatus,
|
||||||
generateSteeringAckMessage,
|
generateSteeringAckMessage,
|
||||||
buildUserSteeringHintPrompt,
|
buildUserSteeringHintPrompt,
|
||||||
|
logBillingEvent,
|
||||||
|
ApiKeyUpdatedEvent,
|
||||||
} from '@google/gemini-cli-core';
|
} from '@google/gemini-cli-core';
|
||||||
import { validateAuthMethod } from '../config/auth.js';
|
import { validateAuthMethod } from '../config/auth.js';
|
||||||
import process from 'node:process';
|
import process from 'node:process';
|
||||||
@@ -391,6 +394,9 @@ export const AppContainer = (props: AppContainerProps) => {
|
|||||||
? { remaining, limit, resetTime }
|
? { remaining, limit, resetTime }
|
||||||
: undefined;
|
: undefined;
|
||||||
});
|
});
|
||||||
|
const [paidTier, setPaidTier] = useState<GeminiUserTier | undefined>(
|
||||||
|
undefined,
|
||||||
|
);
|
||||||
|
|
||||||
const [isConfigInitialized, setConfigInitialized] = useState(false);
|
const [isConfigInitialized, setConfigInitialized] = useState(false);
|
||||||
|
|
||||||
@@ -686,10 +692,17 @@ export const AppContainer = (props: AppContainerProps) => {
|
|||||||
handleProQuotaChoice,
|
handleProQuotaChoice,
|
||||||
validationRequest,
|
validationRequest,
|
||||||
handleValidationChoice,
|
handleValidationChoice,
|
||||||
|
// G1 AI Credits
|
||||||
|
overageMenuRequest,
|
||||||
|
handleOverageMenuChoice,
|
||||||
|
emptyWalletRequest,
|
||||||
|
handleEmptyWalletChoice,
|
||||||
} = useQuotaAndFallback({
|
} = useQuotaAndFallback({
|
||||||
config,
|
config,
|
||||||
historyManager,
|
historyManager,
|
||||||
userTier,
|
userTier,
|
||||||
|
paidTier,
|
||||||
|
settings,
|
||||||
setModelSwitchedFromQuotaError,
|
setModelSwitchedFromQuotaError,
|
||||||
onShowAuthSelection: () => setAuthState(AuthState.Updating),
|
onShowAuthSelection: () => setAuthState(AuthState.Updating),
|
||||||
});
|
});
|
||||||
@@ -729,6 +742,8 @@ export const AppContainer = (props: AppContainerProps) => {
|
|||||||
const handleAuthSelect = useCallback(
|
const handleAuthSelect = useCallback(
|
||||||
async (authType: AuthType | undefined, scope: LoadableSettingScope) => {
|
async (authType: AuthType | undefined, scope: LoadableSettingScope) => {
|
||||||
if (authType) {
|
if (authType) {
|
||||||
|
const previousAuthType =
|
||||||
|
config.getContentGeneratorConfig()?.authType ?? 'unknown';
|
||||||
if (authType === AuthType.LOGIN_WITH_GOOGLE) {
|
if (authType === AuthType.LOGIN_WITH_GOOGLE) {
|
||||||
setAuthContext({ requiresRestart: true });
|
setAuthContext({ requiresRestart: true });
|
||||||
} else {
|
} else {
|
||||||
@@ -741,6 +756,10 @@ export const AppContainer = (props: AppContainerProps) => {
|
|||||||
config.setRemoteAdminSettings(undefined);
|
config.setRemoteAdminSettings(undefined);
|
||||||
await config.refreshAuth(authType);
|
await config.refreshAuth(authType);
|
||||||
setAuthState(AuthState.Authenticated);
|
setAuthState(AuthState.Authenticated);
|
||||||
|
logBillingEvent(
|
||||||
|
config,
|
||||||
|
new ApiKeyUpdatedEvent(previousAuthType, authType),
|
||||||
|
);
|
||||||
} catch (e) {
|
} catch (e) {
|
||||||
if (e instanceof ChangeAuthRequestedError) {
|
if (e instanceof ChangeAuthRequestedError) {
|
||||||
return;
|
return;
|
||||||
@@ -803,6 +822,7 @@ Logging in with Google... Restarting Gemini CLI to continue.
|
|||||||
// Only sync when not currently authenticating
|
// Only sync when not currently authenticating
|
||||||
if (authState === AuthState.Authenticated) {
|
if (authState === AuthState.Authenticated) {
|
||||||
setUserTier(config.getUserTier());
|
setUserTier(config.getUserTier());
|
||||||
|
setPaidTier(config.getUserPaidTier());
|
||||||
}
|
}
|
||||||
}, [config, authState]);
|
}, [config, authState]);
|
||||||
|
|
||||||
@@ -2006,6 +2026,8 @@ Logging in with Google... Restarting Gemini CLI to continue.
|
|||||||
showIdeRestartPrompt ||
|
showIdeRestartPrompt ||
|
||||||
!!proQuotaRequest ||
|
!!proQuotaRequest ||
|
||||||
!!validationRequest ||
|
!!validationRequest ||
|
||||||
|
!!overageMenuRequest ||
|
||||||
|
!!emptyWalletRequest ||
|
||||||
isSessionBrowserOpen ||
|
isSessionBrowserOpen ||
|
||||||
authState === AuthState.AwaitingApiKeyInput ||
|
authState === AuthState.AwaitingApiKeyInput ||
|
||||||
!!newAgents;
|
!!newAgents;
|
||||||
@@ -2033,6 +2055,8 @@ Logging in with Google... Restarting Gemini CLI to continue.
|
|||||||
hasLoopDetectionConfirmationRequest ||
|
hasLoopDetectionConfirmationRequest ||
|
||||||
!!proQuotaRequest ||
|
!!proQuotaRequest ||
|
||||||
!!validationRequest ||
|
!!validationRequest ||
|
||||||
|
!!overageMenuRequest ||
|
||||||
|
!!emptyWalletRequest ||
|
||||||
!!customDialog;
|
!!customDialog;
|
||||||
|
|
||||||
const allowPlanMode =
|
const allowPlanMode =
|
||||||
@@ -2243,6 +2267,9 @@ Logging in with Google... Restarting Gemini CLI to continue.
|
|||||||
stats: quotaStats,
|
stats: quotaStats,
|
||||||
proQuotaRequest,
|
proQuotaRequest,
|
||||||
validationRequest,
|
validationRequest,
|
||||||
|
// G1 AI Credits dialog state
|
||||||
|
overageMenuRequest,
|
||||||
|
emptyWalletRequest,
|
||||||
},
|
},
|
||||||
contextFileNames,
|
contextFileNames,
|
||||||
errorCount,
|
errorCount,
|
||||||
@@ -2367,6 +2394,8 @@ Logging in with Google... Restarting Gemini CLI to continue.
|
|||||||
quotaStats,
|
quotaStats,
|
||||||
proQuotaRequest,
|
proQuotaRequest,
|
||||||
validationRequest,
|
validationRequest,
|
||||||
|
overageMenuRequest,
|
||||||
|
emptyWalletRequest,
|
||||||
contextFileNames,
|
contextFileNames,
|
||||||
errorCount,
|
errorCount,
|
||||||
availableTerminalHeight,
|
availableTerminalHeight,
|
||||||
@@ -2448,6 +2477,9 @@ Logging in with Google... Restarting Gemini CLI to continue.
|
|||||||
handleClearScreen,
|
handleClearScreen,
|
||||||
handleProQuotaChoice,
|
handleProQuotaChoice,
|
||||||
handleValidationChoice,
|
handleValidationChoice,
|
||||||
|
// G1 AI Credits handlers
|
||||||
|
handleOverageMenuChoice,
|
||||||
|
handleEmptyWalletChoice,
|
||||||
openSessionBrowser,
|
openSessionBrowser,
|
||||||
closeSessionBrowser,
|
closeSessionBrowser,
|
||||||
handleResumeSession,
|
handleResumeSession,
|
||||||
@@ -2534,6 +2566,8 @@ Logging in with Google... Restarting Gemini CLI to continue.
|
|||||||
handleClearScreen,
|
handleClearScreen,
|
||||||
handleProQuotaChoice,
|
handleProQuotaChoice,
|
||||||
handleValidationChoice,
|
handleValidationChoice,
|
||||||
|
handleOverageMenuChoice,
|
||||||
|
handleEmptyWalletChoice,
|
||||||
openSessionBrowser,
|
openSessionBrowser,
|
||||||
closeSessionBrowser,
|
closeSessionBrowser,
|
||||||
handleResumeSession,
|
handleResumeSession,
|
||||||
|
|||||||
@@ -39,11 +39,18 @@ describe('statsCommand', () => {
|
|||||||
mockContext.session.stats.sessionStartTime = startTime;
|
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');
|
if (!statsCommand.action) throw new Error('Command has no action');
|
||||||
|
|
||||||
// eslint-disable-next-line @typescript-eslint/no-floating-promises
|
mockContext.services.config = {
|
||||||
statsCommand.action(mockContext, '');
|
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(
|
const expectedDuration = formatDuration(
|
||||||
endTime.getTime() - startTime.getTime(),
|
endTime.getTime() - startTime.getTime(),
|
||||||
@@ -55,6 +62,7 @@ describe('statsCommand', () => {
|
|||||||
tier: undefined,
|
tier: undefined,
|
||||||
userEmail: 'mock@example.com',
|
userEmail: 'mock@example.com',
|
||||||
currentModel: undefined,
|
currentModel: undefined,
|
||||||
|
creditBalance: undefined,
|
||||||
});
|
});
|
||||||
});
|
});
|
||||||
|
|
||||||
@@ -78,6 +86,8 @@ describe('statsCommand', () => {
|
|||||||
getQuotaRemaining: mockGetQuotaRemaining,
|
getQuotaRemaining: mockGetQuotaRemaining,
|
||||||
getQuotaLimit: mockGetQuotaLimit,
|
getQuotaLimit: mockGetQuotaLimit,
|
||||||
getQuotaResetTime: mockGetQuotaResetTime,
|
getQuotaResetTime: mockGetQuotaResetTime,
|
||||||
|
getUserPaidTier: vi.fn(),
|
||||||
|
refreshAvailableCredits: vi.fn(),
|
||||||
} as unknown as Config;
|
} as unknown as Config;
|
||||||
|
|
||||||
await statsCommand.action(mockContext, '');
|
await statsCommand.action(mockContext, '');
|
||||||
|
|||||||
@@ -11,7 +11,10 @@ import type {
|
|||||||
} from '../types.js';
|
} from '../types.js';
|
||||||
import { MessageType } from '../types.js';
|
import { MessageType } from '../types.js';
|
||||||
import { formatDuration } from '../utils/formatters.js';
|
import { formatDuration } from '../utils/formatters.js';
|
||||||
import { UserAccountManager } from '@google/gemini-cli-core';
|
import {
|
||||||
|
UserAccountManager,
|
||||||
|
getG1CreditBalance,
|
||||||
|
} from '@google/gemini-cli-core';
|
||||||
import {
|
import {
|
||||||
type CommandContext,
|
type CommandContext,
|
||||||
type SlashCommand,
|
type SlashCommand,
|
||||||
@@ -27,8 +30,10 @@ function getUserIdentity(context: CommandContext) {
|
|||||||
const userEmail = cachedAccount ?? undefined;
|
const userEmail = cachedAccount ?? undefined;
|
||||||
|
|
||||||
const tier = context.services.config?.getUserTierName();
|
const tier = context.services.config?.getUserTierName();
|
||||||
|
const paidTier = context.services.config?.getUserPaidTier();
|
||||||
|
const creditBalance = getG1CreditBalance(paidTier) ?? undefined;
|
||||||
|
|
||||||
return { selectedAuthType, userEmail, tier };
|
return { selectedAuthType, userEmail, tier, creditBalance };
|
||||||
}
|
}
|
||||||
|
|
||||||
async function defaultSessionView(context: CommandContext) {
|
async function defaultSessionView(context: CommandContext) {
|
||||||
@@ -43,7 +48,8 @@ async function defaultSessionView(context: CommandContext) {
|
|||||||
}
|
}
|
||||||
const wallDuration = now.getTime() - sessionStartTime.getTime();
|
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 currentModel = context.services.config?.getModel();
|
||||||
|
|
||||||
const statsItem: HistoryItemStats = {
|
const statsItem: HistoryItemStats = {
|
||||||
@@ -53,10 +59,14 @@ async function defaultSessionView(context: CommandContext) {
|
|||||||
userEmail,
|
userEmail,
|
||||||
tier,
|
tier,
|
||||||
currentModel,
|
currentModel,
|
||||||
|
creditBalance,
|
||||||
};
|
};
|
||||||
|
|
||||||
if (context.services.config) {
|
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) {
|
if (quota) {
|
||||||
statsItem.quotas = quota;
|
statsItem.quotas = quota;
|
||||||
statsItem.pooledRemaining = context.services.config.getQuotaRemaining();
|
statsItem.pooledRemaining = context.services.config.getQuotaRemaining();
|
||||||
|
|||||||
@@ -80,6 +80,8 @@ describe('DialogManager', () => {
|
|||||||
stats: undefined,
|
stats: undefined,
|
||||||
proQuotaRequest: null,
|
proQuotaRequest: null,
|
||||||
validationRequest: null,
|
validationRequest: null,
|
||||||
|
overageMenuRequest: null,
|
||||||
|
emptyWalletRequest: null,
|
||||||
},
|
},
|
||||||
shouldShowIdePrompt: false,
|
shouldShowIdePrompt: false,
|
||||||
isFolderTrustDialogOpen: false,
|
isFolderTrustDialogOpen: false,
|
||||||
@@ -132,6 +134,8 @@ describe('DialogManager', () => {
|
|||||||
resolve: vi.fn(),
|
resolve: vi.fn(),
|
||||||
},
|
},
|
||||||
validationRequest: null,
|
validationRequest: null,
|
||||||
|
overageMenuRequest: null,
|
||||||
|
emptyWalletRequest: null,
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
'ProQuotaDialog',
|
'ProQuotaDialog',
|
||||||
|
|||||||
@@ -18,6 +18,8 @@ import { EditorSettingsDialog } from './EditorSettingsDialog.js';
|
|||||||
import { PrivacyNotice } from '../privacy/PrivacyNotice.js';
|
import { PrivacyNotice } from '../privacy/PrivacyNotice.js';
|
||||||
import { ProQuotaDialog } from './ProQuotaDialog.js';
|
import { ProQuotaDialog } from './ProQuotaDialog.js';
|
||||||
import { ValidationDialog } from './ValidationDialog.js';
|
import { ValidationDialog } from './ValidationDialog.js';
|
||||||
|
import { OverageMenuDialog } from './OverageMenuDialog.js';
|
||||||
|
import { EmptyWalletDialog } from './EmptyWalletDialog.js';
|
||||||
import { runExitCleanup } from '../../utils/cleanup.js';
|
import { runExitCleanup } from '../../utils/cleanup.js';
|
||||||
import { RELAUNCH_EXIT_CODE } from '../../utils/processUtils.js';
|
import { RELAUNCH_EXIT_CODE } from '../../utils/processUtils.js';
|
||||||
import { SessionBrowser } from './SessionBrowser.js';
|
import { SessionBrowser } from './SessionBrowser.js';
|
||||||
@@ -152,6 +154,28 @@ export const DialogManager = ({
|
|||||||
/>
|
/>
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
if (uiState.quota.overageMenuRequest) {
|
||||||
|
return (
|
||||||
|
<OverageMenuDialog
|
||||||
|
failedModel={uiState.quota.overageMenuRequest.failedModel}
|
||||||
|
fallbackModel={uiState.quota.overageMenuRequest.fallbackModel}
|
||||||
|
resetTime={uiState.quota.overageMenuRequest.resetTime}
|
||||||
|
creditBalance={uiState.quota.overageMenuRequest.creditBalance}
|
||||||
|
onChoice={uiActions.handleOverageMenuChoice}
|
||||||
|
/>
|
||||||
|
);
|
||||||
|
}
|
||||||
|
if (uiState.quota.emptyWalletRequest) {
|
||||||
|
return (
|
||||||
|
<EmptyWalletDialog
|
||||||
|
failedModel={uiState.quota.emptyWalletRequest.failedModel}
|
||||||
|
fallbackModel={uiState.quota.emptyWalletRequest.fallbackModel}
|
||||||
|
resetTime={uiState.quota.emptyWalletRequest.resetTime}
|
||||||
|
onGetCredits={uiState.quota.emptyWalletRequest.onGetCredits}
|
||||||
|
onChoice={uiActions.handleEmptyWalletChoice}
|
||||||
|
/>
|
||||||
|
);
|
||||||
|
}
|
||||||
if (uiState.shouldShowIdePrompt) {
|
if (uiState.shouldShowIdePrompt) {
|
||||||
return (
|
return (
|
||||||
<IdeIntegrationNudge
|
<IdeIntegrationNudge
|
||||||
|
|||||||
@@ -0,0 +1,218 @@
|
|||||||
|
/**
|
||||||
|
* @license
|
||||||
|
* Copyright 2026 Google LLC
|
||||||
|
* SPDX-License-Identifier: Apache-2.0
|
||||||
|
*/
|
||||||
|
|
||||||
|
import { renderWithProviders } from '../../test-utils/render.js';
|
||||||
|
import { waitFor } from '../../test-utils/async.js';
|
||||||
|
import { act } from 'react';
|
||||||
|
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
|
||||||
|
import { EmptyWalletDialog } from './EmptyWalletDialog.js';
|
||||||
|
|
||||||
|
const writeKey = (stdin: { write: (data: string) => void }, key: string) => {
|
||||||
|
act(() => {
|
||||||
|
stdin.write(key);
|
||||||
|
});
|
||||||
|
};
|
||||||
|
|
||||||
|
describe('EmptyWalletDialog', () => {
|
||||||
|
const mockOnChoice = vi.fn();
|
||||||
|
const mockOnGetCredits = vi.fn();
|
||||||
|
|
||||||
|
beforeEach(() => {
|
||||||
|
vi.clearAllMocks();
|
||||||
|
});
|
||||||
|
|
||||||
|
afterEach(() => {
|
||||||
|
vi.restoreAllMocks();
|
||||||
|
});
|
||||||
|
|
||||||
|
describe('rendering', () => {
|
||||||
|
it('should match snapshot with fallback available', async () => {
|
||||||
|
const { lastFrame, unmount, waitUntilReady } = renderWithProviders(
|
||||||
|
<EmptyWalletDialog
|
||||||
|
failedModel="gemini-2.5-pro"
|
||||||
|
fallbackModel="gemini-3-flash-preview"
|
||||||
|
resetTime="2:00 PM"
|
||||||
|
onChoice={mockOnChoice}
|
||||||
|
/>,
|
||||||
|
);
|
||||||
|
await waitUntilReady();
|
||||||
|
|
||||||
|
expect(lastFrame()).toMatchSnapshot();
|
||||||
|
unmount();
|
||||||
|
});
|
||||||
|
|
||||||
|
it('should match snapshot without fallback', async () => {
|
||||||
|
const { lastFrame, unmount, waitUntilReady } = renderWithProviders(
|
||||||
|
<EmptyWalletDialog
|
||||||
|
failedModel="gemini-2.5-pro"
|
||||||
|
onChoice={mockOnChoice}
|
||||||
|
/>,
|
||||||
|
);
|
||||||
|
await waitUntilReady();
|
||||||
|
|
||||||
|
expect(lastFrame()).toMatchSnapshot();
|
||||||
|
unmount();
|
||||||
|
});
|
||||||
|
|
||||||
|
it('should display the model name and usage limit message', async () => {
|
||||||
|
const { lastFrame, unmount, waitUntilReady } = renderWithProviders(
|
||||||
|
<EmptyWalletDialog
|
||||||
|
failedModel="gemini-2.5-pro"
|
||||||
|
onChoice={mockOnChoice}
|
||||||
|
/>,
|
||||||
|
);
|
||||||
|
await waitUntilReady();
|
||||||
|
|
||||||
|
const output = lastFrame() ?? '';
|
||||||
|
expect(output).toContain('gemini-2.5-pro');
|
||||||
|
expect(output).toContain('Usage limit reached');
|
||||||
|
unmount();
|
||||||
|
});
|
||||||
|
|
||||||
|
it('should display purchase prompt and credits update notice', async () => {
|
||||||
|
const { lastFrame, unmount, waitUntilReady } = renderWithProviders(
|
||||||
|
<EmptyWalletDialog
|
||||||
|
failedModel="gemini-2.5-pro"
|
||||||
|
onChoice={mockOnChoice}
|
||||||
|
/>,
|
||||||
|
);
|
||||||
|
await waitUntilReady();
|
||||||
|
|
||||||
|
const output = lastFrame() ?? '';
|
||||||
|
expect(output).toContain('purchase more AI Credits');
|
||||||
|
expect(output).toContain(
|
||||||
|
'Newly purchased AI credits may take a few minutes to update',
|
||||||
|
);
|
||||||
|
unmount();
|
||||||
|
});
|
||||||
|
|
||||||
|
it('should display reset time when provided', async () => {
|
||||||
|
const { lastFrame, unmount, waitUntilReady } = renderWithProviders(
|
||||||
|
<EmptyWalletDialog
|
||||||
|
failedModel="gemini-2.5-pro"
|
||||||
|
resetTime="3:45 PM"
|
||||||
|
onChoice={mockOnChoice}
|
||||||
|
/>,
|
||||||
|
);
|
||||||
|
await waitUntilReady();
|
||||||
|
|
||||||
|
const output = lastFrame() ?? '';
|
||||||
|
expect(output).toContain('3:45 PM');
|
||||||
|
expect(output).toContain('Access resets at');
|
||||||
|
unmount();
|
||||||
|
});
|
||||||
|
|
||||||
|
it('should not display reset time when not provided', async () => {
|
||||||
|
const { lastFrame, unmount, waitUntilReady } = renderWithProviders(
|
||||||
|
<EmptyWalletDialog
|
||||||
|
failedModel="gemini-2.5-pro"
|
||||||
|
onChoice={mockOnChoice}
|
||||||
|
/>,
|
||||||
|
);
|
||||||
|
await waitUntilReady();
|
||||||
|
|
||||||
|
const output = lastFrame() ?? '';
|
||||||
|
expect(output).not.toContain('Access resets at');
|
||||||
|
unmount();
|
||||||
|
});
|
||||||
|
|
||||||
|
it('should display slash command hints', async () => {
|
||||||
|
const { lastFrame, unmount, waitUntilReady } = renderWithProviders(
|
||||||
|
<EmptyWalletDialog
|
||||||
|
failedModel="gemini-2.5-pro"
|
||||||
|
onChoice={mockOnChoice}
|
||||||
|
/>,
|
||||||
|
);
|
||||||
|
await waitUntilReady();
|
||||||
|
|
||||||
|
const output = lastFrame() ?? '';
|
||||||
|
expect(output).toContain('/stats');
|
||||||
|
expect(output).toContain('/model');
|
||||||
|
expect(output).toContain('/auth');
|
||||||
|
unmount();
|
||||||
|
});
|
||||||
|
});
|
||||||
|
|
||||||
|
describe('onChoice handling', () => {
|
||||||
|
it('should call onGetCredits and onChoice when get_credits is selected', async () => {
|
||||||
|
// get_credits is the first item, so just press Enter
|
||||||
|
const { unmount, stdin, waitUntilReady } = renderWithProviders(
|
||||||
|
<EmptyWalletDialog
|
||||||
|
failedModel="gemini-2.5-pro"
|
||||||
|
onChoice={mockOnChoice}
|
||||||
|
onGetCredits={mockOnGetCredits}
|
||||||
|
/>,
|
||||||
|
);
|
||||||
|
await waitUntilReady();
|
||||||
|
|
||||||
|
writeKey(stdin, '\r');
|
||||||
|
|
||||||
|
await waitFor(() => {
|
||||||
|
expect(mockOnGetCredits).toHaveBeenCalled();
|
||||||
|
expect(mockOnChoice).toHaveBeenCalledWith('get_credits');
|
||||||
|
});
|
||||||
|
unmount();
|
||||||
|
});
|
||||||
|
|
||||||
|
it('should call onChoice without onGetCredits when onGetCredits is not provided', async () => {
|
||||||
|
const { unmount, stdin, waitUntilReady } = renderWithProviders(
|
||||||
|
<EmptyWalletDialog
|
||||||
|
failedModel="gemini-2.5-pro"
|
||||||
|
onChoice={mockOnChoice}
|
||||||
|
/>,
|
||||||
|
);
|
||||||
|
await waitUntilReady();
|
||||||
|
|
||||||
|
writeKey(stdin, '\r');
|
||||||
|
|
||||||
|
await waitFor(() => {
|
||||||
|
expect(mockOnChoice).toHaveBeenCalledWith('get_credits');
|
||||||
|
});
|
||||||
|
unmount();
|
||||||
|
});
|
||||||
|
|
||||||
|
it('should call onChoice with use_fallback when selected', async () => {
|
||||||
|
// With fallback: items are [get_credits, use_fallback, stop]
|
||||||
|
// use_fallback is the second item: Down + Enter
|
||||||
|
const { unmount, stdin, waitUntilReady } = renderWithProviders(
|
||||||
|
<EmptyWalletDialog
|
||||||
|
failedModel="gemini-2.5-pro"
|
||||||
|
fallbackModel="gemini-3-flash-preview"
|
||||||
|
onChoice={mockOnChoice}
|
||||||
|
/>,
|
||||||
|
);
|
||||||
|
await waitUntilReady();
|
||||||
|
|
||||||
|
writeKey(stdin, '\x1b[B'); // Down arrow
|
||||||
|
writeKey(stdin, '\r');
|
||||||
|
|
||||||
|
await waitFor(() => {
|
||||||
|
expect(mockOnChoice).toHaveBeenCalledWith('use_fallback');
|
||||||
|
});
|
||||||
|
unmount();
|
||||||
|
});
|
||||||
|
|
||||||
|
it('should call onChoice with stop when selected', async () => {
|
||||||
|
// Without fallback: items are [get_credits, stop]
|
||||||
|
// stop is the second item: Down + Enter
|
||||||
|
const { unmount, stdin, waitUntilReady } = renderWithProviders(
|
||||||
|
<EmptyWalletDialog
|
||||||
|
failedModel="gemini-2.5-pro"
|
||||||
|
onChoice={mockOnChoice}
|
||||||
|
/>,
|
||||||
|
);
|
||||||
|
await waitUntilReady();
|
||||||
|
|
||||||
|
writeKey(stdin, '\x1b[B'); // Down arrow
|
||||||
|
writeKey(stdin, '\r');
|
||||||
|
|
||||||
|
await waitFor(() => {
|
||||||
|
expect(mockOnChoice).toHaveBeenCalledWith('stop');
|
||||||
|
});
|
||||||
|
unmount();
|
||||||
|
});
|
||||||
|
});
|
||||||
|
});
|
||||||
@@ -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>
|
||||||
|
);
|
||||||
|
}
|
||||||
@@ -177,6 +177,8 @@ describe('<Footer />', () => {
|
|||||||
},
|
},
|
||||||
proQuotaRequest: null,
|
proQuotaRequest: null,
|
||||||
validationRequest: null,
|
validationRequest: null,
|
||||||
|
overageMenuRequest: null,
|
||||||
|
emptyWalletRequest: null,
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
@@ -203,6 +205,8 @@ describe('<Footer />', () => {
|
|||||||
},
|
},
|
||||||
proQuotaRequest: null,
|
proQuotaRequest: null,
|
||||||
validationRequest: null,
|
validationRequest: null,
|
||||||
|
overageMenuRequest: null,
|
||||||
|
emptyWalletRequest: null,
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
@@ -229,6 +233,8 @@ describe('<Footer />', () => {
|
|||||||
},
|
},
|
||||||
proQuotaRequest: null,
|
proQuotaRequest: null,
|
||||||
validationRequest: null,
|
validationRequest: null,
|
||||||
|
overageMenuRequest: null,
|
||||||
|
emptyWalletRequest: null,
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
|
|||||||
@@ -146,6 +146,7 @@ export const HistoryItemDisplay: React.FC<HistoryItemDisplayProps> = ({
|
|||||||
}
|
}
|
||||||
: undefined
|
: undefined
|
||||||
}
|
}
|
||||||
|
creditBalance={itemForDisplay.creditBalance}
|
||||||
/>
|
/>
|
||||||
)}
|
)}
|
||||||
{itemForDisplay.type === 'model_stats' && (
|
{itemForDisplay.type === 'model_stats' && (
|
||||||
|
|||||||
@@ -0,0 +1,228 @@
|
|||||||
|
/**
|
||||||
|
* @license
|
||||||
|
* Copyright 2026 Google LLC
|
||||||
|
* SPDX-License-Identifier: Apache-2.0
|
||||||
|
*/
|
||||||
|
|
||||||
|
import { renderWithProviders } from '../../test-utils/render.js';
|
||||||
|
import { waitFor } from '../../test-utils/async.js';
|
||||||
|
import { act } from 'react';
|
||||||
|
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
|
||||||
|
import { OverageMenuDialog } from './OverageMenuDialog.js';
|
||||||
|
|
||||||
|
const writeKey = (stdin: { write: (data: string) => void }, key: string) => {
|
||||||
|
act(() => {
|
||||||
|
stdin.write(key);
|
||||||
|
});
|
||||||
|
};
|
||||||
|
|
||||||
|
describe('OverageMenuDialog', () => {
|
||||||
|
const mockOnChoice = vi.fn();
|
||||||
|
|
||||||
|
beforeEach(() => {
|
||||||
|
vi.clearAllMocks();
|
||||||
|
});
|
||||||
|
|
||||||
|
afterEach(() => {
|
||||||
|
vi.restoreAllMocks();
|
||||||
|
});
|
||||||
|
|
||||||
|
describe('rendering', () => {
|
||||||
|
it('should match snapshot with fallback available', async () => {
|
||||||
|
const { lastFrame, unmount, waitUntilReady } = renderWithProviders(
|
||||||
|
<OverageMenuDialog
|
||||||
|
failedModel="gemini-2.5-pro"
|
||||||
|
fallbackModel="gemini-3-flash-preview"
|
||||||
|
resetTime="2:00 PM"
|
||||||
|
creditBalance={500}
|
||||||
|
onChoice={mockOnChoice}
|
||||||
|
/>,
|
||||||
|
);
|
||||||
|
await waitUntilReady();
|
||||||
|
|
||||||
|
expect(lastFrame()).toMatchSnapshot();
|
||||||
|
unmount();
|
||||||
|
});
|
||||||
|
|
||||||
|
it('should match snapshot without fallback', async () => {
|
||||||
|
const { lastFrame, unmount, waitUntilReady } = renderWithProviders(
|
||||||
|
<OverageMenuDialog
|
||||||
|
failedModel="gemini-2.5-pro"
|
||||||
|
creditBalance={500}
|
||||||
|
onChoice={mockOnChoice}
|
||||||
|
/>,
|
||||||
|
);
|
||||||
|
await waitUntilReady();
|
||||||
|
|
||||||
|
expect(lastFrame()).toMatchSnapshot();
|
||||||
|
unmount();
|
||||||
|
});
|
||||||
|
|
||||||
|
it('should display the credit balance', async () => {
|
||||||
|
const { lastFrame, unmount, waitUntilReady } = renderWithProviders(
|
||||||
|
<OverageMenuDialog
|
||||||
|
failedModel="gemini-2.5-pro"
|
||||||
|
creditBalance={200}
|
||||||
|
onChoice={mockOnChoice}
|
||||||
|
/>,
|
||||||
|
);
|
||||||
|
await waitUntilReady();
|
||||||
|
|
||||||
|
const output = lastFrame() ?? '';
|
||||||
|
expect(output).toContain('200');
|
||||||
|
expect(output).toContain('AI Credits available');
|
||||||
|
unmount();
|
||||||
|
});
|
||||||
|
|
||||||
|
it('should display the model name', async () => {
|
||||||
|
const { lastFrame, unmount, waitUntilReady } = renderWithProviders(
|
||||||
|
<OverageMenuDialog
|
||||||
|
failedModel="gemini-2.5-pro"
|
||||||
|
creditBalance={100}
|
||||||
|
onChoice={mockOnChoice}
|
||||||
|
/>,
|
||||||
|
);
|
||||||
|
await waitUntilReady();
|
||||||
|
|
||||||
|
const output = lastFrame() ?? '';
|
||||||
|
expect(output).toContain('gemini-2.5-pro');
|
||||||
|
expect(output).toContain('Usage limit reached');
|
||||||
|
unmount();
|
||||||
|
});
|
||||||
|
|
||||||
|
it('should display reset time when provided', async () => {
|
||||||
|
const { lastFrame, unmount, waitUntilReady } = renderWithProviders(
|
||||||
|
<OverageMenuDialog
|
||||||
|
failedModel="gemini-2.5-pro"
|
||||||
|
resetTime="3:45 PM"
|
||||||
|
creditBalance={100}
|
||||||
|
onChoice={mockOnChoice}
|
||||||
|
/>,
|
||||||
|
);
|
||||||
|
await waitUntilReady();
|
||||||
|
|
||||||
|
const output = lastFrame() ?? '';
|
||||||
|
expect(output).toContain('3:45 PM');
|
||||||
|
expect(output).toContain('Access resets at');
|
||||||
|
unmount();
|
||||||
|
});
|
||||||
|
|
||||||
|
it('should not display reset time when not provided', async () => {
|
||||||
|
const { lastFrame, unmount, waitUntilReady } = renderWithProviders(
|
||||||
|
<OverageMenuDialog
|
||||||
|
failedModel="gemini-2.5-pro"
|
||||||
|
creditBalance={100}
|
||||||
|
onChoice={mockOnChoice}
|
||||||
|
/>,
|
||||||
|
);
|
||||||
|
await waitUntilReady();
|
||||||
|
|
||||||
|
const output = lastFrame() ?? '';
|
||||||
|
expect(output).not.toContain('Access resets at');
|
||||||
|
unmount();
|
||||||
|
});
|
||||||
|
|
||||||
|
it('should display slash command hints', async () => {
|
||||||
|
const { lastFrame, unmount, waitUntilReady } = renderWithProviders(
|
||||||
|
<OverageMenuDialog
|
||||||
|
failedModel="gemini-2.5-pro"
|
||||||
|
creditBalance={100}
|
||||||
|
onChoice={mockOnChoice}
|
||||||
|
/>,
|
||||||
|
);
|
||||||
|
await waitUntilReady();
|
||||||
|
|
||||||
|
const output = lastFrame() ?? '';
|
||||||
|
expect(output).toContain('/stats');
|
||||||
|
expect(output).toContain('/model');
|
||||||
|
expect(output).toContain('/auth');
|
||||||
|
unmount();
|
||||||
|
});
|
||||||
|
});
|
||||||
|
|
||||||
|
describe('onChoice handling', () => {
|
||||||
|
it('should call onChoice with use_credits when selected', async () => {
|
||||||
|
// use_credits is the first item, so just press Enter
|
||||||
|
const { unmount, stdin, waitUntilReady } = renderWithProviders(
|
||||||
|
<OverageMenuDialog
|
||||||
|
failedModel="gemini-2.5-pro"
|
||||||
|
creditBalance={100}
|
||||||
|
onChoice={mockOnChoice}
|
||||||
|
/>,
|
||||||
|
);
|
||||||
|
await waitUntilReady();
|
||||||
|
|
||||||
|
writeKey(stdin, '\r');
|
||||||
|
|
||||||
|
await waitFor(() => {
|
||||||
|
expect(mockOnChoice).toHaveBeenCalledWith('use_credits');
|
||||||
|
});
|
||||||
|
unmount();
|
||||||
|
});
|
||||||
|
|
||||||
|
it('should call onChoice with manage when selected', async () => {
|
||||||
|
// manage is the second item: Down + Enter
|
||||||
|
const { unmount, stdin, waitUntilReady } = renderWithProviders(
|
||||||
|
<OverageMenuDialog
|
||||||
|
failedModel="gemini-2.5-pro"
|
||||||
|
creditBalance={100}
|
||||||
|
onChoice={mockOnChoice}
|
||||||
|
/>,
|
||||||
|
);
|
||||||
|
await waitUntilReady();
|
||||||
|
|
||||||
|
writeKey(stdin, '\x1b[B'); // Down arrow
|
||||||
|
writeKey(stdin, '\r');
|
||||||
|
|
||||||
|
await waitFor(() => {
|
||||||
|
expect(mockOnChoice).toHaveBeenCalledWith('manage');
|
||||||
|
});
|
||||||
|
unmount();
|
||||||
|
});
|
||||||
|
|
||||||
|
it('should call onChoice with use_fallback when selected', async () => {
|
||||||
|
// With fallback: items are [use_credits, manage, use_fallback, stop]
|
||||||
|
// use_fallback is the third item: Down x2 + Enter
|
||||||
|
const { unmount, stdin, waitUntilReady } = renderWithProviders(
|
||||||
|
<OverageMenuDialog
|
||||||
|
failedModel="gemini-2.5-pro"
|
||||||
|
fallbackModel="gemini-3-flash-preview"
|
||||||
|
creditBalance={100}
|
||||||
|
onChoice={mockOnChoice}
|
||||||
|
/>,
|
||||||
|
);
|
||||||
|
await waitUntilReady();
|
||||||
|
|
||||||
|
writeKey(stdin, '\x1b[B'); // Down arrow
|
||||||
|
writeKey(stdin, '\x1b[B'); // Down arrow
|
||||||
|
writeKey(stdin, '\r');
|
||||||
|
|
||||||
|
await waitFor(() => {
|
||||||
|
expect(mockOnChoice).toHaveBeenCalledWith('use_fallback');
|
||||||
|
});
|
||||||
|
unmount();
|
||||||
|
});
|
||||||
|
|
||||||
|
it('should call onChoice with stop when selected', async () => {
|
||||||
|
// Without fallback: items are [use_credits, manage, stop]
|
||||||
|
// stop is the third item: Down x2 + Enter
|
||||||
|
const { unmount, stdin, waitUntilReady } = renderWithProviders(
|
||||||
|
<OverageMenuDialog
|
||||||
|
failedModel="gemini-2.5-pro"
|
||||||
|
creditBalance={100}
|
||||||
|
onChoice={mockOnChoice}
|
||||||
|
/>,
|
||||||
|
);
|
||||||
|
await waitUntilReady();
|
||||||
|
|
||||||
|
writeKey(stdin, '\x1b[B'); // Down arrow
|
||||||
|
writeKey(stdin, '\x1b[B'); // Down arrow
|
||||||
|
writeKey(stdin, '\r');
|
||||||
|
|
||||||
|
await waitFor(() => {
|
||||||
|
expect(mockOnChoice).toHaveBeenCalledWith('stop');
|
||||||
|
});
|
||||||
|
unmount();
|
||||||
|
});
|
||||||
|
});
|
||||||
|
});
|
||||||
@@ -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>
|
||||||
|
);
|
||||||
|
}
|
||||||
@@ -395,6 +395,7 @@ interface StatsDisplayProps {
|
|||||||
tier?: string;
|
tier?: string;
|
||||||
currentModel?: string;
|
currentModel?: string;
|
||||||
quotaStats?: QuotaStats;
|
quotaStats?: QuotaStats;
|
||||||
|
creditBalance?: number;
|
||||||
}
|
}
|
||||||
|
|
||||||
export const StatsDisplay: React.FC<StatsDisplayProps> = ({
|
export const StatsDisplay: React.FC<StatsDisplayProps> = ({
|
||||||
@@ -407,6 +408,7 @@ export const StatsDisplay: React.FC<StatsDisplayProps> = ({
|
|||||||
tier,
|
tier,
|
||||||
currentModel,
|
currentModel,
|
||||||
quotaStats,
|
quotaStats,
|
||||||
|
creditBalance,
|
||||||
}) => {
|
}) => {
|
||||||
const { stats } = useSessionStats();
|
const { stats } = useSessionStats();
|
||||||
const { metrics } = stats;
|
const { metrics } = stats;
|
||||||
@@ -488,6 +490,17 @@ export const StatsDisplay: React.FC<StatsDisplayProps> = ({
|
|||||||
<Text color={theme.text.primary}>{tier}</Text>
|
<Text color={theme.text.primary}>{tier}</Text>
|
||||||
</StatRow>
|
</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:">
|
<StatRow title="Tool Calls:">
|
||||||
<Text color={theme.text.primary}>
|
<Text color={theme.text.primary}>
|
||||||
{tools.totalCalls} ({' '}
|
{tools.totalCalls} ({' '}
|
||||||
|
|||||||
@@ -0,0 +1,49 @@
|
|||||||
|
// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
|
||||||
|
|
||||||
|
exports[`EmptyWalletDialog > rendering > should match snapshot with fallback available 1`] = `
|
||||||
|
"╭──────────────────────────────────────────────────────────────────────────────────────────────────╮
|
||||||
|
│ │
|
||||||
|
│ Usage limit reached for gemini-2.5-pro. │
|
||||||
|
│ Access resets at 2:00 PM. │
|
||||||
|
│ /stats model for usage details │
|
||||||
|
│ /model to switch models. │
|
||||||
|
│ /auth to switch to API key. │
|
||||||
|
│ │
|
||||||
|
│ To continue using this model now, purchase more AI Credits. │
|
||||||
|
│ │
|
||||||
|
│ Newly purchased AI credits may take a few minutes to update. │
|
||||||
|
│ │
|
||||||
|
│ How would you like to proceed? │
|
||||||
|
│ │
|
||||||
|
│ │
|
||||||
|
│ ● 1. Get AI Credits - Open browser to purchase credits │
|
||||||
|
│ 2. Switch to gemini-3-flash-preview │
|
||||||
|
│ 3. Stop - Abort request │
|
||||||
|
│ │
|
||||||
|
│ │
|
||||||
|
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
|
||||||
|
"
|
||||||
|
`;
|
||||||
|
|
||||||
|
exports[`EmptyWalletDialog > rendering > should match snapshot without fallback 1`] = `
|
||||||
|
"╭──────────────────────────────────────────────────────────────────────────────────────────────────╮
|
||||||
|
│ │
|
||||||
|
│ Usage limit reached for gemini-2.5-pro. │
|
||||||
|
│ /stats model for usage details │
|
||||||
|
│ /model to switch models. │
|
||||||
|
│ /auth to switch to API key. │
|
||||||
|
│ │
|
||||||
|
│ To continue using this model now, purchase more AI Credits. │
|
||||||
|
│ │
|
||||||
|
│ Newly purchased AI credits may take a few minutes to update. │
|
||||||
|
│ │
|
||||||
|
│ How would you like to proceed? │
|
||||||
|
│ │
|
||||||
|
│ │
|
||||||
|
│ ● 1. Get AI Credits - Open browser to purchase credits │
|
||||||
|
│ 2. Stop - Abort request │
|
||||||
|
│ │
|
||||||
|
│ │
|
||||||
|
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
|
||||||
|
"
|
||||||
|
`;
|
||||||
@@ -0,0 +1,47 @@
|
|||||||
|
// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
|
||||||
|
|
||||||
|
exports[`OverageMenuDialog > rendering > should match snapshot with fallback available 1`] = `
|
||||||
|
"╭──────────────────────────────────────────────────────────────────────────────────────────────────╮
|
||||||
|
│ │
|
||||||
|
│ Usage limit reached for gemini-2.5-pro. │
|
||||||
|
│ Access resets at 2:00 PM. │
|
||||||
|
│ /stats model for usage details │
|
||||||
|
│ /model to switch models. │
|
||||||
|
│ /auth to switch to API key. │
|
||||||
|
│ │
|
||||||
|
│ You have 500 AI Credits available. │
|
||||||
|
│ │
|
||||||
|
│ How would you like to proceed? │
|
||||||
|
│ │
|
||||||
|
│ │
|
||||||
|
│ ● 1. Use AI Credits - Continue this request (Overage) │
|
||||||
|
│ 2. Manage - View balance and purchase more credits │
|
||||||
|
│ 3. Switch to gemini-3-flash-preview │
|
||||||
|
│ 4. Stop - Abort request │
|
||||||
|
│ │
|
||||||
|
│ │
|
||||||
|
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
|
||||||
|
"
|
||||||
|
`;
|
||||||
|
|
||||||
|
exports[`OverageMenuDialog > rendering > should match snapshot without fallback 1`] = `
|
||||||
|
"╭──────────────────────────────────────────────────────────────────────────────────────────────────╮
|
||||||
|
│ │
|
||||||
|
│ Usage limit reached for gemini-2.5-pro. │
|
||||||
|
│ /stats model for usage details │
|
||||||
|
│ /model to switch models. │
|
||||||
|
│ /auth to switch to API key. │
|
||||||
|
│ │
|
||||||
|
│ You have 500 AI Credits available. │
|
||||||
|
│ │
|
||||||
|
│ How would you like to proceed? │
|
||||||
|
│ │
|
||||||
|
│ │
|
||||||
|
│ ● 1. Use AI Credits - Continue this request (Overage) │
|
||||||
|
│ 2. Manage - View balance and purchase more credits │
|
||||||
|
│ 3. Stop - Abort request │
|
||||||
|
│ │
|
||||||
|
│ │
|
||||||
|
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
|
||||||
|
"
|
||||||
|
`;
|
||||||
@@ -18,6 +18,7 @@ import type { AuthState } from '../types.js';
|
|||||||
import { type PermissionsDialogProps } from '../components/PermissionsModifyTrustDialog.js';
|
import { type PermissionsDialogProps } from '../components/PermissionsModifyTrustDialog.js';
|
||||||
import type { SessionInfo } from '../../utils/sessionUtils.js';
|
import type { SessionInfo } from '../../utils/sessionUtils.js';
|
||||||
import { type NewAgentsChoice } from '../components/NewAgentsNotification.js';
|
import { type NewAgentsChoice } from '../components/NewAgentsNotification.js';
|
||||||
|
import type { OverageMenuIntent, EmptyWalletIntent } from './UIStateContext.js';
|
||||||
|
|
||||||
export interface UIActions {
|
export interface UIActions {
|
||||||
handleThemeSelect: (
|
handleThemeSelect: (
|
||||||
@@ -62,6 +63,8 @@ export interface UIActions {
|
|||||||
choice: 'retry_later' | 'retry_once' | 'retry_always' | 'upgrade',
|
choice: 'retry_later' | 'retry_once' | 'retry_always' | 'upgrade',
|
||||||
) => void;
|
) => void;
|
||||||
handleValidationChoice: (choice: 'verify' | 'change_auth' | 'cancel') => void;
|
handleValidationChoice: (choice: 'verify' | 'change_auth' | 'cancel') => void;
|
||||||
|
handleOverageMenuChoice: (choice: OverageMenuIntent) => void;
|
||||||
|
handleEmptyWalletChoice: (choice: EmptyWalletIntent) => void;
|
||||||
openSessionBrowser: () => void;
|
openSessionBrowser: () => void;
|
||||||
closeSessionBrowser: () => void;
|
closeSessionBrowser: () => void;
|
||||||
handleResumeSession: (session: SessionInfo) => Promise<void>;
|
handleResumeSession: (session: SessionInfo) => Promise<void>;
|
||||||
|
|||||||
@@ -54,6 +54,34 @@ export interface ValidationDialogRequest {
|
|||||||
resolve: (intent: ValidationIntent) => void;
|
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 UseHistoryManagerReturn } from '../hooks/useHistoryManager.js';
|
||||||
import { type RestartReason } from '../hooks/useIdeTrustListener.js';
|
import { type RestartReason } from '../hooks/useIdeTrustListener.js';
|
||||||
import type { TerminalBackgroundColor } from '../utils/terminalCapabilityManager.js';
|
import type { TerminalBackgroundColor } from '../utils/terminalCapabilityManager.js';
|
||||||
@@ -64,6 +92,9 @@ export interface QuotaState {
|
|||||||
stats: QuotaStats | undefined;
|
stats: QuotaStats | undefined;
|
||||||
proQuotaRequest: ProQuotaDialogRequest | null;
|
proQuotaRequest: ProQuotaDialogRequest | null;
|
||||||
validationRequest: ValidationDialogRequest | null;
|
validationRequest: ValidationDialogRequest | null;
|
||||||
|
// G1 AI Credits overage flow
|
||||||
|
overageMenuRequest: OverageMenuDialogRequest | null;
|
||||||
|
emptyWalletRequest: EmptyWalletDialogRequest | null;
|
||||||
}
|
}
|
||||||
|
|
||||||
export interface UIState {
|
export interface UIState {
|
||||||
|
|||||||
@@ -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);
|
||||||
|
});
|
||||||
|
});
|
||||||
@@ -0,0 +1,290 @@
|
|||||||
|
/**
|
||||||
|
* @license
|
||||||
|
* Copyright 2026 Google LLC
|
||||||
|
* SPDX-License-Identifier: Apache-2.0
|
||||||
|
*/
|
||||||
|
|
||||||
|
import {
|
||||||
|
type Config,
|
||||||
|
type FallbackIntent,
|
||||||
|
type GeminiUserTier,
|
||||||
|
type OverageOption,
|
||||||
|
getG1CreditBalance,
|
||||||
|
shouldAutoUseCredits,
|
||||||
|
shouldShowOverageMenu,
|
||||||
|
shouldShowEmptyWalletMenu,
|
||||||
|
openBrowserSecurely,
|
||||||
|
logBillingEvent,
|
||||||
|
OverageMenuShownEvent,
|
||||||
|
OverageOptionSelectedEvent,
|
||||||
|
EmptyWalletMenuShownEvent,
|
||||||
|
CreditPurchaseClickEvent,
|
||||||
|
buildG1Url,
|
||||||
|
G1_UTM_CAMPAIGNS,
|
||||||
|
UserAccountManager,
|
||||||
|
recordOverageOptionSelected,
|
||||||
|
recordCreditPurchaseClick,
|
||||||
|
} from '@google/gemini-cli-core';
|
||||||
|
import { MessageType } from '../types.js';
|
||||||
|
import type { UseHistoryManagerReturn } from './useHistoryManager.js';
|
||||||
|
import type {
|
||||||
|
OverageMenuIntent,
|
||||||
|
EmptyWalletIntent,
|
||||||
|
EmptyWalletDialogRequest,
|
||||||
|
} from '../contexts/UIStateContext.js';
|
||||||
|
|
||||||
|
interface CreditsFlowArgs {
|
||||||
|
config: Config;
|
||||||
|
paidTier: GeminiUserTier;
|
||||||
|
overageStrategy: 'ask' | 'always' | 'never';
|
||||||
|
failedModel: string;
|
||||||
|
fallbackModel: string;
|
||||||
|
usageLimitReachedModel: string;
|
||||||
|
resetTime: string | undefined;
|
||||||
|
historyManager: UseHistoryManagerReturn;
|
||||||
|
setModelSwitchedFromQuotaError: (value: boolean) => void;
|
||||||
|
isDialogPending: React.MutableRefObject<boolean>;
|
||||||
|
setOverageMenuRequest: (
|
||||||
|
req: {
|
||||||
|
failedModel: string;
|
||||||
|
fallbackModel: string;
|
||||||
|
resetTime: string | undefined;
|
||||||
|
creditBalance: number;
|
||||||
|
resolve: (intent: OverageMenuIntent) => void;
|
||||||
|
} | null,
|
||||||
|
) => void;
|
||||||
|
setEmptyWalletRequest: (req: EmptyWalletDialogRequest | null) => void;
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Handles the G1 AI Credits flow when a quota error occurs.
|
||||||
|
* Returns a FallbackIntent if the credits flow handled the error,
|
||||||
|
* or null to fall through to the default ProQuotaDialog.
|
||||||
|
*/
|
||||||
|
export async function handleCreditsFlow(
|
||||||
|
args: CreditsFlowArgs,
|
||||||
|
): Promise<FallbackIntent | null> {
|
||||||
|
const creditBalance = getG1CreditBalance(args.paidTier);
|
||||||
|
|
||||||
|
// creditBalance is null when user is not eligible for G1 credits.
|
||||||
|
if (creditBalance == null) {
|
||||||
|
return null;
|
||||||
|
}
|
||||||
|
|
||||||
|
const { overageStrategy } = args;
|
||||||
|
|
||||||
|
// If credits are already auto-enabled (strategy='always'), the request
|
||||||
|
// that just failed already included enabledCreditTypes — credits didn't
|
||||||
|
// help. Fall through to ProQuotaDialog which offers the Flash downgrade.
|
||||||
|
if (shouldAutoUseCredits(overageStrategy, creditBalance)) {
|
||||||
|
return null;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Show overage menu when strategy is 'ask' and credits > 0
|
||||||
|
if (shouldShowOverageMenu(overageStrategy, creditBalance)) {
|
||||||
|
return handleOverageMenu(args, creditBalance);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Show empty wallet when credits === 0 and strategy isn't 'never'
|
||||||
|
if (shouldShowEmptyWalletMenu(overageStrategy, creditBalance)) {
|
||||||
|
return handleEmptyWalletMenu(args);
|
||||||
|
}
|
||||||
|
|
||||||
|
return null;
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// Overage menu flow
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
async function handleOverageMenu(
|
||||||
|
args: CreditsFlowArgs,
|
||||||
|
creditBalance: number,
|
||||||
|
): Promise<FallbackIntent> {
|
||||||
|
const {
|
||||||
|
config,
|
||||||
|
fallbackModel,
|
||||||
|
usageLimitReachedModel,
|
||||||
|
overageStrategy,
|
||||||
|
resetTime,
|
||||||
|
isDialogPending,
|
||||||
|
setOverageMenuRequest,
|
||||||
|
setModelSwitchedFromQuotaError,
|
||||||
|
historyManager,
|
||||||
|
} = args;
|
||||||
|
|
||||||
|
logBillingEvent(
|
||||||
|
config,
|
||||||
|
new OverageMenuShownEvent(
|
||||||
|
usageLimitReachedModel,
|
||||||
|
creditBalance,
|
||||||
|
overageStrategy,
|
||||||
|
),
|
||||||
|
);
|
||||||
|
|
||||||
|
if (isDialogPending.current) {
|
||||||
|
return 'stop';
|
||||||
|
}
|
||||||
|
isDialogPending.current = true;
|
||||||
|
|
||||||
|
setModelSwitchedFromQuotaError(true);
|
||||||
|
config.setQuotaErrorOccurred(true);
|
||||||
|
|
||||||
|
const overageIntent = await new Promise<OverageMenuIntent>((resolve) => {
|
||||||
|
setOverageMenuRequest({
|
||||||
|
failedModel: usageLimitReachedModel,
|
||||||
|
fallbackModel,
|
||||||
|
resetTime,
|
||||||
|
creditBalance,
|
||||||
|
resolve,
|
||||||
|
});
|
||||||
|
});
|
||||||
|
|
||||||
|
setOverageMenuRequest(null);
|
||||||
|
isDialogPending.current = false;
|
||||||
|
|
||||||
|
logOverageOptionSelected(
|
||||||
|
config,
|
||||||
|
usageLimitReachedModel,
|
||||||
|
overageIntent,
|
||||||
|
creditBalance,
|
||||||
|
);
|
||||||
|
|
||||||
|
switch (overageIntent) {
|
||||||
|
case 'use_credits':
|
||||||
|
setModelSwitchedFromQuotaError(false);
|
||||||
|
config.setQuotaErrorOccurred(false);
|
||||||
|
config.setOverageStrategy('always');
|
||||||
|
historyManager.addItem(
|
||||||
|
{
|
||||||
|
type: MessageType.INFO,
|
||||||
|
text: `Using AI Credits for this request.`,
|
||||||
|
},
|
||||||
|
Date.now(),
|
||||||
|
);
|
||||||
|
return 'retry_with_credits';
|
||||||
|
|
||||||
|
case 'use_fallback':
|
||||||
|
return 'retry_always';
|
||||||
|
|
||||||
|
case 'manage':
|
||||||
|
logCreditPurchaseClick(config, 'manage', usageLimitReachedModel);
|
||||||
|
await openG1Url('activity', G1_UTM_CAMPAIGNS.MANAGE_ACTIVITY);
|
||||||
|
return 'stop';
|
||||||
|
|
||||||
|
case 'stop':
|
||||||
|
default:
|
||||||
|
return 'stop';
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// Empty wallet flow
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
async function handleEmptyWalletMenu(
|
||||||
|
args: CreditsFlowArgs,
|
||||||
|
): Promise<FallbackIntent> {
|
||||||
|
const {
|
||||||
|
config,
|
||||||
|
fallbackModel,
|
||||||
|
usageLimitReachedModel,
|
||||||
|
resetTime,
|
||||||
|
isDialogPending,
|
||||||
|
setEmptyWalletRequest,
|
||||||
|
setModelSwitchedFromQuotaError,
|
||||||
|
} = args;
|
||||||
|
|
||||||
|
logBillingEvent(
|
||||||
|
config,
|
||||||
|
new EmptyWalletMenuShownEvent(usageLimitReachedModel),
|
||||||
|
);
|
||||||
|
|
||||||
|
if (isDialogPending.current) {
|
||||||
|
return 'stop';
|
||||||
|
}
|
||||||
|
isDialogPending.current = true;
|
||||||
|
|
||||||
|
setModelSwitchedFromQuotaError(true);
|
||||||
|
config.setQuotaErrorOccurred(true);
|
||||||
|
|
||||||
|
const emptyWalletIntent = await new Promise<EmptyWalletIntent>((resolve) => {
|
||||||
|
setEmptyWalletRequest({
|
||||||
|
failedModel: usageLimitReachedModel,
|
||||||
|
fallbackModel,
|
||||||
|
resetTime,
|
||||||
|
onGetCredits: () => {
|
||||||
|
logCreditPurchaseClick(
|
||||||
|
config,
|
||||||
|
'empty_wallet_menu',
|
||||||
|
usageLimitReachedModel,
|
||||||
|
);
|
||||||
|
void openG1Url('credits', G1_UTM_CAMPAIGNS.EMPTY_WALLET_ADD_CREDITS);
|
||||||
|
},
|
||||||
|
resolve,
|
||||||
|
});
|
||||||
|
});
|
||||||
|
|
||||||
|
setEmptyWalletRequest(null);
|
||||||
|
isDialogPending.current = false;
|
||||||
|
|
||||||
|
switch (emptyWalletIntent) {
|
||||||
|
case 'get_credits':
|
||||||
|
args.historyManager.addItem(
|
||||||
|
{
|
||||||
|
type: MessageType.INFO,
|
||||||
|
text: 'Newly purchased AI credits may take a few minutes to update. Run /stats to check your balance.',
|
||||||
|
},
|
||||||
|
Date.now(),
|
||||||
|
);
|
||||||
|
return 'stop';
|
||||||
|
|
||||||
|
case 'use_fallback':
|
||||||
|
return 'retry_always';
|
||||||
|
|
||||||
|
case 'stop':
|
||||||
|
default:
|
||||||
|
return 'stop';
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// Telemetry helpers
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
function logOverageOptionSelected(
|
||||||
|
config: Config,
|
||||||
|
model: string,
|
||||||
|
option: OverageOption,
|
||||||
|
creditBalance: number,
|
||||||
|
): void {
|
||||||
|
logBillingEvent(
|
||||||
|
config,
|
||||||
|
new OverageOptionSelectedEvent(model, option, creditBalance),
|
||||||
|
);
|
||||||
|
recordOverageOptionSelected(config, {
|
||||||
|
selected_option: option,
|
||||||
|
model,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
function logCreditPurchaseClick(
|
||||||
|
config: Config,
|
||||||
|
source: 'overage_menu' | 'empty_wallet_menu' | 'manage',
|
||||||
|
model: string,
|
||||||
|
): void {
|
||||||
|
logBillingEvent(config, new CreditPurchaseClickEvent(source, model));
|
||||||
|
recordCreditPurchaseClick(config, { source, model });
|
||||||
|
}
|
||||||
|
|
||||||
|
async function openG1Url(
|
||||||
|
path: 'activity' | 'credits',
|
||||||
|
campaign: string,
|
||||||
|
): Promise<void> {
|
||||||
|
try {
|
||||||
|
const userEmail = new UserAccountManager().getCachedGoogleAccount() ?? '';
|
||||||
|
await openBrowserSecurely(buildG1Url(path, userEmail, campaign));
|
||||||
|
} catch {
|
||||||
|
// Ignore browser open errors
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -14,7 +14,8 @@ import {
|
|||||||
type Mock,
|
type Mock,
|
||||||
} from 'vitest';
|
} from 'vitest';
|
||||||
import { act } from 'react';
|
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 {
|
import {
|
||||||
type Config,
|
type Config,
|
||||||
type FallbackModelHandler,
|
type FallbackModelHandler,
|
||||||
@@ -29,6 +30,12 @@ import {
|
|||||||
ModelNotFoundError,
|
ModelNotFoundError,
|
||||||
DEFAULT_GEMINI_MODEL,
|
DEFAULT_GEMINI_MODEL,
|
||||||
DEFAULT_GEMINI_FLASH_MODEL,
|
DEFAULT_GEMINI_FLASH_MODEL,
|
||||||
|
getG1CreditBalance,
|
||||||
|
shouldAutoUseCredits,
|
||||||
|
shouldShowOverageMenu,
|
||||||
|
shouldShowEmptyWalletMenu,
|
||||||
|
logBillingEvent,
|
||||||
|
G1_CREDIT_TYPE,
|
||||||
} from '@google/gemini-cli-core';
|
} from '@google/gemini-cli-core';
|
||||||
import { useQuotaAndFallback } from './useQuotaAndFallback.js';
|
import { useQuotaAndFallback } from './useQuotaAndFallback.js';
|
||||||
import type { UseHistoryManagerReturn } from './useHistoryManager.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
|
// Use a type alias for SpyInstance as it's not directly exported
|
||||||
type SpyInstance = ReturnType<typeof vi.spyOn>;
|
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', () => {
|
describe('useQuotaAndFallback', () => {
|
||||||
let mockConfig: Config;
|
let mockConfig: Config;
|
||||||
let mockHistoryManager: UseHistoryManagerReturn;
|
let mockHistoryManager: UseHistoryManagerReturn;
|
||||||
@@ -74,10 +94,16 @@ describe('useQuotaAndFallback', () => {
|
|||||||
vi.spyOn(mockConfig, 'setModel');
|
vi.spyOn(mockConfig, 'setModel');
|
||||||
vi.spyOn(mockConfig, 'setActiveModel');
|
vi.spyOn(mockConfig, 'setActiveModel');
|
||||||
vi.spyOn(mockConfig, 'activateFallbackMode');
|
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(() => {
|
afterEach(() => {
|
||||||
vi.clearAllMocks();
|
vi.restoreAllMocks();
|
||||||
});
|
});
|
||||||
|
|
||||||
it('should register a fallback handler on initialization', () => {
|
it('should register a fallback handler on initialization', () => {
|
||||||
@@ -88,6 +114,8 @@ describe('useQuotaAndFallback', () => {
|
|||||||
userTier: UserTierId.FREE,
|
userTier: UserTierId.FREE,
|
||||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||||
onShowAuthSelection: mockOnShowAuthSelection,
|
onShowAuthSelection: mockOnShowAuthSelection,
|
||||||
|
paidTier: null,
|
||||||
|
settings: mockSettings,
|
||||||
}),
|
}),
|
||||||
);
|
);
|
||||||
|
|
||||||
@@ -109,6 +137,8 @@ describe('useQuotaAndFallback', () => {
|
|||||||
userTier: UserTierId.FREE,
|
userTier: UserTierId.FREE,
|
||||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||||
onShowAuthSelection: mockOnShowAuthSelection,
|
onShowAuthSelection: mockOnShowAuthSelection,
|
||||||
|
paidTier: null,
|
||||||
|
settings: mockSettings,
|
||||||
}),
|
}),
|
||||||
);
|
);
|
||||||
|
|
||||||
@@ -140,6 +170,8 @@ describe('useQuotaAndFallback', () => {
|
|||||||
userTier: UserTierId.FREE,
|
userTier: UserTierId.FREE,
|
||||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||||
onShowAuthSelection: mockOnShowAuthSelection,
|
onShowAuthSelection: mockOnShowAuthSelection,
|
||||||
|
paidTier: null,
|
||||||
|
settings: mockSettings,
|
||||||
}),
|
}),
|
||||||
);
|
);
|
||||||
|
|
||||||
@@ -193,6 +225,8 @@ describe('useQuotaAndFallback', () => {
|
|||||||
userTier: UserTierId.FREE,
|
userTier: UserTierId.FREE,
|
||||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||||
onShowAuthSelection: mockOnShowAuthSelection,
|
onShowAuthSelection: mockOnShowAuthSelection,
|
||||||
|
paidTier: null,
|
||||||
|
settings: mockSettings,
|
||||||
}),
|
}),
|
||||||
);
|
);
|
||||||
|
|
||||||
@@ -232,6 +266,8 @@ describe('useQuotaAndFallback', () => {
|
|||||||
userTier: UserTierId.FREE,
|
userTier: UserTierId.FREE,
|
||||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||||
onShowAuthSelection: mockOnShowAuthSelection,
|
onShowAuthSelection: mockOnShowAuthSelection,
|
||||||
|
paidTier: null,
|
||||||
|
settings: mockSettings,
|
||||||
}),
|
}),
|
||||||
);
|
);
|
||||||
|
|
||||||
@@ -264,6 +300,8 @@ describe('useQuotaAndFallback', () => {
|
|||||||
userTier: UserTierId.FREE,
|
userTier: UserTierId.FREE,
|
||||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||||
onShowAuthSelection: mockOnShowAuthSelection,
|
onShowAuthSelection: mockOnShowAuthSelection,
|
||||||
|
paidTier: null,
|
||||||
|
settings: mockSettings,
|
||||||
}),
|
}),
|
||||||
);
|
);
|
||||||
|
|
||||||
@@ -330,6 +368,8 @@ describe('useQuotaAndFallback', () => {
|
|||||||
setModelSwitchedFromQuotaError:
|
setModelSwitchedFromQuotaError:
|
||||||
mockSetModelSwitchedFromQuotaError,
|
mockSetModelSwitchedFromQuotaError,
|
||||||
onShowAuthSelection: mockOnShowAuthSelection,
|
onShowAuthSelection: mockOnShowAuthSelection,
|
||||||
|
paidTier: null,
|
||||||
|
settings: mockSettings,
|
||||||
}),
|
}),
|
||||||
);
|
);
|
||||||
|
|
||||||
@@ -385,6 +425,8 @@ describe('useQuotaAndFallback', () => {
|
|||||||
userTier: UserTierId.FREE,
|
userTier: UserTierId.FREE,
|
||||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||||
onShowAuthSelection: mockOnShowAuthSelection,
|
onShowAuthSelection: mockOnShowAuthSelection,
|
||||||
|
paidTier: null,
|
||||||
|
settings: mockSettings,
|
||||||
}),
|
}),
|
||||||
);
|
);
|
||||||
|
|
||||||
@@ -430,6 +472,8 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
|
|||||||
userTier: UserTierId.FREE,
|
userTier: UserTierId.FREE,
|
||||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||||
onShowAuthSelection: mockOnShowAuthSelection,
|
onShowAuthSelection: mockOnShowAuthSelection,
|
||||||
|
paidTier: null,
|
||||||
|
settings: mockSettings,
|
||||||
}),
|
}),
|
||||||
);
|
);
|
||||||
|
|
||||||
@@ -464,6 +508,243 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
|
|||||||
});
|
});
|
||||||
});
|
});
|
||||||
|
|
||||||
|
describe('G1 AI Credits Flow', () => {
|
||||||
|
const mockPaidTier = {
|
||||||
|
id: UserTierId.STANDARD,
|
||||||
|
userTier: UserTierId.STANDARD,
|
||||||
|
availableCredits: [
|
||||||
|
{
|
||||||
|
creditType: G1_CREDIT_TYPE,
|
||||||
|
creditAmount: '100',
|
||||||
|
},
|
||||||
|
],
|
||||||
|
};
|
||||||
|
|
||||||
|
beforeEach(() => {
|
||||||
|
// Default to having credits
|
||||||
|
vi.mocked(getG1CreditBalance).mockReturnValue(100);
|
||||||
|
});
|
||||||
|
|
||||||
|
it('should fall through to ProQuotaDialog if credits are already active (strategy=always)', async () => {
|
||||||
|
// If shouldAutoUseCredits is true, credits were already active on the
|
||||||
|
// failed request — they didn't help. Fall through to ProQuotaDialog
|
||||||
|
// so the user can downgrade to Flash instead of retrying infinitely.
|
||||||
|
vi.mocked(shouldAutoUseCredits).mockReturnValue(true);
|
||||||
|
|
||||||
|
const { result } = renderHook(() =>
|
||||||
|
useQuotaAndFallback({
|
||||||
|
config: mockConfig,
|
||||||
|
historyManager: mockHistoryManager,
|
||||||
|
userTier: UserTierId.STANDARD,
|
||||||
|
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||||
|
onShowAuthSelection: mockOnShowAuthSelection,
|
||||||
|
paidTier: mockPaidTier,
|
||||||
|
settings: mockSettings,
|
||||||
|
}),
|
||||||
|
);
|
||||||
|
|
||||||
|
const handler = setFallbackHandlerSpy.mock
|
||||||
|
.calls[0][0] as FallbackModelHandler;
|
||||||
|
|
||||||
|
const error = new TerminalQuotaError(
|
||||||
|
'pro quota',
|
||||||
|
mockGoogleApiError,
|
||||||
|
1000 * 60 * 5,
|
||||||
|
);
|
||||||
|
|
||||||
|
const intentPromise = handler(
|
||||||
|
PREVIEW_GEMINI_MODEL,
|
||||||
|
'gemini-flash',
|
||||||
|
error,
|
||||||
|
);
|
||||||
|
|
||||||
|
// Since credits didn't help, the ProQuotaDialog should be shown
|
||||||
|
await waitFor(() => {
|
||||||
|
expect(result.current.proQuotaRequest).not.toBeNull();
|
||||||
|
});
|
||||||
|
|
||||||
|
// Resolve it to verify the flow completes
|
||||||
|
act(() => {
|
||||||
|
result.current.handleProQuotaChoice('stop');
|
||||||
|
});
|
||||||
|
|
||||||
|
const intent = await intentPromise;
|
||||||
|
expect(intent).toBe('stop');
|
||||||
|
});
|
||||||
|
|
||||||
|
it('should show overage menu if balance > 0 and not auto-using', async () => {
|
||||||
|
vi.mocked(shouldAutoUseCredits).mockReturnValue(false);
|
||||||
|
vi.mocked(shouldShowOverageMenu).mockReturnValue(true);
|
||||||
|
|
||||||
|
const { result } = renderHook(() =>
|
||||||
|
useQuotaAndFallback({
|
||||||
|
config: mockConfig,
|
||||||
|
historyManager: mockHistoryManager,
|
||||||
|
userTier: UserTierId.STANDARD,
|
||||||
|
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||||
|
onShowAuthSelection: mockOnShowAuthSelection,
|
||||||
|
paidTier: mockPaidTier,
|
||||||
|
settings: mockSettings,
|
||||||
|
}),
|
||||||
|
);
|
||||||
|
|
||||||
|
const handler = setFallbackHandlerSpy.mock
|
||||||
|
.calls[0][0] as FallbackModelHandler;
|
||||||
|
|
||||||
|
let promise: Promise<FallbackIntent | null>;
|
||||||
|
act(() => {
|
||||||
|
promise = handler(
|
||||||
|
PREVIEW_GEMINI_MODEL,
|
||||||
|
'gemini-flash',
|
||||||
|
new TerminalQuotaError('pro quota', mockGoogleApiError),
|
||||||
|
);
|
||||||
|
});
|
||||||
|
|
||||||
|
expect(result.current.overageMenuRequest).not.toBeNull();
|
||||||
|
expect(result.current.overageMenuRequest?.creditBalance).toBe(100);
|
||||||
|
expect(logBillingEvent).toHaveBeenCalled();
|
||||||
|
|
||||||
|
// Simulate choosing "Use Credits"
|
||||||
|
await act(async () => {
|
||||||
|
result.current.handleOverageMenuChoice('use_credits');
|
||||||
|
await promise!;
|
||||||
|
});
|
||||||
|
|
||||||
|
const intent = await promise!;
|
||||||
|
expect(intent).toBe('retry_with_credits');
|
||||||
|
});
|
||||||
|
|
||||||
|
it('should handle use_fallback from overage menu', async () => {
|
||||||
|
vi.mocked(shouldAutoUseCredits).mockReturnValue(false);
|
||||||
|
vi.mocked(shouldShowOverageMenu).mockReturnValue(true);
|
||||||
|
|
||||||
|
const { result } = renderHook(() =>
|
||||||
|
useQuotaAndFallback({
|
||||||
|
config: mockConfig,
|
||||||
|
historyManager: mockHistoryManager,
|
||||||
|
userTier: UserTierId.STANDARD,
|
||||||
|
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||||
|
onShowAuthSelection: mockOnShowAuthSelection,
|
||||||
|
paidTier: mockPaidTier,
|
||||||
|
settings: mockSettings,
|
||||||
|
}),
|
||||||
|
);
|
||||||
|
|
||||||
|
const handler = setFallbackHandlerSpy.mock
|
||||||
|
.calls[0][0] as FallbackModelHandler;
|
||||||
|
|
||||||
|
let promise: Promise<FallbackIntent | null>;
|
||||||
|
act(() => {
|
||||||
|
promise = handler(
|
||||||
|
PREVIEW_GEMINI_MODEL,
|
||||||
|
'gemini-flash',
|
||||||
|
new TerminalQuotaError('pro quota', mockGoogleApiError),
|
||||||
|
);
|
||||||
|
});
|
||||||
|
|
||||||
|
// Simulate choosing "Switch to fallback"
|
||||||
|
await act(async () => {
|
||||||
|
result.current.handleOverageMenuChoice('use_fallback');
|
||||||
|
await promise!;
|
||||||
|
});
|
||||||
|
|
||||||
|
const intent = await promise!;
|
||||||
|
expect(intent).toBe('retry_always');
|
||||||
|
});
|
||||||
|
|
||||||
|
it('should show empty wallet menu if balance is 0', async () => {
|
||||||
|
vi.mocked(getG1CreditBalance).mockReturnValue(0);
|
||||||
|
vi.mocked(shouldAutoUseCredits).mockReturnValue(false);
|
||||||
|
vi.mocked(shouldShowOverageMenu).mockReturnValue(false);
|
||||||
|
vi.mocked(shouldShowEmptyWalletMenu).mockReturnValue(true);
|
||||||
|
|
||||||
|
const { result } = renderHook(() =>
|
||||||
|
useQuotaAndFallback({
|
||||||
|
config: mockConfig,
|
||||||
|
historyManager: mockHistoryManager,
|
||||||
|
userTier: UserTierId.STANDARD,
|
||||||
|
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||||
|
onShowAuthSelection: mockOnShowAuthSelection,
|
||||||
|
paidTier: { ...mockPaidTier, availableCredits: [] },
|
||||||
|
settings: mockSettings,
|
||||||
|
}),
|
||||||
|
);
|
||||||
|
|
||||||
|
const handler = setFallbackHandlerSpy.mock
|
||||||
|
.calls[0][0] as FallbackModelHandler;
|
||||||
|
|
||||||
|
let promise: Promise<FallbackIntent | null>;
|
||||||
|
act(() => {
|
||||||
|
promise = handler(
|
||||||
|
PREVIEW_GEMINI_MODEL,
|
||||||
|
'gemini-flash',
|
||||||
|
new TerminalQuotaError('pro quota', mockGoogleApiError),
|
||||||
|
);
|
||||||
|
});
|
||||||
|
|
||||||
|
expect(result.current.emptyWalletRequest).not.toBeNull();
|
||||||
|
expect(logBillingEvent).toHaveBeenCalled();
|
||||||
|
|
||||||
|
// Simulate choosing "Stop"
|
||||||
|
await act(async () => {
|
||||||
|
result.current.handleEmptyWalletChoice('stop');
|
||||||
|
await promise!;
|
||||||
|
});
|
||||||
|
|
||||||
|
const intent = await promise!;
|
||||||
|
expect(intent).toBe('stop');
|
||||||
|
});
|
||||||
|
|
||||||
|
it('should add info message to history when get_credits is selected', async () => {
|
||||||
|
vi.mocked(getG1CreditBalance).mockReturnValue(0);
|
||||||
|
vi.mocked(shouldAutoUseCredits).mockReturnValue(false);
|
||||||
|
vi.mocked(shouldShowOverageMenu).mockReturnValue(false);
|
||||||
|
vi.mocked(shouldShowEmptyWalletMenu).mockReturnValue(true);
|
||||||
|
|
||||||
|
const { result } = renderHook(() =>
|
||||||
|
useQuotaAndFallback({
|
||||||
|
config: mockConfig,
|
||||||
|
historyManager: mockHistoryManager,
|
||||||
|
userTier: UserTierId.STANDARD,
|
||||||
|
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||||
|
onShowAuthSelection: mockOnShowAuthSelection,
|
||||||
|
paidTier: { ...mockPaidTier, availableCredits: [] },
|
||||||
|
settings: mockSettings,
|
||||||
|
}),
|
||||||
|
);
|
||||||
|
|
||||||
|
const handler = setFallbackHandlerSpy.mock
|
||||||
|
.calls[0][0] as FallbackModelHandler;
|
||||||
|
|
||||||
|
let promise: Promise<FallbackIntent | null>;
|
||||||
|
act(() => {
|
||||||
|
promise = handler(
|
||||||
|
PREVIEW_GEMINI_MODEL,
|
||||||
|
'gemini-flash',
|
||||||
|
new TerminalQuotaError('pro quota', mockGoogleApiError),
|
||||||
|
);
|
||||||
|
});
|
||||||
|
|
||||||
|
expect(result.current.emptyWalletRequest).not.toBeNull();
|
||||||
|
|
||||||
|
// Simulate choosing "Get AI Credits"
|
||||||
|
await act(async () => {
|
||||||
|
result.current.handleEmptyWalletChoice('get_credits');
|
||||||
|
await promise!;
|
||||||
|
});
|
||||||
|
|
||||||
|
const intent = await promise!;
|
||||||
|
expect(intent).toBe('stop');
|
||||||
|
expect(mockHistoryManager.addItem).toHaveBeenCalledWith(
|
||||||
|
expect.objectContaining({
|
||||||
|
type: MessageType.INFO,
|
||||||
|
text: expect.stringContaining('few minutes'),
|
||||||
|
}),
|
||||||
|
expect.any(Number),
|
||||||
|
);
|
||||||
|
});
|
||||||
|
});
|
||||||
|
|
||||||
describe('handleProQuotaChoice', () => {
|
describe('handleProQuotaChoice', () => {
|
||||||
it('should do nothing if there is no pending pro quota request', () => {
|
it('should do nothing if there is no pending pro quota request', () => {
|
||||||
const { result } = renderHook(() =>
|
const { result } = renderHook(() =>
|
||||||
@@ -473,6 +754,8 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
|
|||||||
userTier: UserTierId.FREE,
|
userTier: UserTierId.FREE,
|
||||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||||
onShowAuthSelection: mockOnShowAuthSelection,
|
onShowAuthSelection: mockOnShowAuthSelection,
|
||||||
|
paidTier: null,
|
||||||
|
settings: mockSettings,
|
||||||
}),
|
}),
|
||||||
);
|
);
|
||||||
|
|
||||||
@@ -491,6 +774,8 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
|
|||||||
userTier: UserTierId.FREE,
|
userTier: UserTierId.FREE,
|
||||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||||
onShowAuthSelection: mockOnShowAuthSelection,
|
onShowAuthSelection: mockOnShowAuthSelection,
|
||||||
|
paidTier: null,
|
||||||
|
settings: mockSettings,
|
||||||
}),
|
}),
|
||||||
);
|
);
|
||||||
|
|
||||||
@@ -522,6 +807,8 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
|
|||||||
userTier: UserTierId.FREE,
|
userTier: UserTierId.FREE,
|
||||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||||
onShowAuthSelection: mockOnShowAuthSelection,
|
onShowAuthSelection: mockOnShowAuthSelection,
|
||||||
|
paidTier: null,
|
||||||
|
settings: mockSettings,
|
||||||
}),
|
}),
|
||||||
);
|
);
|
||||||
|
|
||||||
@@ -566,6 +853,8 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
|
|||||||
userTier: UserTierId.FREE,
|
userTier: UserTierId.FREE,
|
||||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||||
onShowAuthSelection: mockOnShowAuthSelection,
|
onShowAuthSelection: mockOnShowAuthSelection,
|
||||||
|
paidTier: null,
|
||||||
|
settings: mockSettings,
|
||||||
}),
|
}),
|
||||||
);
|
);
|
||||||
|
|
||||||
@@ -602,6 +891,8 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
|
|||||||
userTier: UserTierId.FREE,
|
userTier: UserTierId.FREE,
|
||||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||||
onShowAuthSelection: mockOnShowAuthSelection,
|
onShowAuthSelection: mockOnShowAuthSelection,
|
||||||
|
paidTier: null,
|
||||||
|
settings: mockSettings,
|
||||||
}),
|
}),
|
||||||
);
|
);
|
||||||
|
|
||||||
@@ -646,6 +937,8 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
|
|||||||
userTier: UserTierId.FREE,
|
userTier: UserTierId.FREE,
|
||||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||||
onShowAuthSelection: mockOnShowAuthSelection,
|
onShowAuthSelection: mockOnShowAuthSelection,
|
||||||
|
paidTier: null,
|
||||||
|
settings: mockSettings,
|
||||||
}),
|
}),
|
||||||
);
|
);
|
||||||
|
|
||||||
@@ -661,6 +954,8 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
|
|||||||
userTier: UserTierId.FREE,
|
userTier: UserTierId.FREE,
|
||||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||||
onShowAuthSelection: mockOnShowAuthSelection,
|
onShowAuthSelection: mockOnShowAuthSelection,
|
||||||
|
paidTier: null,
|
||||||
|
settings: mockSettings,
|
||||||
}),
|
}),
|
||||||
);
|
);
|
||||||
|
|
||||||
@@ -703,6 +998,8 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
|
|||||||
userTier: UserTierId.FREE,
|
userTier: UserTierId.FREE,
|
||||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||||
onShowAuthSelection: mockOnShowAuthSelection,
|
onShowAuthSelection: mockOnShowAuthSelection,
|
||||||
|
paidTier: null,
|
||||||
|
settings: mockSettings,
|
||||||
}),
|
}),
|
||||||
);
|
);
|
||||||
|
|
||||||
@@ -745,6 +1042,8 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
|
|||||||
userTier: UserTierId.FREE,
|
userTier: UserTierId.FREE,
|
||||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||||
onShowAuthSelection: mockOnShowAuthSelection,
|
onShowAuthSelection: mockOnShowAuthSelection,
|
||||||
|
paidTier: null,
|
||||||
|
settings: mockSettings,
|
||||||
}),
|
}),
|
||||||
);
|
);
|
||||||
|
|
||||||
@@ -775,6 +1074,8 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
|
|||||||
userTier: UserTierId.FREE,
|
userTier: UserTierId.FREE,
|
||||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||||
onShowAuthSelection: mockOnShowAuthSelection,
|
onShowAuthSelection: mockOnShowAuthSelection,
|
||||||
|
paidTier: null,
|
||||||
|
settings: mockSettings,
|
||||||
}),
|
}),
|
||||||
);
|
);
|
||||||
|
|
||||||
@@ -805,6 +1106,8 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
|
|||||||
userTier: UserTierId.FREE,
|
userTier: UserTierId.FREE,
|
||||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||||
onShowAuthSelection: mockOnShowAuthSelection,
|
onShowAuthSelection: mockOnShowAuthSelection,
|
||||||
|
paidTier: null,
|
||||||
|
settings: mockSettings,
|
||||||
}),
|
}),
|
||||||
);
|
);
|
||||||
|
|
||||||
|
|||||||
@@ -16,7 +16,9 @@ import {
|
|||||||
type UserTierId,
|
type UserTierId,
|
||||||
VALID_GEMINI_MODELS,
|
VALID_GEMINI_MODELS,
|
||||||
isProModel,
|
isProModel,
|
||||||
|
isOverageEligibleModel,
|
||||||
getDisplayString,
|
getDisplayString,
|
||||||
|
type GeminiUserTier,
|
||||||
} from '@google/gemini-cli-core';
|
} from '@google/gemini-cli-core';
|
||||||
import { useCallback, useEffect, useRef, useState } from 'react';
|
import { useCallback, useEffect, useRef, useState } from 'react';
|
||||||
import { type UseHistoryManagerReturn } from './useHistoryManager.js';
|
import { type UseHistoryManagerReturn } from './useHistoryManager.js';
|
||||||
@@ -24,12 +26,20 @@ import { MessageType } from '../types.js';
|
|||||||
import {
|
import {
|
||||||
type ProQuotaDialogRequest,
|
type ProQuotaDialogRequest,
|
||||||
type ValidationDialogRequest,
|
type ValidationDialogRequest,
|
||||||
|
type OverageMenuDialogRequest,
|
||||||
|
type OverageMenuIntent,
|
||||||
|
type EmptyWalletDialogRequest,
|
||||||
|
type EmptyWalletIntent,
|
||||||
} from '../contexts/UIStateContext.js';
|
} from '../contexts/UIStateContext.js';
|
||||||
|
import type { LoadedSettings } from '../../config/settings.js';
|
||||||
|
import { handleCreditsFlow } from './creditsFlowHandler.js';
|
||||||
|
|
||||||
interface UseQuotaAndFallbackArgs {
|
interface UseQuotaAndFallbackArgs {
|
||||||
config: Config;
|
config: Config;
|
||||||
historyManager: UseHistoryManagerReturn;
|
historyManager: UseHistoryManagerReturn;
|
||||||
userTier: UserTierId | undefined;
|
userTier: UserTierId | undefined;
|
||||||
|
paidTier: GeminiUserTier | null | undefined;
|
||||||
|
settings: LoadedSettings;
|
||||||
setModelSwitchedFromQuotaError: (value: boolean) => void;
|
setModelSwitchedFromQuotaError: (value: boolean) => void;
|
||||||
onShowAuthSelection: () => void;
|
onShowAuthSelection: () => void;
|
||||||
}
|
}
|
||||||
@@ -38,6 +48,8 @@ export function useQuotaAndFallback({
|
|||||||
config,
|
config,
|
||||||
historyManager,
|
historyManager,
|
||||||
userTier,
|
userTier,
|
||||||
|
paidTier,
|
||||||
|
settings,
|
||||||
setModelSwitchedFromQuotaError,
|
setModelSwitchedFromQuotaError,
|
||||||
onShowAuthSelection,
|
onShowAuthSelection,
|
||||||
}: UseQuotaAndFallbackArgs) {
|
}: UseQuotaAndFallbackArgs) {
|
||||||
@@ -45,9 +57,22 @@ export function useQuotaAndFallback({
|
|||||||
useState<ProQuotaDialogRequest | null>(null);
|
useState<ProQuotaDialogRequest | null>(null);
|
||||||
const [validationRequest, setValidationRequest] =
|
const [validationRequest, setValidationRequest] =
|
||||||
useState<ValidationDialogRequest | null>(null);
|
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 isDialogPending = useRef(false);
|
||||||
const isValidationPending = 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
|
// Set up Flash fallback handler
|
||||||
useEffect(() => {
|
useEffect(() => {
|
||||||
const fallbackHandler: FallbackModelHandler = async (
|
const fallbackHandler: FallbackModelHandler = async (
|
||||||
@@ -63,12 +88,52 @@ export function useQuotaAndFallback({
|
|||||||
const usageLimitReachedModel = isProModel(failedModel)
|
const usageLimitReachedModel = isProModel(failedModel)
|
||||||
? 'all Pro models'
|
? 'all Pro models'
|
||||||
: failedModel;
|
: failedModel;
|
||||||
|
|
||||||
if (error instanceof TerminalQuotaError) {
|
if (error instanceof TerminalQuotaError) {
|
||||||
isTerminalQuotaError = true;
|
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 = [
|
const messageLines = [
|
||||||
`Usage limit reached for ${usageLimitReachedModel}.`,
|
`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`,
|
`/stats model for usage details`,
|
||||||
`/model to switch models.`,
|
`/model to switch models.`,
|
||||||
contentGeneratorConfig?.authType === AuthType.LOGIN_WITH_GOOGLE
|
contentGeneratorConfig?.authType === AuthType.LOGIN_WITH_GOOGLE
|
||||||
@@ -126,7 +191,16 @@ export function useQuotaAndFallback({
|
|||||||
};
|
};
|
||||||
|
|
||||||
config.setFallbackModelHandler(fallbackHandler);
|
config.setFallbackModelHandler(fallbackHandler);
|
||||||
}, [config, historyManager, userTier, setModelSwitchedFromQuotaError]);
|
}, [
|
||||||
|
config,
|
||||||
|
historyManager,
|
||||||
|
userTier,
|
||||||
|
paidTier,
|
||||||
|
settings,
|
||||||
|
initialOverageStrategy,
|
||||||
|
setModelSwitchedFromQuotaError,
|
||||||
|
onShowAuthSelection,
|
||||||
|
]);
|
||||||
|
|
||||||
// Set up validation handler for 403 VALIDATION_REQUIRED errors
|
// Set up validation handler for 403 VALIDATION_REQUIRED errors
|
||||||
useEffect(() => {
|
useEffect(() => {
|
||||||
@@ -204,11 +278,38 @@ export function useQuotaAndFallback({
|
|||||||
[validationRequest, onShowAuthSelection],
|
[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 {
|
return {
|
||||||
proQuotaRequest,
|
proQuotaRequest,
|
||||||
handleProQuotaChoice,
|
handleProQuotaChoice,
|
||||||
validationRequest,
|
validationRequest,
|
||||||
handleValidationChoice,
|
handleValidationChoice,
|
||||||
|
// G1 AI Credits
|
||||||
|
overageMenuRequest,
|
||||||
|
handleOverageMenuChoice,
|
||||||
|
emptyWalletRequest,
|
||||||
|
handleEmptyWalletChoice,
|
||||||
};
|
};
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -221,5 +322,5 @@ function getResetTimeMessage(delayMs: number): string {
|
|||||||
timeZoneName: 'short',
|
timeZoneName: 'short',
|
||||||
});
|
});
|
||||||
|
|
||||||
return `Access resets at ${timeFormatter.format(resetDate)}.`;
|
return timeFormatter.format(resetDate);
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -201,6 +201,7 @@ export type HistoryItemStats = HistoryItemQuotaBase & {
|
|||||||
type: 'stats';
|
type: 'stats';
|
||||||
duration: string;
|
duration: string;
|
||||||
quotas?: RetrieveUserQuotaResponse;
|
quotas?: RetrieveUserQuotaResponse;
|
||||||
|
creditBalance?: number;
|
||||||
};
|
};
|
||||||
|
|
||||||
export type HistoryItemModelStats = HistoryItemQuotaBase & {
|
export type HistoryItemModelStats = HistoryItemQuotaBase & {
|
||||||
|
|||||||
@@ -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(true);
|
||||||
|
});
|
||||||
|
|
||||||
|
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);
|
||||||
|
});
|
||||||
|
});
|
||||||
|
});
|
||||||
@@ -0,0 +1,186 @@
|
|||||||
|
/**
|
||||||
|
* @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,
|
||||||
|
PREVIEW_GEMINI_3_1_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,
|
||||||
|
PREVIEW_GEMINI_3_1_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
|
||||||
|
);
|
||||||
|
}
|
||||||
@@ -0,0 +1,7 @@
|
|||||||
|
/**
|
||||||
|
* @license
|
||||||
|
* Copyright 2026 Google LLC
|
||||||
|
* SPDX-License-Identifier: Apache-2.0
|
||||||
|
*/
|
||||||
|
|
||||||
|
export * from './billing.js';
|
||||||
@@ -73,6 +73,8 @@ describe('codeAssist', () => {
|
|||||||
'session-123',
|
'session-123',
|
||||||
'free-tier',
|
'free-tier',
|
||||||
'free-tier-name',
|
'free-tier-name',
|
||||||
|
undefined,
|
||||||
|
mockConfig,
|
||||||
);
|
);
|
||||||
expect(generator).toBeInstanceOf(MockedCodeAssistServer);
|
expect(generator).toBeInstanceOf(MockedCodeAssistServer);
|
||||||
});
|
});
|
||||||
@@ -103,6 +105,8 @@ describe('codeAssist', () => {
|
|||||||
undefined, // No session ID
|
undefined, // No session ID
|
||||||
'free-tier',
|
'free-tier',
|
||||||
'free-tier-name',
|
'free-tier-name',
|
||||||
|
undefined,
|
||||||
|
mockConfig,
|
||||||
);
|
);
|
||||||
expect(generator).toBeInstanceOf(MockedCodeAssistServer);
|
expect(generator).toBeInstanceOf(MockedCodeAssistServer);
|
||||||
});
|
});
|
||||||
|
|||||||
@@ -36,6 +36,8 @@ export async function createCodeAssistContentGenerator(
|
|||||||
sessionId,
|
sessionId,
|
||||||
userData.userTier,
|
userData.userTier,
|
||||||
userData.userTierName,
|
userData.userTierName,
|
||||||
|
userData.paidTier,
|
||||||
|
config,
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -28,12 +28,14 @@ import type {
|
|||||||
} from '@google/genai';
|
} from '@google/genai';
|
||||||
import { GenerateContentResponse } from '@google/genai';
|
import { GenerateContentResponse } from '@google/genai';
|
||||||
import { debugLogger } from '../utils/debugLogger.js';
|
import { debugLogger } from '../utils/debugLogger.js';
|
||||||
|
import type { Credits } from './types.js';
|
||||||
|
|
||||||
export interface CAGenerateContentRequest {
|
export interface CAGenerateContentRequest {
|
||||||
model: string;
|
model: string;
|
||||||
project?: string;
|
project?: string;
|
||||||
user_prompt_id?: string;
|
user_prompt_id?: string;
|
||||||
request: VertexGenerateContentRequest;
|
request: VertexGenerateContentRequest;
|
||||||
|
enabled_credit_types?: string[];
|
||||||
}
|
}
|
||||||
|
|
||||||
interface VertexGenerateContentRequest {
|
interface VertexGenerateContentRequest {
|
||||||
@@ -75,6 +77,8 @@ interface VertexGenerationConfig {
|
|||||||
export interface CaGenerateContentResponse {
|
export interface CaGenerateContentResponse {
|
||||||
response?: VertexGenerateContentResponse;
|
response?: VertexGenerateContentResponse;
|
||||||
traceId?: string;
|
traceId?: string;
|
||||||
|
consumedCredits?: Credits[];
|
||||||
|
remainingCredits?: Credits[];
|
||||||
}
|
}
|
||||||
|
|
||||||
interface VertexGenerateContentResponse {
|
interface VertexGenerateContentResponse {
|
||||||
@@ -127,12 +131,14 @@ export function toGenerateContentRequest(
|
|||||||
userPromptId: string,
|
userPromptId: string,
|
||||||
project?: string,
|
project?: string,
|
||||||
sessionId?: string,
|
sessionId?: string,
|
||||||
|
enabledCreditTypes?: string[],
|
||||||
): CAGenerateContentRequest {
|
): CAGenerateContentRequest {
|
||||||
return {
|
return {
|
||||||
model: req.model,
|
model: req.model,
|
||||||
project,
|
project,
|
||||||
user_prompt_id: userPromptId,
|
user_prompt_id: userPromptId,
|
||||||
request: toVertexGenerateContentRequest(req, sessionId),
|
request: toVertexGenerateContentRequest(req, sessionId),
|
||||||
|
enabled_credit_types: enabledCreditTypes,
|
||||||
};
|
};
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -305,3 +311,16 @@ function toVertexGenerationConfig(
|
|||||||
thinkingConfig: config.thinkingConfig,
|
thinkingConfig: config.thinkingConfig,
|
||||||
};
|
};
|
||||||
}
|
}
|
||||||
|
|
||||||
|
export function fromGenerateContentResponseUsage(
|
||||||
|
metadata?: GenerateContentResponseUsageMetadata,
|
||||||
|
): GenerateContentResponseUsageMetadata | undefined {
|
||||||
|
if (!metadata) {
|
||||||
|
return undefined;
|
||||||
|
}
|
||||||
|
return {
|
||||||
|
promptTokenCount: metadata.promptTokenCount,
|
||||||
|
candidatesTokenCount: metadata.candidatesTokenCount,
|
||||||
|
totalTokenCount: metadata.totalTokenCount,
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|||||||
@@ -21,7 +21,10 @@ import type {
|
|||||||
ConversationInteraction,
|
ConversationInteraction,
|
||||||
StreamingLatency,
|
StreamingLatency,
|
||||||
RecordCodeAssistMetricsRequest,
|
RecordCodeAssistMetricsRequest,
|
||||||
|
GeminiUserTier,
|
||||||
|
Credits,
|
||||||
} from './types.js';
|
} from './types.js';
|
||||||
|
import { UserTierId } from './types.js';
|
||||||
import type {
|
import type {
|
||||||
ListExperimentsRequest,
|
ListExperimentsRequest,
|
||||||
ListExperimentsResponse,
|
ListExperimentsResponse,
|
||||||
@@ -37,7 +40,15 @@ import type {
|
|||||||
import * as readline from 'node:readline';
|
import * as readline from 'node:readline';
|
||||||
import { Readable } from 'node:stream';
|
import { Readable } from 'node:stream';
|
||||||
import type { ContentGenerator } from '../core/contentGenerator.js';
|
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 {
|
import type {
|
||||||
CaCountTokenResponse,
|
CaCountTokenResponse,
|
||||||
CaGenerateContentResponse,
|
CaGenerateContentResponse,
|
||||||
@@ -72,6 +83,8 @@ export class CodeAssistServer implements ContentGenerator {
|
|||||||
readonly sessionId?: string,
|
readonly sessionId?: string,
|
||||||
readonly userTier?: UserTierId,
|
readonly userTier?: UserTierId,
|
||||||
readonly userTierName?: string,
|
readonly userTierName?: string,
|
||||||
|
readonly paidTier?: GeminiUserTier,
|
||||||
|
readonly config?: Config,
|
||||||
) {}
|
) {}
|
||||||
|
|
||||||
async generateContentStream(
|
async generateContentStream(
|
||||||
@@ -80,6 +93,19 @@ export class CodeAssistServer implements ContentGenerator {
|
|||||||
// eslint-disable-next-line @typescript-eslint/no-unused-vars
|
// eslint-disable-next-line @typescript-eslint/no-unused-vars
|
||||||
role: LlmRole,
|
role: LlmRole,
|
||||||
): Promise<AsyncGenerator<GenerateContentResponse>> {
|
): 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 =
|
const responses =
|
||||||
await this.requestStreamingPost<CaGenerateContentResponse>(
|
await this.requestStreamingPost<CaGenerateContentResponse>(
|
||||||
'streamGenerateContent',
|
'streamGenerateContent',
|
||||||
@@ -88,6 +114,7 @@ export class CodeAssistServer implements ContentGenerator {
|
|||||||
userPromptId,
|
userPromptId,
|
||||||
this.projectId,
|
this.projectId,
|
||||||
this.sessionId,
|
this.sessionId,
|
||||||
|
enabledCreditTypes,
|
||||||
),
|
),
|
||||||
req.config?.abortSignal,
|
req.config?.abortSignal,
|
||||||
);
|
);
|
||||||
@@ -99,6 +126,9 @@ export class CodeAssistServer implements ContentGenerator {
|
|||||||
return (async function* (
|
return (async function* (
|
||||||
server: CodeAssistServer,
|
server: CodeAssistServer,
|
||||||
): AsyncGenerator<GenerateContentResponse> {
|
): AsyncGenerator<GenerateContentResponse> {
|
||||||
|
let totalConsumed = 0;
|
||||||
|
let lastRemaining = 0;
|
||||||
|
|
||||||
for await (const response of responses) {
|
for await (const response of responses) {
|
||||||
if (isFirst) {
|
if (isFirst) {
|
||||||
streamingLatency.firstMessageLatency = formatProtoJsonDuration(
|
streamingLatency.firstMessageLatency = formatProtoJsonDuration(
|
||||||
@@ -121,8 +151,38 @@ export class CodeAssistServer implements ContentGenerator {
|
|||||||
req.config?.abortSignal,
|
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;
|
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);
|
})(this);
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -140,6 +200,7 @@ export class CodeAssistServer implements ContentGenerator {
|
|||||||
userPromptId,
|
userPromptId,
|
||||||
this.projectId,
|
this.projectId,
|
||||||
this.sessionId,
|
this.sessionId,
|
||||||
|
undefined,
|
||||||
),
|
),
|
||||||
req.config?.abortSignal,
|
req.config?.abortSignal,
|
||||||
GENERATE_CONTENT_RETRY_DELAY_IN_MILLISECONDS,
|
GENERATE_CONTENT_RETRY_DELAY_IN_MILLISECONDS,
|
||||||
@@ -160,9 +221,29 @@ export class CodeAssistServer implements ContentGenerator {
|
|||||||
req.config?.abortSignal,
|
req.config?.abortSignal,
|
||||||
);
|
);
|
||||||
|
|
||||||
|
if (response.remainingCredits) {
|
||||||
|
this.updateCredits(response.remainingCredits);
|
||||||
|
}
|
||||||
|
|
||||||
return translatedResponse;
|
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(
|
async onboardUser(
|
||||||
req: OnboardUserRequest,
|
req: OnboardUserRequest,
|
||||||
): Promise<LongRunningOperationResponse> {
|
): Promise<LongRunningOperationResponse> {
|
||||||
@@ -192,6 +273,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(
|
async fetchAdminControls(
|
||||||
req: FetchAdminControlsRequest,
|
req: FetchAdminControlsRequest,
|
||||||
): Promise<FetchAdminControlsResponse> {
|
): Promise<FetchAdminControlsResponse> {
|
||||||
|
|||||||
@@ -52,6 +52,7 @@ export interface UserData {
|
|||||||
projectId: string;
|
projectId: string;
|
||||||
userTier: UserTierId;
|
userTier: UserTierId;
|
||||||
userTierName?: string;
|
userTierName?: string;
|
||||||
|
paidTier?: GeminiUserTier;
|
||||||
}
|
}
|
||||||
|
|
||||||
/**
|
/**
|
||||||
@@ -146,6 +147,7 @@ export async function setupUser(
|
|||||||
loadRes.currentTier.id ??
|
loadRes.currentTier.id ??
|
||||||
UserTierId.STANDARD,
|
UserTierId.STANDARD,
|
||||||
userTierName: loadRes.paidTier?.name ?? loadRes.currentTier.name,
|
userTierName: loadRes.paidTier?.name ?? loadRes.currentTier.name,
|
||||||
|
paidTier: loadRes.paidTier ?? undefined,
|
||||||
};
|
};
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -157,6 +159,7 @@ export async function setupUser(
|
|||||||
userTier:
|
userTier:
|
||||||
loadRes.paidTier?.id ?? loadRes.currentTier.id ?? UserTierId.STANDARD,
|
loadRes.paidTier?.id ?? loadRes.currentTier.id ?? UserTierId.STANDARD,
|
||||||
userTierName: loadRes.paidTier?.name ?? loadRes.currentTier.name,
|
userTierName: loadRes.paidTier?.name ?? loadRes.currentTier.name,
|
||||||
|
paidTier: loadRes.paidTier ?? undefined,
|
||||||
};
|
};
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -39,11 +39,41 @@ export type ClientMetadataPluginType =
|
|||||||
| 'AIPLUGIN_INTELLIJ'
|
| 'AIPLUGIN_INTELLIJ'
|
||||||
| 'AIPLUGIN_STUDIO';
|
| '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 {
|
export interface LoadCodeAssistRequest {
|
||||||
cloudaicompanionProject?: string;
|
cloudaicompanionProject?: string;
|
||||||
metadata: ClientMetadata;
|
metadata: ClientMetadata;
|
||||||
|
mode?: LoadCodeAssistMode;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
export type LoadCodeAssistMode =
|
||||||
|
| 'MODE_UNSPECIFIED'
|
||||||
|
| 'FULL_ELIGIBILITY_CHECK'
|
||||||
|
| 'HEALTH_CHECK';
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* Represents LoadCodeAssistResponse proto json field
|
* Represents LoadCodeAssistResponse proto json field
|
||||||
* http://google3/google/internal/cloud/code/v1internal/cloudcode.proto;l=224
|
* http://google3/google/internal/cloud/code/v1internal/cloudcode.proto;l=224
|
||||||
@@ -69,6 +99,8 @@ export interface GeminiUserTier {
|
|||||||
privacyNotice?: PrivacyNotice;
|
privacyNotice?: PrivacyNotice;
|
||||||
hasAcceptedTos?: boolean;
|
hasAcceptedTos?: boolean;
|
||||||
hasOnboardedPreviously?: boolean;
|
hasOnboardedPreviously?: boolean;
|
||||||
|
/** Available AI credits for this tier (e.g., Google One AI credits) */
|
||||||
|
availableCredits?: AvailableCredits[];
|
||||||
}
|
}
|
||||||
|
|
||||||
/**
|
/**
|
||||||
|
|||||||
@@ -13,6 +13,7 @@ import type {
|
|||||||
ContentGenerator,
|
ContentGenerator,
|
||||||
ContentGeneratorConfig,
|
ContentGeneratorConfig,
|
||||||
} from '../core/contentGenerator.js';
|
} from '../core/contentGenerator.js';
|
||||||
|
import type { OverageStrategy } from '../billing/billing.js';
|
||||||
import {
|
import {
|
||||||
AuthType,
|
AuthType,
|
||||||
createContentGenerator,
|
createContentGenerator,
|
||||||
@@ -116,6 +117,7 @@ import {
|
|||||||
import { HookSystem } from '../hooks/index.js';
|
import { HookSystem } from '../hooks/index.js';
|
||||||
import type {
|
import type {
|
||||||
UserTierId,
|
UserTierId,
|
||||||
|
GeminiUserTier,
|
||||||
RetrieveUserQuotaResponse,
|
RetrieveUserQuotaResponse,
|
||||||
AdminControlsSettings,
|
AdminControlsSettings,
|
||||||
} from '../code_assist/types.js';
|
} from '../code_assist/types.js';
|
||||||
@@ -573,6 +575,9 @@ export interface ConfigParameters {
|
|||||||
agents?: AgentSettings;
|
agents?: AgentSettings;
|
||||||
}>;
|
}>;
|
||||||
enableConseca?: boolean;
|
enableConseca?: boolean;
|
||||||
|
billing?: {
|
||||||
|
overageStrategy?: OverageStrategy;
|
||||||
|
};
|
||||||
}
|
}
|
||||||
|
|
||||||
export class Config {
|
export class Config {
|
||||||
@@ -754,6 +759,10 @@ export class Config {
|
|||||||
}>)
|
}>)
|
||||||
| undefined;
|
| undefined;
|
||||||
|
|
||||||
|
private readonly billing: {
|
||||||
|
overageStrategy: OverageStrategy;
|
||||||
|
};
|
||||||
|
|
||||||
private readonly enableAgents: boolean;
|
private readonly enableAgents: boolean;
|
||||||
private agents: AgentSettings;
|
private agents: AgentSettings;
|
||||||
private readonly enableEventDrivenScheduler: boolean;
|
private readonly enableEventDrivenScheduler: boolean;
|
||||||
@@ -1001,6 +1010,10 @@ export class Config {
|
|||||||
this.onModelChange = params.onModelChange;
|
this.onModelChange = params.onModelChange;
|
||||||
this.onReload = params.onReload;
|
this.onReload = params.onReload;
|
||||||
|
|
||||||
|
this.billing = {
|
||||||
|
overageStrategy: params.billing?.overageStrategy ?? 'ask',
|
||||||
|
};
|
||||||
|
|
||||||
if (params.contextFileName) {
|
if (params.contextFileName) {
|
||||||
setGeminiMdFilename(params.contextFileName);
|
setGeminiMdFilename(params.contextFileName);
|
||||||
}
|
}
|
||||||
@@ -1264,6 +1277,10 @@ export class Config {
|
|||||||
return this.contentGenerator?.userTierName;
|
return this.contentGenerator?.userTierName;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
getUserPaidTier(): GeminiUserTier | undefined {
|
||||||
|
return this.contentGenerator?.paidTier;
|
||||||
|
}
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* Provides access to the BaseLlmClient for stateless LLM operations.
|
* Provides access to the BaseLlmClient for stateless LLM operations.
|
||||||
*/
|
*/
|
||||||
@@ -1578,6 +1595,19 @@ export class Config {
|
|||||||
this.hasAccessToPreviewModel = hasAccess;
|
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> {
|
async refreshUserQuota(): Promise<RetrieveUserQuotaResponse | undefined> {
|
||||||
const codeAssistServer = getCodeAssistServer(this);
|
const codeAssistServer = getCodeAssistServer(this);
|
||||||
if (!codeAssistServer || !codeAssistServer.projectId) {
|
if (!codeAssistServer || !codeAssistServer.projectId) {
|
||||||
@@ -2052,6 +2082,19 @@ export class Config {
|
|||||||
return this.telemetrySettings.outfile;
|
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 {
|
getTelemetryUseCollector(): boolean {
|
||||||
return this.telemetrySettings.useCollector ?? false;
|
return this.telemetrySettings.useCollector ?? false;
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -17,7 +17,7 @@ import { createCodeAssistContentGenerator } from '../code_assist/codeAssist.js';
|
|||||||
import type { Config } from '../config/config.js';
|
import type { Config } from '../config/config.js';
|
||||||
import { loadApiKey } from './apiKeyCredentialStorage.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 { LoggingContentGenerator } from './loggingContentGenerator.js';
|
||||||
import { InstallationManager } from '../utils/installationManager.js';
|
import { InstallationManager } from '../utils/installationManager.js';
|
||||||
import { FakeContentGenerator } from './fakeContentGenerator.js';
|
import { FakeContentGenerator } from './fakeContentGenerator.js';
|
||||||
@@ -49,6 +49,8 @@ export interface ContentGenerator {
|
|||||||
userTier?: UserTierId;
|
userTier?: UserTierId;
|
||||||
|
|
||||||
userTierName?: string;
|
userTierName?: string;
|
||||||
|
|
||||||
|
paidTier?: GeminiUserTier;
|
||||||
}
|
}
|
||||||
|
|
||||||
export enum AuthType {
|
export enum AuthType {
|
||||||
|
|||||||
@@ -14,7 +14,7 @@ import {
|
|||||||
} from '@google/genai';
|
} from '@google/genai';
|
||||||
import { promises } from 'node:fs';
|
import { promises } from 'node:fs';
|
||||||
import type { ContentGenerator } from './contentGenerator.js';
|
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 { safeJsonStringify } from '../utils/safeJsonStringify.js';
|
||||||
import type { LlmRole } from '../telemetry/types.js';
|
import type { LlmRole } from '../telemetry/types.js';
|
||||||
|
|
||||||
@@ -44,6 +44,7 @@ export class FakeContentGenerator implements ContentGenerator {
|
|||||||
private callCounter = 0;
|
private callCounter = 0;
|
||||||
userTier?: UserTierId;
|
userTier?: UserTierId;
|
||||||
userTierName?: string;
|
userTierName?: string;
|
||||||
|
paidTier?: GeminiUserTier;
|
||||||
|
|
||||||
constructor(private readonly responses: FakeResponse[]) {}
|
constructor(private readonly responses: FakeResponse[]) {}
|
||||||
|
|
||||||
|
|||||||
@@ -24,7 +24,7 @@ import {
|
|||||||
} from '../telemetry/types.js';
|
} from '../telemetry/types.js';
|
||||||
import type { LlmRole } from '../telemetry/llmRole.js';
|
import type { LlmRole } from '../telemetry/llmRole.js';
|
||||||
import type { Config } from '../config/config.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 {
|
import {
|
||||||
logApiError,
|
logApiError,
|
||||||
logApiRequest,
|
logApiRequest,
|
||||||
@@ -163,6 +163,10 @@ export class LoggingContentGenerator implements ContentGenerator {
|
|||||||
return this.wrapped.userTierName;
|
return this.wrapped.userTierName;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
get paidTier(): GeminiUserTier | undefined {
|
||||||
|
return this.wrapped.paidTier;
|
||||||
|
}
|
||||||
|
|
||||||
private logApiRequest(
|
private logApiRequest(
|
||||||
contents: Content[],
|
contents: Content[],
|
||||||
model: string,
|
model: string,
|
||||||
|
|||||||
@@ -140,6 +140,9 @@ async function processIntent(
|
|||||||
// based on the availability service state (which is updated before this).
|
// based on the availability service state (which is updated before this).
|
||||||
return true;
|
return true;
|
||||||
|
|
||||||
|
case 'retry_with_credits':
|
||||||
|
return true;
|
||||||
|
|
||||||
case 'stop':
|
case 'stop':
|
||||||
// Do not switch model on stop. User wants to stay on current model (and stop).
|
// Do not switch model on stop. User wants to stay on current model (and stop).
|
||||||
return false;
|
return false;
|
||||||
|
|||||||
@@ -17,6 +17,7 @@ import type {
|
|||||||
export type FallbackIntent =
|
export type FallbackIntent =
|
||||||
| 'retry_always' // Retry with fallback model and stick to it for future requests.
|
| 'retry_always' // Retry with fallback model and stick to it for future requests.
|
||||||
| 'retry_once' // Retry with fallback model for this request only.
|
| '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.
|
| '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.
|
| '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.
|
| 'upgrade'; // Give user an option to upgrade the tier.
|
||||||
|
|||||||
@@ -18,6 +18,7 @@ export * from './policy/policy-engine.js';
|
|||||||
export * from './policy/toml-loader.js';
|
export * from './policy/toml-loader.js';
|
||||||
export * from './policy/config.js';
|
export * from './policy/config.js';
|
||||||
export * from './policy/integrity.js';
|
export * from './policy/integrity.js';
|
||||||
|
export * from './billing/index.js';
|
||||||
export * from './confirmation-bus/types.js';
|
export * from './confirmation-bus/types.js';
|
||||||
export * from './confirmation-bus/message-bus.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/userAccountManager.js';
|
||||||
export * from './utils/authConsent.js';
|
export * from './utils/authConsent.js';
|
||||||
export * from './utils/googleQuotaErrors.js';
|
export * from './utils/googleQuotaErrors.js';
|
||||||
|
export * from './utils/googleErrors.js';
|
||||||
export * from './utils/fileUtils.js';
|
export * from './utils/fileUtils.js';
|
||||||
export * from './utils/planUtils.js';
|
export * from './utils/planUtils.js';
|
||||||
export * from './utils/approvalModeUtils.js';
|
export * from './utils/approvalModeUtils.js';
|
||||||
@@ -98,6 +100,7 @@ export * from './utils/ignorePatterns.js';
|
|||||||
export * from './utils/partUtils.js';
|
export * from './utils/partUtils.js';
|
||||||
export * from './utils/promptIdContext.js';
|
export * from './utils/promptIdContext.js';
|
||||||
export * from './utils/thoughtUtils.js';
|
export * from './utils/thoughtUtils.js';
|
||||||
|
export * from './utils/secure-browser-launcher.js';
|
||||||
export * from './utils/debugLogger.js';
|
export * from './utils/debugLogger.js';
|
||||||
export * from './utils/events.js';
|
export * from './utils/events.js';
|
||||||
export * from './utils/extensionLoader.js';
|
export * from './utils/extensionLoader.js';
|
||||||
@@ -183,6 +186,8 @@ export { OAuthUtils } from './mcp/oauth-utils.js';
|
|||||||
|
|
||||||
// Export telemetry functions
|
// Export telemetry functions
|
||||||
export * from './telemetry/index.js';
|
export * from './telemetry/index.js';
|
||||||
|
export * from './telemetry/billingEvents.js';
|
||||||
|
export { logBillingEvent } from './telemetry/loggers.js';
|
||||||
export { sessionId, createSessionId } from './utils/session.js';
|
export { sessionId, createSessionId } from './utils/session.js';
|
||||||
export * from './utils/compatibility.js';
|
export * from './utils/compatibility.js';
|
||||||
export * from './utils/browser.js';
|
export * from './utils/browser.js';
|
||||||
|
|||||||
@@ -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');
|
||||||
|
});
|
||||||
|
});
|
||||||
|
});
|
||||||
@@ -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;
|
||||||
@@ -39,6 +39,7 @@ describe('conseca-logger', () => {
|
|||||||
getTelemetryLogPromptsEnabled: vi.fn().mockReturnValue(true),
|
getTelemetryLogPromptsEnabled: vi.fn().mockReturnValue(true),
|
||||||
isInteractive: vi.fn().mockReturnValue(true),
|
isInteractive: vi.fn().mockReturnValue(true),
|
||||||
getExperiments: vi.fn().mockReturnValue({ experimentIds: [] }),
|
getExperiments: vi.fn().mockReturnValue({ experimentIds: [] }),
|
||||||
|
getContentGeneratorConfig: vi.fn().mockReturnValue({ authType: 'oauth' }),
|
||||||
} as unknown as Config;
|
} as unknown as Config;
|
||||||
|
|
||||||
mockLogger = {
|
mockLogger = {
|
||||||
|
|||||||
@@ -77,6 +77,7 @@ export type { TelemetryEvent } from './types.js';
|
|||||||
export { SpanStatusCode, ValueType } from '@opentelemetry/api';
|
export { SpanStatusCode, ValueType } from '@opentelemetry/api';
|
||||||
export { SemanticAttributes } from '@opentelemetry/semantic-conventions';
|
export { SemanticAttributes } from '@opentelemetry/semantic-conventions';
|
||||||
export * from './uiTelemetry.js';
|
export * from './uiTelemetry.js';
|
||||||
|
export * from './billingEvents.js';
|
||||||
export {
|
export {
|
||||||
MemoryMonitor,
|
MemoryMonitor,
|
||||||
initializeMemoryMonitor,
|
initializeMemoryMonitor,
|
||||||
@@ -145,6 +146,9 @@ export {
|
|||||||
GenAiOperationName,
|
GenAiOperationName,
|
||||||
GenAiProviderName,
|
GenAiProviderName,
|
||||||
GenAiTokenType,
|
GenAiTokenType,
|
||||||
|
// Billing metrics functions
|
||||||
|
recordOverageOptionSelected,
|
||||||
|
recordCreditPurchaseClick,
|
||||||
} from './metrics.js';
|
} from './metrics.js';
|
||||||
export { runInDevTraceSpan, type SpanMetadata } from './trace.js';
|
export { runInDevTraceSpan, type SpanMetadata } from './trace.js';
|
||||||
export { startupProfiler, StartupProfiler } from './startupProfiler.js';
|
export { startupProfiler, StartupProfiler } from './startupProfiler.js';
|
||||||
|
|||||||
@@ -280,6 +280,7 @@ describe('loggers', () => {
|
|||||||
isInteractive: () => false,
|
isInteractive: () => false,
|
||||||
getExperiments: () => undefined,
|
getExperiments: () => undefined,
|
||||||
getExperimentsAsync: async () => undefined,
|
getExperimentsAsync: async () => undefined,
|
||||||
|
getContentGeneratorConfig: () => undefined,
|
||||||
} as unknown as Config;
|
} as unknown as Config;
|
||||||
|
|
||||||
it('should log a user prompt', () => {
|
it('should log a user prompt', () => {
|
||||||
@@ -319,6 +320,7 @@ describe('loggers', () => {
|
|||||||
isInteractive: () => false,
|
isInteractive: () => false,
|
||||||
getExperiments: () => undefined,
|
getExperiments: () => undefined,
|
||||||
getExperimentsAsync: async () => undefined,
|
getExperimentsAsync: async () => undefined,
|
||||||
|
getContentGeneratorConfig: () => undefined,
|
||||||
} as unknown as Config;
|
} as unknown as Config;
|
||||||
const event = new UserPromptEvent(
|
const event = new UserPromptEvent(
|
||||||
11,
|
11,
|
||||||
@@ -356,7 +358,8 @@ describe('loggers', () => {
|
|||||||
isInteractive: () => false,
|
isInteractive: () => false,
|
||||||
getExperiments: () => undefined,
|
getExperiments: () => undefined,
|
||||||
getExperimentsAsync: async () => undefined,
|
getExperimentsAsync: async () => undefined,
|
||||||
} as Config;
|
getContentGeneratorConfig: () => undefined,
|
||||||
|
} as unknown as Config;
|
||||||
|
|
||||||
const mockMetrics = {
|
const mockMetrics = {
|
||||||
recordApiResponseMetrics: vi.fn(),
|
recordApiResponseMetrics: vi.fn(),
|
||||||
@@ -558,7 +561,8 @@ describe('loggers', () => {
|
|||||||
isInteractive: () => false,
|
isInteractive: () => false,
|
||||||
getExperiments: () => undefined,
|
getExperiments: () => undefined,
|
||||||
getExperimentsAsync: async () => undefined,
|
getExperimentsAsync: async () => undefined,
|
||||||
} as Config;
|
getContentGeneratorConfig: () => undefined,
|
||||||
|
} as unknown as Config;
|
||||||
|
|
||||||
const mockMetrics = {
|
const mockMetrics = {
|
||||||
recordApiResponseMetrics: vi.fn(),
|
recordApiResponseMetrics: vi.fn(),
|
||||||
@@ -996,6 +1000,7 @@ describe('loggers', () => {
|
|||||||
isInteractive: () => false,
|
isInteractive: () => false,
|
||||||
getExperiments: () => undefined,
|
getExperiments: () => undefined,
|
||||||
getExperimentsAsync: async () => undefined,
|
getExperimentsAsync: async () => undefined,
|
||||||
|
getContentGeneratorConfig: () => undefined,
|
||||||
} as unknown as Config;
|
} as unknown as Config;
|
||||||
|
|
||||||
it('should log flash fallback event', () => {
|
it('should log flash fallback event', () => {
|
||||||
@@ -1025,6 +1030,7 @@ describe('loggers', () => {
|
|||||||
isInteractive: () => false,
|
isInteractive: () => false,
|
||||||
getExperiments: () => undefined,
|
getExperiments: () => undefined,
|
||||||
getExperimentsAsync: async () => undefined,
|
getExperimentsAsync: async () => undefined,
|
||||||
|
getContentGeneratorConfig: () => undefined,
|
||||||
} as unknown as Config;
|
} as unknown as Config;
|
||||||
|
|
||||||
beforeEach(() => {
|
beforeEach(() => {
|
||||||
@@ -1121,7 +1127,8 @@ describe('loggers', () => {
|
|||||||
isInteractive: () => false,
|
isInteractive: () => false,
|
||||||
getExperiments: () => undefined,
|
getExperiments: () => undefined,
|
||||||
getExperimentsAsync: async () => undefined,
|
getExperimentsAsync: async () => undefined,
|
||||||
} as Config;
|
getContentGeneratorConfig: () => undefined,
|
||||||
|
} as unknown as Config;
|
||||||
|
|
||||||
const mockMetrics = {
|
const mockMetrics = {
|
||||||
recordToolCallMetrics: vi.fn(),
|
recordToolCallMetrics: vi.fn(),
|
||||||
@@ -1741,7 +1748,8 @@ describe('loggers', () => {
|
|||||||
isInteractive: () => false,
|
isInteractive: () => false,
|
||||||
getExperiments: () => undefined,
|
getExperiments: () => undefined,
|
||||||
getExperimentsAsync: async () => undefined,
|
getExperimentsAsync: async () => undefined,
|
||||||
} as Config;
|
getContentGeneratorConfig: () => undefined,
|
||||||
|
} as unknown as Config;
|
||||||
|
|
||||||
const mockMetrics = {
|
const mockMetrics = {
|
||||||
recordFileOperationMetric: vi.fn(),
|
recordFileOperationMetric: vi.fn(),
|
||||||
@@ -1803,6 +1811,7 @@ describe('loggers', () => {
|
|||||||
isInteractive: () => false,
|
isInteractive: () => false,
|
||||||
getExperiments: () => undefined,
|
getExperiments: () => undefined,
|
||||||
getExperimentsAsync: async () => undefined,
|
getExperimentsAsync: async () => undefined,
|
||||||
|
getContentGeneratorConfig: () => undefined,
|
||||||
} as unknown as Config;
|
} as unknown as Config;
|
||||||
|
|
||||||
it('should log a tool output truncated event', () => {
|
it('should log a tool output truncated event', () => {
|
||||||
@@ -1842,6 +1851,7 @@ describe('loggers', () => {
|
|||||||
isInteractive: () => false,
|
isInteractive: () => false,
|
||||||
getExperiments: () => undefined,
|
getExperiments: () => undefined,
|
||||||
getExperimentsAsync: async () => undefined,
|
getExperimentsAsync: async () => undefined,
|
||||||
|
getContentGeneratorConfig: () => undefined,
|
||||||
} as unknown as Config;
|
} as unknown as Config;
|
||||||
|
|
||||||
beforeEach(() => {
|
beforeEach(() => {
|
||||||
@@ -2099,6 +2109,7 @@ describe('loggers', () => {
|
|||||||
isInteractive: () => false,
|
isInteractive: () => false,
|
||||||
getExperiments: () => undefined,
|
getExperiments: () => undefined,
|
||||||
getExperimentsAsync: async () => undefined,
|
getExperimentsAsync: async () => undefined,
|
||||||
|
getContentGeneratorConfig: () => undefined,
|
||||||
} as unknown as Config;
|
} as unknown as Config;
|
||||||
|
|
||||||
beforeEach(() => {
|
beforeEach(() => {
|
||||||
@@ -2146,6 +2157,7 @@ describe('loggers', () => {
|
|||||||
isInteractive: () => false,
|
isInteractive: () => false,
|
||||||
getExperiments: () => undefined,
|
getExperiments: () => undefined,
|
||||||
getExperimentsAsync: async () => undefined,
|
getExperimentsAsync: async () => undefined,
|
||||||
|
getContentGeneratorConfig: () => undefined,
|
||||||
} as unknown as Config;
|
} as unknown as Config;
|
||||||
|
|
||||||
beforeEach(() => {
|
beforeEach(() => {
|
||||||
@@ -2193,6 +2205,7 @@ describe('loggers', () => {
|
|||||||
isInteractive: () => false,
|
isInteractive: () => false,
|
||||||
getExperiments: () => undefined,
|
getExperiments: () => undefined,
|
||||||
getExperimentsAsync: async () => undefined,
|
getExperimentsAsync: async () => undefined,
|
||||||
|
getContentGeneratorConfig: () => undefined,
|
||||||
} as unknown as Config;
|
} as unknown as Config;
|
||||||
|
|
||||||
beforeEach(() => {
|
beforeEach(() => {
|
||||||
@@ -2231,6 +2244,7 @@ describe('loggers', () => {
|
|||||||
isInteractive: () => false,
|
isInteractive: () => false,
|
||||||
getExperiments: () => undefined,
|
getExperiments: () => undefined,
|
||||||
getExperimentsAsync: async () => undefined,
|
getExperimentsAsync: async () => undefined,
|
||||||
|
getContentGeneratorConfig: () => undefined,
|
||||||
} as unknown as Config;
|
} as unknown as Config;
|
||||||
|
|
||||||
beforeEach(() => {
|
beforeEach(() => {
|
||||||
@@ -2284,6 +2298,7 @@ describe('loggers', () => {
|
|||||||
isInteractive: () => false,
|
isInteractive: () => false,
|
||||||
getExperiments: () => undefined,
|
getExperiments: () => undefined,
|
||||||
getExperimentsAsync: async () => undefined,
|
getExperimentsAsync: async () => undefined,
|
||||||
|
getContentGeneratorConfig: () => undefined,
|
||||||
} as unknown as Config;
|
} as unknown as Config;
|
||||||
|
|
||||||
beforeEach(() => {
|
beforeEach(() => {
|
||||||
@@ -2322,6 +2337,7 @@ describe('loggers', () => {
|
|||||||
getExperiments: () => undefined,
|
getExperiments: () => undefined,
|
||||||
getExperimentsAsync: async () => undefined,
|
getExperimentsAsync: async () => undefined,
|
||||||
getTelemetryLogPromptsEnabled: () => false,
|
getTelemetryLogPromptsEnabled: () => false,
|
||||||
|
getContentGeneratorConfig: () => undefined,
|
||||||
} as unknown as Config;
|
} as unknown as Config;
|
||||||
|
|
||||||
beforeEach(() => {
|
beforeEach(() => {
|
||||||
|
|||||||
@@ -84,6 +84,7 @@ import type { UiEvent } from './uiTelemetry.js';
|
|||||||
import { uiTelemetryService } from './uiTelemetry.js';
|
import { uiTelemetryService } from './uiTelemetry.js';
|
||||||
import { ClearcutLogger } from './clearcut-logger/clearcut-logger.js';
|
import { ClearcutLogger } from './clearcut-logger/clearcut-logger.js';
|
||||||
import { debugLogger } from '../utils/debugLogger.js';
|
import { debugLogger } from '../utils/debugLogger.js';
|
||||||
|
import type { BillingTelemetryEvent } from './billingEvents.js';
|
||||||
|
|
||||||
export function logCliConfiguration(
|
export function logCliConfiguration(
|
||||||
config: Config,
|
config: Config,
|
||||||
@@ -827,3 +828,17 @@ export function logTokenStorageInitialization(
|
|||||||
recordTokenStorageInitialization(config, event);
|
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);
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|||||||
@@ -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 EVENT_HOOK_CALL_LATENCY = 'gemini_cli.hook_call.latency';
|
||||||
const KEYCHAIN_AVAILABILITY_COUNT = 'gemini_cli.keychain.availability.count';
|
const KEYCHAIN_AVAILABILITY_COUNT = 'gemini_cli.keychain.availability.count';
|
||||||
const TOKEN_STORAGE_TYPE_COUNT = 'gemini_cli.token_storage.type.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
|
// Agent Metrics
|
||||||
const AGENT_RUN_COUNT = 'gemini_cli.agent.run.count';
|
const AGENT_RUN_COUNT = 'gemini_cli.agent.run.count';
|
||||||
@@ -259,6 +261,26 @@ const COUNTER_DEFINITIONS = {
|
|||||||
forced: boolean;
|
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;
|
} as const;
|
||||||
|
|
||||||
const HISTOGRAM_DEFINITIONS = {
|
const HISTOGRAM_DEFINITIONS = {
|
||||||
@@ -597,6 +619,8 @@ let hookCallCounter: Counter | undefined;
|
|||||||
let hookCallLatencyHistogram: Histogram | undefined;
|
let hookCallLatencyHistogram: Histogram | undefined;
|
||||||
let keychainAvailabilityCounter: Counter | undefined;
|
let keychainAvailabilityCounter: Counter | undefined;
|
||||||
let tokenStorageTypeCounter: Counter | undefined;
|
let tokenStorageTypeCounter: Counter | undefined;
|
||||||
|
let overageOptionCounter: Counter | undefined;
|
||||||
|
let creditPurchaseCounter: Counter | undefined;
|
||||||
|
|
||||||
// OpenTelemetry GenAI Semantic Convention Metrics
|
// OpenTelemetry GenAI Semantic Convention Metrics
|
||||||
let genAiClientTokenUsageHistogram: Histogram | undefined;
|
let genAiClientTokenUsageHistogram: Histogram | undefined;
|
||||||
@@ -1334,3 +1358,31 @@ export function recordTokenStorageInitialization(
|
|||||||
forced: event.forced,
|
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,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|||||||
@@ -32,6 +32,7 @@ function createMockConfig(logPromptsEnabled: boolean): Config {
|
|||||||
getModel: () => 'gemini-1.5-flash',
|
getModel: () => 'gemini-1.5-flash',
|
||||||
isInteractive: () => true,
|
isInteractive: () => true,
|
||||||
getUserEmail: () => undefined,
|
getUserEmail: () => undefined,
|
||||||
|
getContentGeneratorConfig: () => undefined,
|
||||||
} as unknown as Config;
|
} as unknown as Config;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -77,6 +77,7 @@ describe('Telemetry SDK', () => {
|
|||||||
isInteractive: () => false,
|
isInteractive: () => false,
|
||||||
getExperiments: () => undefined,
|
getExperiments: () => undefined,
|
||||||
getExperimentsAsync: async () => undefined,
|
getExperimentsAsync: async () => undefined,
|
||||||
|
getContentGeneratorConfig: () => undefined,
|
||||||
} as unknown as Config;
|
} as unknown as Config;
|
||||||
});
|
});
|
||||||
|
|
||||||
|
|||||||
@@ -15,11 +15,13 @@ const installationManager = new InstallationManager();
|
|||||||
export function getCommonAttributes(config: Config): Attributes {
|
export function getCommonAttributes(config: Config): Attributes {
|
||||||
const email = userAccountManager.getCachedGoogleAccount();
|
const email = userAccountManager.getCachedGoogleAccount();
|
||||||
const experiments = config.getExperiments();
|
const experiments = config.getExperiments();
|
||||||
|
const authType = config.getContentGeneratorConfig()?.authType;
|
||||||
return {
|
return {
|
||||||
'session.id': config.getSessionId(),
|
'session.id': config.getSessionId(),
|
||||||
'installation.id': installationManager.getInstallationId(),
|
'installation.id': installationManager.getInstallationId(),
|
||||||
interactive: config.isInteractive(),
|
interactive: config.isInteractive(),
|
||||||
...(email && { 'user.email': email }),
|
...(email && { 'user.email': email }),
|
||||||
|
...(authType && { auth_type: authType }),
|
||||||
...(experiments &&
|
...(experiments &&
|
||||||
experiments.experimentIds.length > 0 && {
|
experiments.experimentIds.length > 0 && {
|
||||||
'experiments.ids': experiments.experimentIds,
|
'experiments.ids': experiments.experimentIds,
|
||||||
|
|||||||
@@ -16,8 +16,8 @@
|
|||||||
export interface ErrorInfo {
|
export interface ErrorInfo {
|
||||||
'@type': 'type.googleapis.com/google.rpc.ErrorInfo';
|
'@type': 'type.googleapis.com/google.rpc.ErrorInfo';
|
||||||
reason: string;
|
reason: string;
|
||||||
domain: string;
|
domain?: string;
|
||||||
metadata: { [key: string]: string };
|
metadata?: { [key: string]: string };
|
||||||
}
|
}
|
||||||
|
|
||||||
export interface RetryInfo {
|
export interface RetryInfo {
|
||||||
|
|||||||
@@ -313,6 +313,26 @@ describe('classifyGoogleError', () => {
|
|||||||
expect(result).toBeInstanceOf(TerminalQuotaError);
|
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', () => {
|
it('should prioritize daily limit over retry info', () => {
|
||||||
const apiError: GoogleApiError = {
|
const apiError: GoogleApiError = {
|
||||||
code: 429,
|
code: 429,
|
||||||
|
|||||||
@@ -19,17 +19,24 @@ import { getErrorStatus, ModelNotFoundError } from './httpErrors.js';
|
|||||||
*/
|
*/
|
||||||
export class TerminalQuotaError extends Error {
|
export class TerminalQuotaError extends Error {
|
||||||
retryDelayMs?: number;
|
retryDelayMs?: number;
|
||||||
|
reason?: string;
|
||||||
|
|
||||||
constructor(
|
constructor(
|
||||||
message: string,
|
message: string,
|
||||||
override readonly cause: GoogleApiError,
|
override readonly cause: GoogleApiError,
|
||||||
retryDelaySeconds?: number,
|
retryDelaySeconds?: number,
|
||||||
|
reason?: string,
|
||||||
) {
|
) {
|
||||||
super(message);
|
super(message);
|
||||||
this.name = 'TerminalQuotaError';
|
this.name = 'TerminalQuotaError';
|
||||||
this.retryDelayMs = retryDelaySeconds
|
this.retryDelayMs = retryDelaySeconds
|
||||||
? retryDelaySeconds * 1000
|
? retryDelaySeconds * 1000
|
||||||
: undefined;
|
: undefined;
|
||||||
|
this.reason = reason;
|
||||||
|
}
|
||||||
|
|
||||||
|
get isInsufficientCredits(): boolean {
|
||||||
|
return this.reason === 'INSUFFICIENT_G1_CREDITS_BALANCE';
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -121,6 +128,7 @@ function classifyValidationRequiredError(
|
|||||||
}
|
}
|
||||||
|
|
||||||
if (
|
if (
|
||||||
|
!errorInfo.domain ||
|
||||||
!CLOUDCODE_DOMAINS.includes(errorInfo.domain) ||
|
!CLOUDCODE_DOMAINS.includes(errorInfo.domain) ||
|
||||||
errorInfo.reason !== 'VALIDATION_REQUIRED'
|
errorInfo.reason !== 'VALIDATION_REQUIRED'
|
||||||
) {
|
) {
|
||||||
@@ -293,6 +301,16 @@ export function classifyGoogleError(error: unknown): unknown {
|
|||||||
}
|
}
|
||||||
|
|
||||||
if (errorInfo) {
|
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
|
// New Cloud Code API quota handling
|
||||||
if (errorInfo.domain) {
|
if (errorInfo.domain) {
|
||||||
const validDomains = [
|
const validDomains = [
|
||||||
@@ -313,6 +331,7 @@ export function classifyGoogleError(error: unknown): unknown {
|
|||||||
`${googleApiError.message}`,
|
`${googleApiError.message}`,
|
||||||
googleApiError,
|
googleApiError,
|
||||||
delaySeconds,
|
delaySeconds,
|
||||||
|
errorInfo.reason,
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -520,6 +520,24 @@
|
|||||||
"markdownDescription": "Telemetry configuration.\n\n- Category: `Advanced`\n- Requires restart: `yes`",
|
"markdownDescription": "Telemetry configuration.\n\n- Category: `Advanced`\n- Requires restart: `yes`",
|
||||||
"$ref": "#/$defs/TelemetrySettings"
|
"$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": {
|
"model": {
|
||||||
"title": "Model",
|
"title": "Model",
|
||||||
"description": "Settings related to the generative model.",
|
"description": "Settings related to the generative model.",
|
||||||
|
|||||||
Reference in New Issue
Block a user