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

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

View File

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

View File

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

View File

@@ -0,0 +1,7 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
export * from './billing.js';