mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-04-22 19:14:33 -07:00
feat: implement G1 AI credits overage flow with billing telemetry
Adds end-to-end support for Google One AI credits in quota exhaustion flows: - New billing module (packages/core/src/billing/) with credit balance checking, overage strategy management, and G1 URL construction - OverageMenuDialog and EmptyWalletDialog UI components for quota exhaustion with credit purchase options - Credits flow handler extracted to creditsFlowHandler.ts with overage menu, empty wallet, and auto-use-credits logic - Server-side credit tracking: enabledCreditTypes on requests, consumed/remaining credits from streaming responses - Billing telemetry events (overage menu shown, option selected, credits used, credit purchase click, API key updated) - OpenTelemetry metrics for overage option and credit purchase counters - Credit balance display in /stats command with refresh support - Settings: general.overageStrategy (ask/always/never) for credit usage - Error handling: INSUFFICIENT_G1_CREDITS_BALANCE as terminal error regardless of domain field presence - Persistent info message after
This commit is contained in:
Gaurav Ghosh
@@ -14,7 +14,8 @@ import {
|
||||
type Mock,
|
||||
} from 'vitest';
|
||||
import { act } from 'react';
|
||||
import { renderHook } from '../../test-utils/render.js';
|
||||
import { renderHook, mockSettings } from '../../test-utils/render.js';
|
||||
import { waitFor } from '../../test-utils/async.js';
|
||||
import {
|
||||
type Config,
|
||||
type FallbackModelHandler,
|
||||
@@ -29,6 +30,12 @@ import {
|
||||
ModelNotFoundError,
|
||||
DEFAULT_GEMINI_MODEL,
|
||||
DEFAULT_GEMINI_FLASH_MODEL,
|
||||
getG1CreditBalance,
|
||||
shouldAutoUseCredits,
|
||||
shouldShowOverageMenu,
|
||||
shouldShowEmptyWalletMenu,
|
||||
logBillingEvent,
|
||||
G1_CREDIT_TYPE,
|
||||
} from '@google/gemini-cli-core';
|
||||
import { useQuotaAndFallback } from './useQuotaAndFallback.js';
|
||||
import type { UseHistoryManagerReturn } from './useHistoryManager.js';
|
||||
@@ -37,6 +44,19 @@ import { MessageType } from '../types.js';
|
||||
// Use a type alias for SpyInstance as it's not directly exported
|
||||
type SpyInstance = ReturnType<typeof vi.spyOn>;
|
||||
|
||||
vi.mock('@google/gemini-cli-core', async (importOriginal) => {
|
||||
const actual =
|
||||
await importOriginal<typeof import('@google/gemini-cli-core')>();
|
||||
return {
|
||||
...actual,
|
||||
getG1CreditBalance: vi.fn(),
|
||||
shouldAutoUseCredits: vi.fn(),
|
||||
shouldShowOverageMenu: vi.fn(),
|
||||
shouldShowEmptyWalletMenu: vi.fn(),
|
||||
logBillingEvent: vi.fn(),
|
||||
};
|
||||
});
|
||||
|
||||
describe('useQuotaAndFallback', () => {
|
||||
let mockConfig: Config;
|
||||
let mockHistoryManager: UseHistoryManagerReturn;
|
||||
@@ -74,6 +94,12 @@ describe('useQuotaAndFallback', () => {
|
||||
vi.spyOn(mockConfig, 'setModel');
|
||||
vi.spyOn(mockConfig, 'setActiveModel');
|
||||
vi.spyOn(mockConfig, 'activateFallbackMode');
|
||||
|
||||
// Mock billing utility functions
|
||||
vi.mocked(getG1CreditBalance).mockReturnValue(0);
|
||||
vi.mocked(shouldAutoUseCredits).mockReturnValue(false);
|
||||
vi.mocked(shouldShowOverageMenu).mockReturnValue(false);
|
||||
vi.mocked(shouldShowEmptyWalletMenu).mockReturnValue(false);
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
@@ -88,6 +114,8 @@ describe('useQuotaAndFallback', () => {
|
||||
userTier: UserTierId.FREE,
|
||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||
onShowAuthSelection: mockOnShowAuthSelection,
|
||||
paidTier: null,
|
||||
settings: mockSettings,
|
||||
}),
|
||||
);
|
||||
|
||||
@@ -105,6 +133,8 @@ describe('useQuotaAndFallback', () => {
|
||||
userTier: UserTierId.FREE,
|
||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||
onShowAuthSelection: mockOnShowAuthSelection,
|
||||
paidTier: null,
|
||||
settings: mockSettings,
|
||||
}),
|
||||
);
|
||||
return setFallbackHandlerSpy.mock.calls[0][0] as FallbackModelHandler;
|
||||
@@ -132,6 +162,8 @@ describe('useQuotaAndFallback', () => {
|
||||
userTier: UserTierId.FREE,
|
||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||
onShowAuthSelection: mockOnShowAuthSelection,
|
||||
paidTier: null,
|
||||
settings: mockSettings,
|
||||
}),
|
||||
);
|
||||
|
||||
@@ -185,6 +217,8 @@ describe('useQuotaAndFallback', () => {
|
||||
userTier: UserTierId.FREE,
|
||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||
onShowAuthSelection: mockOnShowAuthSelection,
|
||||
paidTier: null,
|
||||
settings: mockSettings,
|
||||
}),
|
||||
);
|
||||
|
||||
@@ -224,6 +258,8 @@ describe('useQuotaAndFallback', () => {
|
||||
userTier: UserTierId.FREE,
|
||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||
onShowAuthSelection: mockOnShowAuthSelection,
|
||||
paidTier: null,
|
||||
settings: mockSettings,
|
||||
}),
|
||||
);
|
||||
|
||||
@@ -256,6 +292,8 @@ describe('useQuotaAndFallback', () => {
|
||||
userTier: UserTierId.FREE,
|
||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||
onShowAuthSelection: mockOnShowAuthSelection,
|
||||
paidTier: null,
|
||||
settings: mockSettings,
|
||||
}),
|
||||
);
|
||||
|
||||
@@ -322,6 +360,8 @@ describe('useQuotaAndFallback', () => {
|
||||
setModelSwitchedFromQuotaError:
|
||||
mockSetModelSwitchedFromQuotaError,
|
||||
onShowAuthSelection: mockOnShowAuthSelection,
|
||||
paidTier: null,
|
||||
settings: mockSettings,
|
||||
}),
|
||||
);
|
||||
|
||||
@@ -377,6 +417,8 @@ describe('useQuotaAndFallback', () => {
|
||||
userTier: UserTierId.FREE,
|
||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||
onShowAuthSelection: mockOnShowAuthSelection,
|
||||
paidTier: null,
|
||||
settings: mockSettings,
|
||||
}),
|
||||
);
|
||||
|
||||
@@ -422,6 +464,8 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
|
||||
userTier: UserTierId.FREE,
|
||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||
onShowAuthSelection: mockOnShowAuthSelection,
|
||||
paidTier: null,
|
||||
settings: mockSettings,
|
||||
}),
|
||||
);
|
||||
|
||||
@@ -456,6 +500,243 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
|
||||
});
|
||||
});
|
||||
|
||||
describe('G1 AI Credits Flow', () => {
|
||||
const mockPaidTier = {
|
||||
id: UserTierId.STANDARD,
|
||||
userTier: UserTierId.STANDARD,
|
||||
availableCredits: [
|
||||
{
|
||||
creditType: G1_CREDIT_TYPE,
|
||||
creditAmount: '100',
|
||||
},
|
||||
],
|
||||
};
|
||||
|
||||
beforeEach(() => {
|
||||
// Default to having credits
|
||||
vi.mocked(getG1CreditBalance).mockReturnValue(100);
|
||||
});
|
||||
|
||||
it('should fall through to ProQuotaDialog if credits are already active (strategy=always)', async () => {
|
||||
// If shouldAutoUseCredits is true, credits were already active on the
|
||||
// failed request — they didn't help. Fall through to ProQuotaDialog
|
||||
// so the user can downgrade to Flash instead of retrying infinitely.
|
||||
vi.mocked(shouldAutoUseCredits).mockReturnValue(true);
|
||||
|
||||
const { result } = renderHook(() =>
|
||||
useQuotaAndFallback({
|
||||
config: mockConfig,
|
||||
historyManager: mockHistoryManager,
|
||||
userTier: UserTierId.STANDARD,
|
||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||
onShowAuthSelection: mockOnShowAuthSelection,
|
||||
paidTier: mockPaidTier,
|
||||
settings: mockSettings,
|
||||
}),
|
||||
);
|
||||
|
||||
const handler = setFallbackHandlerSpy.mock
|
||||
.calls[0][0] as FallbackModelHandler;
|
||||
|
||||
const error = new TerminalQuotaError(
|
||||
'pro quota',
|
||||
mockGoogleApiError,
|
||||
1000 * 60 * 5,
|
||||
);
|
||||
|
||||
const intentPromise = handler(
|
||||
PREVIEW_GEMINI_MODEL,
|
||||
'gemini-flash',
|
||||
error,
|
||||
);
|
||||
|
||||
// Since credits didn't help, the ProQuotaDialog should be shown
|
||||
await waitFor(() => {
|
||||
expect(result.current.proQuotaRequest).not.toBeNull();
|
||||
});
|
||||
|
||||
// Resolve it to verify the flow completes
|
||||
act(() => {
|
||||
result.current.handleProQuotaChoice('stop');
|
||||
});
|
||||
|
||||
const intent = await intentPromise;
|
||||
expect(intent).toBe('stop');
|
||||
});
|
||||
|
||||
it('should show overage menu if balance > 0 and not auto-using', async () => {
|
||||
vi.mocked(shouldAutoUseCredits).mockReturnValue(false);
|
||||
vi.mocked(shouldShowOverageMenu).mockReturnValue(true);
|
||||
|
||||
const { result } = renderHook(() =>
|
||||
useQuotaAndFallback({
|
||||
config: mockConfig,
|
||||
historyManager: mockHistoryManager,
|
||||
userTier: UserTierId.STANDARD,
|
||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||
onShowAuthSelection: mockOnShowAuthSelection,
|
||||
paidTier: mockPaidTier,
|
||||
settings: mockSettings,
|
||||
}),
|
||||
);
|
||||
|
||||
const handler = setFallbackHandlerSpy.mock
|
||||
.calls[0][0] as FallbackModelHandler;
|
||||
|
||||
let promise: Promise<FallbackIntent | null>;
|
||||
act(() => {
|
||||
promise = handler(
|
||||
PREVIEW_GEMINI_MODEL,
|
||||
'gemini-flash',
|
||||
new TerminalQuotaError('pro quota', mockGoogleApiError),
|
||||
);
|
||||
});
|
||||
|
||||
expect(result.current.overageMenuRequest).not.toBeNull();
|
||||
expect(result.current.overageMenuRequest?.creditBalance).toBe(100);
|
||||
expect(logBillingEvent).toHaveBeenCalled();
|
||||
|
||||
// Simulate choosing "Use Credits"
|
||||
await act(async () => {
|
||||
result.current.handleOverageMenuChoice('use_credits');
|
||||
await promise!;
|
||||
});
|
||||
|
||||
const intent = await promise!;
|
||||
expect(intent).toBe('retry_with_credits');
|
||||
});
|
||||
|
||||
it('should handle use_fallback from overage menu', async () => {
|
||||
vi.mocked(shouldAutoUseCredits).mockReturnValue(false);
|
||||
vi.mocked(shouldShowOverageMenu).mockReturnValue(true);
|
||||
|
||||
const { result } = renderHook(() =>
|
||||
useQuotaAndFallback({
|
||||
config: mockConfig,
|
||||
historyManager: mockHistoryManager,
|
||||
userTier: UserTierId.STANDARD,
|
||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||
onShowAuthSelection: mockOnShowAuthSelection,
|
||||
paidTier: mockPaidTier,
|
||||
settings: mockSettings,
|
||||
}),
|
||||
);
|
||||
|
||||
const handler = setFallbackHandlerSpy.mock
|
||||
.calls[0][0] as FallbackModelHandler;
|
||||
|
||||
let promise: Promise<FallbackIntent | null>;
|
||||
act(() => {
|
||||
promise = handler(
|
||||
PREVIEW_GEMINI_MODEL,
|
||||
'gemini-flash',
|
||||
new TerminalQuotaError('pro quota', mockGoogleApiError),
|
||||
);
|
||||
});
|
||||
|
||||
// Simulate choosing "Switch to fallback"
|
||||
await act(async () => {
|
||||
result.current.handleOverageMenuChoice('use_fallback');
|
||||
await promise!;
|
||||
});
|
||||
|
||||
const intent = await promise!;
|
||||
expect(intent).toBe('retry_always');
|
||||
});
|
||||
|
||||
it('should show empty wallet menu if balance is 0', async () => {
|
||||
vi.mocked(getG1CreditBalance).mockReturnValue(0);
|
||||
vi.mocked(shouldAutoUseCredits).mockReturnValue(false);
|
||||
vi.mocked(shouldShowOverageMenu).mockReturnValue(false);
|
||||
vi.mocked(shouldShowEmptyWalletMenu).mockReturnValue(true);
|
||||
|
||||
const { result } = renderHook(() =>
|
||||
useQuotaAndFallback({
|
||||
config: mockConfig,
|
||||
historyManager: mockHistoryManager,
|
||||
userTier: UserTierId.STANDARD,
|
||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||
onShowAuthSelection: mockOnShowAuthSelection,
|
||||
paidTier: { ...mockPaidTier, availableCredits: [] },
|
||||
settings: mockSettings,
|
||||
}),
|
||||
);
|
||||
|
||||
const handler = setFallbackHandlerSpy.mock
|
||||
.calls[0][0] as FallbackModelHandler;
|
||||
|
||||
let promise: Promise<FallbackIntent | null>;
|
||||
act(() => {
|
||||
promise = handler(
|
||||
PREVIEW_GEMINI_MODEL,
|
||||
'gemini-flash',
|
||||
new TerminalQuotaError('pro quota', mockGoogleApiError),
|
||||
);
|
||||
});
|
||||
|
||||
expect(result.current.emptyWalletRequest).not.toBeNull();
|
||||
expect(logBillingEvent).toHaveBeenCalled();
|
||||
|
||||
// Simulate choosing "Stop"
|
||||
await act(async () => {
|
||||
result.current.handleEmptyWalletChoice('stop');
|
||||
await promise!;
|
||||
});
|
||||
|
||||
const intent = await promise!;
|
||||
expect(intent).toBe('stop');
|
||||
});
|
||||
|
||||
it('should add info message to history when get_credits is selected', async () => {
|
||||
vi.mocked(getG1CreditBalance).mockReturnValue(0);
|
||||
vi.mocked(shouldAutoUseCredits).mockReturnValue(false);
|
||||
vi.mocked(shouldShowOverageMenu).mockReturnValue(false);
|
||||
vi.mocked(shouldShowEmptyWalletMenu).mockReturnValue(true);
|
||||
|
||||
const { result } = renderHook(() =>
|
||||
useQuotaAndFallback({
|
||||
config: mockConfig,
|
||||
historyManager: mockHistoryManager,
|
||||
userTier: UserTierId.STANDARD,
|
||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||
onShowAuthSelection: mockOnShowAuthSelection,
|
||||
paidTier: { ...mockPaidTier, availableCredits: [] },
|
||||
settings: mockSettings,
|
||||
}),
|
||||
);
|
||||
|
||||
const handler = setFallbackHandlerSpy.mock
|
||||
.calls[0][0] as FallbackModelHandler;
|
||||
|
||||
let promise: Promise<FallbackIntent | null>;
|
||||
act(() => {
|
||||
promise = handler(
|
||||
PREVIEW_GEMINI_MODEL,
|
||||
'gemini-flash',
|
||||
new TerminalQuotaError('pro quota', mockGoogleApiError),
|
||||
);
|
||||
});
|
||||
|
||||
expect(result.current.emptyWalletRequest).not.toBeNull();
|
||||
|
||||
// Simulate choosing "Get AI Credits"
|
||||
await act(async () => {
|
||||
result.current.handleEmptyWalletChoice('get_credits');
|
||||
await promise!;
|
||||
});
|
||||
|
||||
const intent = await promise!;
|
||||
expect(intent).toBe('stop');
|
||||
expect(mockHistoryManager.addItem).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
type: MessageType.INFO,
|
||||
text: expect.stringContaining('few minutes'),
|
||||
}),
|
||||
expect.any(Number),
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe('handleProQuotaChoice', () => {
|
||||
it('should do nothing if there is no pending pro quota request', () => {
|
||||
const { result } = renderHook(() =>
|
||||
@@ -465,6 +746,8 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
|
||||
userTier: UserTierId.FREE,
|
||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||
onShowAuthSelection: mockOnShowAuthSelection,
|
||||
paidTier: null,
|
||||
settings: mockSettings,
|
||||
}),
|
||||
);
|
||||
|
||||
@@ -483,6 +766,8 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
|
||||
userTier: UserTierId.FREE,
|
||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||
onShowAuthSelection: mockOnShowAuthSelection,
|
||||
paidTier: null,
|
||||
settings: mockSettings,
|
||||
}),
|
||||
);
|
||||
|
||||
@@ -514,6 +799,8 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
|
||||
userTier: UserTierId.FREE,
|
||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||
onShowAuthSelection: mockOnShowAuthSelection,
|
||||
paidTier: null,
|
||||
settings: mockSettings,
|
||||
}),
|
||||
);
|
||||
|
||||
@@ -558,6 +845,8 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
|
||||
userTier: UserTierId.FREE,
|
||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||
onShowAuthSelection: mockOnShowAuthSelection,
|
||||
paidTier: null,
|
||||
settings: mockSettings,
|
||||
}),
|
||||
);
|
||||
|
||||
@@ -594,6 +883,8 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
|
||||
userTier: UserTierId.FREE,
|
||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||
onShowAuthSelection: mockOnShowAuthSelection,
|
||||
paidTier: null,
|
||||
settings: mockSettings,
|
||||
}),
|
||||
);
|
||||
|
||||
@@ -638,6 +929,8 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
|
||||
userTier: UserTierId.FREE,
|
||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||
onShowAuthSelection: mockOnShowAuthSelection,
|
||||
paidTier: null,
|
||||
settings: mockSettings,
|
||||
}),
|
||||
);
|
||||
|
||||
@@ -653,6 +946,8 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
|
||||
userTier: UserTierId.FREE,
|
||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||
onShowAuthSelection: mockOnShowAuthSelection,
|
||||
paidTier: null,
|
||||
settings: mockSettings,
|
||||
}),
|
||||
);
|
||||
|
||||
@@ -695,6 +990,8 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
|
||||
userTier: UserTierId.FREE,
|
||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||
onShowAuthSelection: mockOnShowAuthSelection,
|
||||
paidTier: null,
|
||||
settings: mockSettings,
|
||||
}),
|
||||
);
|
||||
|
||||
@@ -737,6 +1034,8 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
|
||||
userTier: UserTierId.FREE,
|
||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||
onShowAuthSelection: mockOnShowAuthSelection,
|
||||
paidTier: null,
|
||||
settings: mockSettings,
|
||||
}),
|
||||
);
|
||||
|
||||
@@ -767,6 +1066,8 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
|
||||
userTier: UserTierId.FREE,
|
||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||
onShowAuthSelection: mockOnShowAuthSelection,
|
||||
paidTier: null,
|
||||
settings: mockSettings,
|
||||
}),
|
||||
);
|
||||
|
||||
@@ -797,6 +1098,8 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
|
||||
userTier: UserTierId.FREE,
|
||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||
onShowAuthSelection: mockOnShowAuthSelection,
|
||||
paidTier: null,
|
||||
settings: mockSettings,
|
||||
}),
|
||||
);
|
||||
|
||||
|
||||
Reference in New Issue
Block a user