mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-04-27 13:34:15 -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:
@@ -0,0 +1,254 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect } from 'vitest';
|
||||
import type { GeminiUserTier } from '../code_assist/types.js';
|
||||
import {
|
||||
buildG1Url,
|
||||
getG1CreditBalance,
|
||||
G1_CREDIT_TYPE,
|
||||
G1_UTM_CAMPAIGNS,
|
||||
isOverageEligibleModel,
|
||||
shouldAutoUseCredits,
|
||||
shouldShowEmptyWalletMenu,
|
||||
shouldShowOverageMenu,
|
||||
wrapInAccountChooser,
|
||||
} from './billing.js';
|
||||
|
||||
describe('billing', () => {
|
||||
describe('wrapInAccountChooser', () => {
|
||||
it('should wrap URL with AccountChooser redirect', () => {
|
||||
const result = wrapInAccountChooser(
|
||||
'user@gmail.com',
|
||||
'https://one.google.com/ai/activity',
|
||||
);
|
||||
expect(result).toBe(
|
||||
'https://accounts.google.com/AccountChooser?Email=user%40gmail.com&continue=https%3A%2F%2Fone.google.com%2Fai%2Factivity',
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle special characters in email', () => {
|
||||
const result = wrapInAccountChooser(
|
||||
'user+test@example.com',
|
||||
'https://example.com',
|
||||
);
|
||||
expect(result).toContain('Email=user%2Btest%40example.com');
|
||||
});
|
||||
});
|
||||
|
||||
describe('buildG1Url', () => {
|
||||
it('should build activity URL with UTM params wrapped in AccountChooser', () => {
|
||||
const result = buildG1Url(
|
||||
'activity',
|
||||
'user@gmail.com',
|
||||
G1_UTM_CAMPAIGNS.MANAGE_ACTIVITY,
|
||||
);
|
||||
|
||||
// Should contain AccountChooser prefix
|
||||
expect(result).toContain('https://accounts.google.com/AccountChooser');
|
||||
expect(result).toContain('Email=user%40gmail.com');
|
||||
|
||||
// The continue URL should contain the G1 activity path and UTM params
|
||||
expect(result).toContain('one.google.com%2Fai%2Factivity');
|
||||
expect(result).toContain('utm_source%3Dgemini_cli');
|
||||
expect(result).toContain(
|
||||
'utm_campaign%3Dhydrogen_cli_settings_ai_credits_activity_page',
|
||||
);
|
||||
});
|
||||
|
||||
it('should build credits URL with UTM params wrapped in AccountChooser', () => {
|
||||
const result = buildG1Url(
|
||||
'credits',
|
||||
'test@example.com',
|
||||
G1_UTM_CAMPAIGNS.EMPTY_WALLET_ADD_CREDITS,
|
||||
);
|
||||
|
||||
expect(result).toContain('https://accounts.google.com/AccountChooser');
|
||||
expect(result).toContain('one.google.com%2Fai%2Fcredits');
|
||||
expect(result).toContain(
|
||||
'utm_campaign%3Dhydrogen_cli_insufficient_credits_add_credits',
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe('getG1CreditBalance', () => {
|
||||
it('should return null for null tier', () => {
|
||||
expect(getG1CreditBalance(null)).toBeNull();
|
||||
});
|
||||
|
||||
it('should return null for undefined tier', () => {
|
||||
expect(getG1CreditBalance(undefined)).toBeNull();
|
||||
});
|
||||
|
||||
it('should return null for tier without availableCredits', () => {
|
||||
const tier: GeminiUserTier = { id: 'PERSONAL' };
|
||||
expect(getG1CreditBalance(tier)).toBeNull();
|
||||
});
|
||||
|
||||
it('should return null for empty availableCredits array', () => {
|
||||
const tier: GeminiUserTier = { id: 'PERSONAL', availableCredits: [] };
|
||||
expect(getG1CreditBalance(tier)).toBeNull();
|
||||
});
|
||||
|
||||
it('should return null when no G1 credit type found', () => {
|
||||
const tier: GeminiUserTier = {
|
||||
id: 'PERSONAL',
|
||||
availableCredits: [
|
||||
{ creditType: 'CREDIT_TYPE_UNSPECIFIED', creditAmount: '100' },
|
||||
],
|
||||
};
|
||||
expect(getG1CreditBalance(tier)).toBeNull();
|
||||
});
|
||||
|
||||
it('should return G1 credit balance when present', () => {
|
||||
const tier: GeminiUserTier = {
|
||||
id: 'PERSONAL',
|
||||
availableCredits: [{ creditType: G1_CREDIT_TYPE, creditAmount: '500' }],
|
||||
};
|
||||
expect(getG1CreditBalance(tier)).toBe(500);
|
||||
});
|
||||
|
||||
it('should return G1 credit balance when multiple credit types present', () => {
|
||||
const tier: GeminiUserTier = {
|
||||
id: 'PERSONAL',
|
||||
availableCredits: [
|
||||
{ creditType: 'CREDIT_TYPE_UNSPECIFIED', creditAmount: '100' },
|
||||
{ creditType: G1_CREDIT_TYPE, creditAmount: '750' },
|
||||
],
|
||||
};
|
||||
expect(getG1CreditBalance(tier)).toBe(750);
|
||||
});
|
||||
|
||||
it('should return 0 for invalid credit amount', () => {
|
||||
const tier: GeminiUserTier = {
|
||||
id: 'PERSONAL',
|
||||
availableCredits: [
|
||||
{ creditType: G1_CREDIT_TYPE, creditAmount: 'invalid' },
|
||||
],
|
||||
};
|
||||
expect(getG1CreditBalance(tier)).toBe(0);
|
||||
});
|
||||
|
||||
it('should handle large credit amounts (int64 as string)', () => {
|
||||
const tier: GeminiUserTier = {
|
||||
id: 'PERSONAL',
|
||||
availableCredits: [
|
||||
{ creditType: G1_CREDIT_TYPE, creditAmount: '9999999999' },
|
||||
],
|
||||
};
|
||||
expect(getG1CreditBalance(tier)).toBe(9999999999);
|
||||
});
|
||||
|
||||
it('should sum multiple credits of the same G1 type', () => {
|
||||
const tier: GeminiUserTier = {
|
||||
id: 'PERSONAL',
|
||||
availableCredits: [
|
||||
{ creditType: G1_CREDIT_TYPE, creditAmount: '1000' },
|
||||
{ creditType: G1_CREDIT_TYPE, creditAmount: '8' },
|
||||
],
|
||||
};
|
||||
expect(getG1CreditBalance(tier)).toBe(1008);
|
||||
});
|
||||
});
|
||||
|
||||
describe('shouldAutoUseCredits', () => {
|
||||
it('should return true when strategy is always and balance > 0', () => {
|
||||
expect(shouldAutoUseCredits('always', 100)).toBe(true);
|
||||
});
|
||||
|
||||
it('should return false when strategy is always but balance is 0', () => {
|
||||
expect(shouldAutoUseCredits('always', 0)).toBe(false);
|
||||
});
|
||||
|
||||
it('should return false when strategy is ask', () => {
|
||||
expect(shouldAutoUseCredits('ask', 100)).toBe(false);
|
||||
});
|
||||
|
||||
it('should return false when strategy is never', () => {
|
||||
expect(shouldAutoUseCredits('never', 100)).toBe(false);
|
||||
});
|
||||
|
||||
it('should return false when creditBalance is null (ineligible)', () => {
|
||||
expect(shouldAutoUseCredits('always', null)).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe('shouldShowOverageMenu', () => {
|
||||
it('should return true when strategy is ask and balance > 0', () => {
|
||||
expect(shouldShowOverageMenu('ask', 100)).toBe(true);
|
||||
});
|
||||
|
||||
it('should return false when strategy is ask but balance is 0', () => {
|
||||
expect(shouldShowOverageMenu('ask', 0)).toBe(false);
|
||||
});
|
||||
|
||||
it('should return false when strategy is always', () => {
|
||||
expect(shouldShowOverageMenu('always', 100)).toBe(false);
|
||||
});
|
||||
|
||||
it('should return false when strategy is never', () => {
|
||||
expect(shouldShowOverageMenu('never', 100)).toBe(false);
|
||||
});
|
||||
|
||||
it('should return false when creditBalance is null (ineligible)', () => {
|
||||
expect(shouldShowOverageMenu('ask', null)).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe('shouldShowEmptyWalletMenu', () => {
|
||||
it('should return true when strategy is ask and balance is 0', () => {
|
||||
expect(shouldShowEmptyWalletMenu('ask', 0)).toBe(true);
|
||||
});
|
||||
|
||||
it('should return true when strategy is always and balance is 0', () => {
|
||||
expect(shouldShowEmptyWalletMenu('always', 0)).toBe(true);
|
||||
});
|
||||
|
||||
it('should return false when strategy is never', () => {
|
||||
expect(shouldShowEmptyWalletMenu('never', 0)).toBe(false);
|
||||
});
|
||||
|
||||
it('should return false when balance > 0', () => {
|
||||
expect(shouldShowEmptyWalletMenu('ask', 100)).toBe(false);
|
||||
});
|
||||
|
||||
it('should return false when creditBalance is null (ineligible)', () => {
|
||||
expect(shouldShowEmptyWalletMenu('ask', null)).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe('isOverageEligibleModel', () => {
|
||||
it('should return true for gemini-3-pro-preview', () => {
|
||||
expect(isOverageEligibleModel('gemini-3-pro-preview')).toBe(true);
|
||||
});
|
||||
|
||||
it('should return true for gemini-3.1-pro-preview', () => {
|
||||
expect(isOverageEligibleModel('gemini-3.1-pro-preview')).toBe(false);
|
||||
});
|
||||
|
||||
it('should return true for gemini-3.1-pro-preview-customtools', () => {
|
||||
expect(isOverageEligibleModel('gemini-3.1-pro-preview-customtools')).toBe(
|
||||
false,
|
||||
);
|
||||
});
|
||||
|
||||
it('should return false for gemini-3-flash-preview', () => {
|
||||
expect(isOverageEligibleModel('gemini-3-flash-preview')).toBe(false);
|
||||
});
|
||||
|
||||
it('should return false for gemini-2.5-pro', () => {
|
||||
expect(isOverageEligibleModel('gemini-2.5-pro')).toBe(false);
|
||||
});
|
||||
|
||||
it('should return false for gemini-2.5-flash', () => {
|
||||
expect(isOverageEligibleModel('gemini-2.5-flash')).toBe(false);
|
||||
});
|
||||
|
||||
it('should return false for custom model names', () => {
|
||||
expect(isOverageEligibleModel('my-custom-model')).toBe(false);
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,180 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type {
|
||||
AvailableCredits,
|
||||
CreditType,
|
||||
GeminiUserTier,
|
||||
} from '../code_assist/types.js';
|
||||
import { PREVIEW_GEMINI_MODEL } from '../config/models.js';
|
||||
|
||||
/**
|
||||
* Strategy for handling quota exhaustion when AI credits are available.
|
||||
* - 'ask': Prompt the user each time
|
||||
* - 'always': Automatically use credits
|
||||
* - 'never': Never use credits, show standard fallback
|
||||
*/
|
||||
export type OverageStrategy = 'ask' | 'always' | 'never';
|
||||
|
||||
/** Credit type for Google One AI credits */
|
||||
export const G1_CREDIT_TYPE: CreditType = 'GOOGLE_ONE_AI';
|
||||
|
||||
/**
|
||||
* The set of models that support AI credits overage billing.
|
||||
* Only these models are eligible for the credits-based retry flow.
|
||||
*/
|
||||
export const OVERAGE_ELIGIBLE_MODELS = new Set([PREVIEW_GEMINI_MODEL]);
|
||||
|
||||
/**
|
||||
* Checks if a model is eligible for AI credits overage billing.
|
||||
* @param model The model name to check.
|
||||
* @returns true if the model supports credits overage, false otherwise.
|
||||
*/
|
||||
export function isOverageEligibleModel(model: string): boolean {
|
||||
return OVERAGE_ELIGIBLE_MODELS.has(model);
|
||||
}
|
||||
|
||||
/** Base URL for Google One AI page */
|
||||
const G1_AI_BASE_URL = 'https://one.google.com/ai';
|
||||
|
||||
/** AccountChooser URL for redirecting with email context */
|
||||
const ACCOUNT_CHOOSER_URL = 'https://accounts.google.com/AccountChooser';
|
||||
|
||||
/** UTM parameters for CLI tracking */
|
||||
const UTM_SOURCE = 'gemini_cli';
|
||||
// TODO: change to 'desktop' when G1 service fix is rolled out
|
||||
const UTM_MEDIUM = 'web';
|
||||
|
||||
/**
|
||||
* Wraps a URL in the AccountChooser redirect to maintain user context.
|
||||
* @param email User's email address for account selection
|
||||
* @param continueUrl The destination URL after account selection
|
||||
* @returns The full AccountChooser redirect URL
|
||||
*/
|
||||
export function wrapInAccountChooser(
|
||||
email: string,
|
||||
continueUrl: string,
|
||||
): string {
|
||||
const params = new URLSearchParams({
|
||||
Email: email,
|
||||
continue: continueUrl,
|
||||
});
|
||||
return `${ACCOUNT_CHOOSER_URL}?${params.toString()}`;
|
||||
}
|
||||
|
||||
/**
|
||||
* UTM campaign identifiers per the design doc.
|
||||
*/
|
||||
export const G1_UTM_CAMPAIGNS = {
|
||||
/** From Interception Flow "Manage" link (user has credits) */
|
||||
MANAGE_ACTIVITY: 'hydrogen_cli_settings_ai_credits_activity_page',
|
||||
/** From "Manage" to add more credits */
|
||||
MANAGE_ADD_CREDITS: 'hydrogen_cli_settings_add_credits',
|
||||
/** From Empty Wallet Flow "Get AI Credits" link */
|
||||
EMPTY_WALLET_ADD_CREDITS: 'hydrogen_cli_insufficient_credits_add_credits',
|
||||
} as const;
|
||||
|
||||
/**
|
||||
* Builds a G1 AI URL with UTM tracking parameters.
|
||||
* @param path The path segment (e.g., 'activity' or 'credits')
|
||||
* @param email User's email for AccountChooser wrapper
|
||||
* @param campaign The UTM campaign identifier
|
||||
* @returns The complete URL wrapped in AccountChooser
|
||||
*/
|
||||
export function buildG1Url(
|
||||
path: 'activity' | 'credits',
|
||||
email: string,
|
||||
campaign: string,
|
||||
): string {
|
||||
const baseUrl = `${G1_AI_BASE_URL}/${path}`;
|
||||
const params = new URLSearchParams({
|
||||
utm_source: UTM_SOURCE,
|
||||
utm_medium: UTM_MEDIUM,
|
||||
utm_campaign: campaign,
|
||||
});
|
||||
const urlWithUtm = `${baseUrl}?${params.toString()}`;
|
||||
return wrapInAccountChooser(email, urlWithUtm);
|
||||
}
|
||||
|
||||
/**
|
||||
* Extracts the G1 AI credit balance from a tier's available credits.
|
||||
* @param tier The user tier to check
|
||||
* @returns The credit amount as a number, 0 if eligible but empty, or null if not eligible
|
||||
*/
|
||||
export function getG1CreditBalance(
|
||||
tier: GeminiUserTier | null | undefined,
|
||||
): number | null {
|
||||
if (!tier?.availableCredits) {
|
||||
return null;
|
||||
}
|
||||
|
||||
const g1Credits = tier.availableCredits.filter(
|
||||
(credit: AvailableCredits) => credit.creditType === G1_CREDIT_TYPE,
|
||||
);
|
||||
|
||||
if (g1Credits.length === 0) {
|
||||
return null;
|
||||
}
|
||||
|
||||
// creditAmount is an int64 represented as string; sum all matching entries
|
||||
return g1Credits.reduce((sum, credit) => {
|
||||
const amount = parseInt(credit.creditAmount ?? '0', 10);
|
||||
return sum + (isNaN(amount) ? 0 : amount);
|
||||
}, 0);
|
||||
}
|
||||
|
||||
export const MIN_CREDIT_BALANCE = 50;
|
||||
|
||||
/**
|
||||
* Determines if credits should be automatically used based on the overage strategy.
|
||||
* @param strategy The configured overage strategy
|
||||
* @param creditBalance The available credit balance
|
||||
* @returns true if credits should be auto-used, false otherwise
|
||||
*/
|
||||
export function shouldAutoUseCredits(
|
||||
strategy: OverageStrategy,
|
||||
creditBalance: number | null,
|
||||
): boolean {
|
||||
return (
|
||||
strategy === 'always' &&
|
||||
creditBalance != null &&
|
||||
creditBalance >= MIN_CREDIT_BALANCE
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Determines if the overage menu should be shown based on the strategy.
|
||||
* @param strategy The configured overage strategy
|
||||
* @param creditBalance The available credit balance
|
||||
* @returns true if the menu should be shown
|
||||
*/
|
||||
export function shouldShowOverageMenu(
|
||||
strategy: OverageStrategy,
|
||||
creditBalance: number | null,
|
||||
): boolean {
|
||||
return (
|
||||
strategy === 'ask' &&
|
||||
creditBalance != null &&
|
||||
creditBalance >= MIN_CREDIT_BALANCE
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Determines if the empty wallet menu should be shown.
|
||||
* @param strategy The configured overage strategy
|
||||
* @param creditBalance The available credit balance
|
||||
* @returns true if the empty wallet menu should be shown
|
||||
*/
|
||||
export function shouldShowEmptyWalletMenu(
|
||||
strategy: OverageStrategy,
|
||||
creditBalance: number | null,
|
||||
): boolean {
|
||||
return (
|
||||
strategy !== 'never' &&
|
||||
creditBalance != null &&
|
||||
creditBalance < MIN_CREDIT_BALANCE
|
||||
);
|
||||
}
|
||||
@@ -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',
|
||||
'free-tier',
|
||||
'free-tier-name',
|
||||
undefined,
|
||||
mockConfig,
|
||||
);
|
||||
expect(generator).toBeInstanceOf(MockedCodeAssistServer);
|
||||
});
|
||||
@@ -103,6 +105,8 @@ describe('codeAssist', () => {
|
||||
undefined, // No session ID
|
||||
'free-tier',
|
||||
'free-tier-name',
|
||||
undefined,
|
||||
mockConfig,
|
||||
);
|
||||
expect(generator).toBeInstanceOf(MockedCodeAssistServer);
|
||||
});
|
||||
|
||||
@@ -36,6 +36,8 @@ export async function createCodeAssistContentGenerator(
|
||||
sessionId,
|
||||
userData.userTier,
|
||||
userData.userTierName,
|
||||
userData.paidTier,
|
||||
config,
|
||||
);
|
||||
}
|
||||
|
||||
|
||||
@@ -27,12 +27,14 @@ import type {
|
||||
ToolConfig,
|
||||
} from '@google/genai';
|
||||
import { GenerateContentResponse } from '@google/genai';
|
||||
import type { Credits } from './types.js';
|
||||
|
||||
export interface CAGenerateContentRequest {
|
||||
model: string;
|
||||
project?: string;
|
||||
user_prompt_id?: string;
|
||||
request: VertexGenerateContentRequest;
|
||||
enabled_credit_types?: string[];
|
||||
}
|
||||
|
||||
interface VertexGenerateContentRequest {
|
||||
@@ -74,6 +76,8 @@ interface VertexGenerationConfig {
|
||||
export interface CaGenerateContentResponse {
|
||||
response: VertexGenerateContentResponse;
|
||||
traceId?: string;
|
||||
consumedCredits?: Credits[];
|
||||
remainingCredits?: Credits[];
|
||||
}
|
||||
|
||||
interface VertexGenerateContentResponse {
|
||||
@@ -121,12 +125,14 @@ export function toGenerateContentRequest(
|
||||
userPromptId: string,
|
||||
project?: string,
|
||||
sessionId?: string,
|
||||
enabledCreditTypes?: string[],
|
||||
): CAGenerateContentRequest {
|
||||
return {
|
||||
model: req.model,
|
||||
project,
|
||||
user_prompt_id: userPromptId,
|
||||
request: toVertexGenerateContentRequest(req, sessionId),
|
||||
enabled_credit_types: enabledCreditTypes,
|
||||
};
|
||||
}
|
||||
|
||||
@@ -299,3 +305,16 @@ function toVertexGenerationConfig(
|
||||
thinkingConfig: config.thinkingConfig,
|
||||
};
|
||||
}
|
||||
|
||||
export function fromGenerateContentResponseUsage(
|
||||
metadata?: GenerateContentResponseUsageMetadata,
|
||||
): GenerateContentResponseUsageMetadata | undefined {
|
||||
if (!metadata) {
|
||||
return undefined;
|
||||
}
|
||||
return {
|
||||
promptTokenCount: metadata.promptTokenCount,
|
||||
candidatesTokenCount: metadata.candidatesTokenCount,
|
||||
totalTokenCount: metadata.totalTokenCount,
|
||||
};
|
||||
}
|
||||
|
||||
@@ -21,7 +21,10 @@ import type {
|
||||
ConversationInteraction,
|
||||
StreamingLatency,
|
||||
RecordCodeAssistMetricsRequest,
|
||||
GeminiUserTier,
|
||||
Credits,
|
||||
} from './types.js';
|
||||
import { UserTierId } from './types.js';
|
||||
import type {
|
||||
ListExperimentsRequest,
|
||||
ListExperimentsResponse,
|
||||
@@ -37,7 +40,15 @@ import type {
|
||||
import * as readline from 'node:readline';
|
||||
import { Readable } from 'node:stream';
|
||||
import type { ContentGenerator } from '../core/contentGenerator.js';
|
||||
import { UserTierId } from './types.js';
|
||||
import type { Config } from '../config/config.js';
|
||||
import {
|
||||
G1_CREDIT_TYPE,
|
||||
getG1CreditBalance,
|
||||
isOverageEligibleModel,
|
||||
shouldAutoUseCredits,
|
||||
} from '../billing/billing.js';
|
||||
import { logBillingEvent } from '../telemetry/loggers.js';
|
||||
import { CreditsUsedEvent } from '../telemetry/billingEvents.js';
|
||||
import type {
|
||||
CaCountTokenResponse,
|
||||
CaGenerateContentResponse,
|
||||
@@ -71,6 +82,8 @@ export class CodeAssistServer implements ContentGenerator {
|
||||
readonly sessionId?: string,
|
||||
readonly userTier?: UserTierId,
|
||||
readonly userTierName?: string,
|
||||
readonly paidTier?: GeminiUserTier,
|
||||
readonly config?: Config,
|
||||
) {}
|
||||
|
||||
async generateContentStream(
|
||||
@@ -79,6 +92,19 @@ export class CodeAssistServer implements ContentGenerator {
|
||||
// eslint-disable-next-line @typescript-eslint/no-unused-vars
|
||||
role: LlmRole,
|
||||
): Promise<AsyncGenerator<GenerateContentResponse>> {
|
||||
const autoUse = this.config
|
||||
? shouldAutoUseCredits(
|
||||
this.config.getBillingSettings().overageStrategy,
|
||||
getG1CreditBalance(this.paidTier),
|
||||
)
|
||||
: false;
|
||||
const modelIsEligible = isOverageEligibleModel(req.model);
|
||||
const shouldEnableCredits = modelIsEligible && autoUse;
|
||||
|
||||
const enabledCreditTypes = shouldEnableCredits
|
||||
? ([G1_CREDIT_TYPE] as string[])
|
||||
: undefined;
|
||||
|
||||
const responses =
|
||||
await this.requestStreamingPost<CaGenerateContentResponse>(
|
||||
'streamGenerateContent',
|
||||
@@ -87,6 +113,7 @@ export class CodeAssistServer implements ContentGenerator {
|
||||
userPromptId,
|
||||
this.projectId,
|
||||
this.sessionId,
|
||||
enabledCreditTypes,
|
||||
),
|
||||
req.config?.abortSignal,
|
||||
);
|
||||
@@ -98,6 +125,9 @@ export class CodeAssistServer implements ContentGenerator {
|
||||
return (async function* (
|
||||
server: CodeAssistServer,
|
||||
): AsyncGenerator<GenerateContentResponse> {
|
||||
let totalConsumed = 0;
|
||||
let lastRemaining = 0;
|
||||
|
||||
for await (const response of responses) {
|
||||
if (isFirst) {
|
||||
streamingLatency.firstMessageLatency = formatProtoJsonDuration(
|
||||
@@ -120,8 +150,38 @@ export class CodeAssistServer implements ContentGenerator {
|
||||
req.config?.abortSignal,
|
||||
);
|
||||
|
||||
if (response.consumedCredits) {
|
||||
for (const credit of response.consumedCredits) {
|
||||
if (credit.creditType === G1_CREDIT_TYPE && credit.creditAmount) {
|
||||
totalConsumed += parseInt(credit.creditAmount, 10) || 0;
|
||||
}
|
||||
}
|
||||
}
|
||||
if (response.remainingCredits) {
|
||||
// Sum all G1 credit entries for consistency with getG1CreditBalance
|
||||
lastRemaining = response.remainingCredits.reduce((sum, credit) => {
|
||||
if (credit.creditType === G1_CREDIT_TYPE && credit.creditAmount) {
|
||||
return sum + (parseInt(credit.creditAmount, 10) || 0);
|
||||
}
|
||||
return sum;
|
||||
}, 0);
|
||||
server.updateCredits(response.remainingCredits);
|
||||
}
|
||||
|
||||
yield translatedResponse;
|
||||
}
|
||||
|
||||
// Emit credits used telemetry after the stream completes
|
||||
if (totalConsumed > 0 && server.config) {
|
||||
logBillingEvent(
|
||||
server.config,
|
||||
new CreditsUsedEvent(
|
||||
req.model ?? 'unknown',
|
||||
totalConsumed,
|
||||
lastRemaining,
|
||||
),
|
||||
);
|
||||
}
|
||||
})(this);
|
||||
}
|
||||
|
||||
@@ -139,6 +199,7 @@ export class CodeAssistServer implements ContentGenerator {
|
||||
userPromptId,
|
||||
this.projectId,
|
||||
this.sessionId,
|
||||
undefined,
|
||||
),
|
||||
req.config?.abortSignal,
|
||||
);
|
||||
@@ -158,9 +219,29 @@ export class CodeAssistServer implements ContentGenerator {
|
||||
req.config?.abortSignal,
|
||||
);
|
||||
|
||||
if (response.remainingCredits) {
|
||||
this.updateCredits(response.remainingCredits);
|
||||
}
|
||||
|
||||
return translatedResponse;
|
||||
}
|
||||
|
||||
private updateCredits(remainingCredits: Credits[]): void {
|
||||
if (!this.paidTier) {
|
||||
return;
|
||||
}
|
||||
|
||||
// Replace the G1 credits entries with the latest remaining amounts.
|
||||
// Non-G1 credits are preserved as-is.
|
||||
const nonG1Credits = (this.paidTier.availableCredits ?? []).filter(
|
||||
(c) => c.creditType !== G1_CREDIT_TYPE,
|
||||
);
|
||||
const updatedG1Credits = remainingCredits.filter(
|
||||
(c) => c.creditType === G1_CREDIT_TYPE,
|
||||
);
|
||||
this.paidTier.availableCredits = [...nonG1Credits, ...updatedG1Credits];
|
||||
}
|
||||
|
||||
async onboardUser(
|
||||
req: OnboardUserRequest,
|
||||
): Promise<LongRunningOperationResponse> {
|
||||
@@ -190,6 +271,25 @@ export class CodeAssistServer implements ContentGenerator {
|
||||
}
|
||||
}
|
||||
|
||||
async refreshAvailableCredits(): Promise<void> {
|
||||
if (!this.paidTier) {
|
||||
return;
|
||||
}
|
||||
const res = await this.loadCodeAssist({
|
||||
cloudaicompanionProject: this.projectId,
|
||||
metadata: {
|
||||
ideType: 'IDE_UNSPECIFIED',
|
||||
platform: 'PLATFORM_UNSPECIFIED',
|
||||
pluginType: 'GEMINI',
|
||||
duetProject: this.projectId,
|
||||
},
|
||||
mode: 'HEALTH_CHECK',
|
||||
});
|
||||
if (res.paidTier?.availableCredits) {
|
||||
this.paidTier.availableCredits = res.paidTier.availableCredits;
|
||||
}
|
||||
}
|
||||
|
||||
async fetchAdminControls(
|
||||
req: FetchAdminControlsRequest,
|
||||
): Promise<FetchAdminControlsResponse> {
|
||||
|
||||
@@ -51,6 +51,7 @@ export interface UserData {
|
||||
projectId: string;
|
||||
userTier: UserTierId;
|
||||
userTierName?: string;
|
||||
paidTier?: GeminiUserTier;
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -136,6 +137,7 @@ export async function setupUser(
|
||||
projectId,
|
||||
userTier: loadRes.paidTier?.id ?? loadRes.currentTier.id,
|
||||
userTierName: loadRes.paidTier?.name ?? loadRes.currentTier.name,
|
||||
paidTier: loadRes.paidTier ?? undefined,
|
||||
};
|
||||
}
|
||||
|
||||
@@ -146,6 +148,7 @@ export async function setupUser(
|
||||
projectId: loadRes.cloudaicompanionProject,
|
||||
userTier: loadRes.paidTier?.id ?? loadRes.currentTier.id,
|
||||
userTierName: loadRes.paidTier?.name ?? loadRes.currentTier.name,
|
||||
paidTier: loadRes.paidTier ?? undefined,
|
||||
};
|
||||
}
|
||||
|
||||
|
||||
@@ -39,11 +39,41 @@ export type ClientMetadataPluginType =
|
||||
| 'AIPLUGIN_INTELLIJ'
|
||||
| 'AIPLUGIN_STUDIO';
|
||||
|
||||
/**
|
||||
* Credit types that can be used for API consumption.
|
||||
*/
|
||||
export type CreditType = 'CREDIT_TYPE_UNSPECIFIED' | 'GOOGLE_ONE_AI';
|
||||
|
||||
/**
|
||||
* Represents a credit amount for a specific credit type.
|
||||
* Used in LoadCodeAssistResponse for available credits and
|
||||
* in GenerateContentResponse for consumed/remaining credits.
|
||||
*/
|
||||
export interface Credits {
|
||||
creditType: CreditType;
|
||||
creditAmount: string; // int64 represented as string in JSON
|
||||
}
|
||||
|
||||
/** Alias for Credits used in available_credits context */
|
||||
export type AvailableCredits = Credits;
|
||||
|
||||
/** Alias for Credits used in consumedCredits context */
|
||||
export type ConsumedCredits = Credits;
|
||||
|
||||
/** Alias for Credits used in remainingCredits context */
|
||||
export type RemainingCredits = Credits;
|
||||
|
||||
export interface LoadCodeAssistRequest {
|
||||
cloudaicompanionProject?: string;
|
||||
metadata: ClientMetadata;
|
||||
mode?: LoadCodeAssistMode;
|
||||
}
|
||||
|
||||
export type LoadCodeAssistMode =
|
||||
| 'MODE_UNSPECIFIED'
|
||||
| 'FULL_ELIGIBILITY_CHECK'
|
||||
| 'HEALTH_CHECK';
|
||||
|
||||
/**
|
||||
* Represents LoadCodeAssistResponse proto json field
|
||||
* http://google3/google/internal/cloud/code/v1internal/cloudcode.proto;l=224
|
||||
@@ -69,6 +99,8 @@ export interface GeminiUserTier {
|
||||
privacyNotice?: PrivacyNotice;
|
||||
hasAcceptedTos?: boolean;
|
||||
hasOnboardedPreviously?: boolean;
|
||||
/** Available AI credits for this tier (e.g., Google One AI credits) */
|
||||
availableCredits?: AvailableCredits[];
|
||||
}
|
||||
|
||||
/**
|
||||
|
||||
@@ -13,6 +13,7 @@ import type {
|
||||
ContentGenerator,
|
||||
ContentGeneratorConfig,
|
||||
} from '../core/contentGenerator.js';
|
||||
import type { OverageStrategy } from '../billing/billing.js';
|
||||
import {
|
||||
AuthType,
|
||||
createContentGenerator,
|
||||
@@ -110,6 +111,7 @@ import { ApprovalMode, type PolicyEngineConfig } from '../policy/types.js';
|
||||
import { HookSystem } from '../hooks/index.js';
|
||||
import type {
|
||||
UserTierId,
|
||||
GeminiUserTier,
|
||||
RetrieveUserQuotaResponse,
|
||||
AdminControlsSettings,
|
||||
} from '../code_assist/types.js';
|
||||
@@ -549,6 +551,9 @@ export interface ConfigParameters {
|
||||
agents?: AgentSettings;
|
||||
}>;
|
||||
enableConseca?: boolean;
|
||||
billing?: {
|
||||
overageStrategy?: OverageStrategy;
|
||||
};
|
||||
}
|
||||
|
||||
export class Config {
|
||||
@@ -725,6 +730,10 @@ export class Config {
|
||||
}>)
|
||||
| undefined;
|
||||
|
||||
private readonly billing: {
|
||||
overageStrategy: OverageStrategy;
|
||||
};
|
||||
|
||||
private readonly enableAgents: boolean;
|
||||
private agents: AgentSettings;
|
||||
private readonly enableEventDrivenScheduler: boolean;
|
||||
@@ -962,6 +971,10 @@ export class Config {
|
||||
this.onModelChange = params.onModelChange;
|
||||
this.onReload = params.onReload;
|
||||
|
||||
this.billing = {
|
||||
overageStrategy: params.billing?.overageStrategy ?? 'ask',
|
||||
};
|
||||
|
||||
if (params.contextFileName) {
|
||||
setGeminiMdFilename(params.contextFileName);
|
||||
}
|
||||
@@ -1224,6 +1237,10 @@ export class Config {
|
||||
return this.contentGenerator?.userTierName;
|
||||
}
|
||||
|
||||
getUserPaidTier(): GeminiUserTier | undefined {
|
||||
return this.contentGenerator?.paidTier;
|
||||
}
|
||||
|
||||
/**
|
||||
* Provides access to the BaseLlmClient for stateless LLM operations.
|
||||
*/
|
||||
@@ -1531,6 +1548,19 @@ export class Config {
|
||||
this.hasAccessToPreviewModel = hasAccess;
|
||||
}
|
||||
|
||||
async refreshAvailableCredits(): Promise<void> {
|
||||
const codeAssistServer = getCodeAssistServer(this);
|
||||
if (!codeAssistServer) {
|
||||
return;
|
||||
}
|
||||
try {
|
||||
await codeAssistServer.refreshAvailableCredits();
|
||||
} catch {
|
||||
// Non-fatal: proceed even if refresh fails.
|
||||
// The actual credit balance will be verified server-side.
|
||||
}
|
||||
}
|
||||
|
||||
async refreshUserQuota(): Promise<RetrieveUserQuotaResponse | undefined> {
|
||||
const codeAssistServer = getCodeAssistServer(this);
|
||||
if (!codeAssistServer || !codeAssistServer.projectId) {
|
||||
@@ -2005,6 +2035,19 @@ export class Config {
|
||||
return this.telemetrySettings.outfile;
|
||||
}
|
||||
|
||||
getBillingSettings(): { overageStrategy: OverageStrategy } {
|
||||
return this.billing;
|
||||
}
|
||||
|
||||
/**
|
||||
* Updates the overage strategy at runtime.
|
||||
* Used to switch from 'ask' to 'always' after the user accepts credits
|
||||
* via the overage dialog, so subsequent API calls auto-include credits.
|
||||
*/
|
||||
setOverageStrategy(strategy: OverageStrategy): void {
|
||||
this.billing.overageStrategy = strategy;
|
||||
}
|
||||
|
||||
getTelemetryUseCollector(): boolean {
|
||||
return this.telemetrySettings.useCollector ?? false;
|
||||
}
|
||||
|
||||
@@ -17,7 +17,7 @@ import { createCodeAssistContentGenerator } from '../code_assist/codeAssist.js';
|
||||
import type { Config } from '../config/config.js';
|
||||
import { loadApiKey } from './apiKeyCredentialStorage.js';
|
||||
|
||||
import type { UserTierId } from '../code_assist/types.js';
|
||||
import type { UserTierId, GeminiUserTier } from '../code_assist/types.js';
|
||||
import { LoggingContentGenerator } from './loggingContentGenerator.js';
|
||||
import { InstallationManager } from '../utils/installationManager.js';
|
||||
import { FakeContentGenerator } from './fakeContentGenerator.js';
|
||||
@@ -49,6 +49,8 @@ export interface ContentGenerator {
|
||||
userTier?: UserTierId;
|
||||
|
||||
userTierName?: string;
|
||||
|
||||
paidTier?: GeminiUserTier;
|
||||
}
|
||||
|
||||
export enum AuthType {
|
||||
|
||||
@@ -14,7 +14,7 @@ import {
|
||||
} from '@google/genai';
|
||||
import { promises } from 'node:fs';
|
||||
import type { ContentGenerator } from './contentGenerator.js';
|
||||
import type { UserTierId } from '../code_assist/types.js';
|
||||
import type { UserTierId, GeminiUserTier } from '../code_assist/types.js';
|
||||
import { safeJsonStringify } from '../utils/safeJsonStringify.js';
|
||||
import type { LlmRole } from '../telemetry/types.js';
|
||||
|
||||
@@ -44,6 +44,7 @@ export class FakeContentGenerator implements ContentGenerator {
|
||||
private callCounter = 0;
|
||||
userTier?: UserTierId;
|
||||
userTierName?: string;
|
||||
paidTier?: GeminiUserTier;
|
||||
|
||||
constructor(private readonly responses: FakeResponse[]) {}
|
||||
|
||||
|
||||
@@ -24,7 +24,7 @@ import {
|
||||
} from '../telemetry/types.js';
|
||||
import type { LlmRole } from '../telemetry/llmRole.js';
|
||||
import type { Config } from '../config/config.js';
|
||||
import type { UserTierId } from '../code_assist/types.js';
|
||||
import type { UserTierId, GeminiUserTier } from '../code_assist/types.js';
|
||||
import {
|
||||
logApiError,
|
||||
logApiRequest,
|
||||
@@ -153,6 +153,10 @@ export class LoggingContentGenerator implements ContentGenerator {
|
||||
return this.wrapped.userTierName;
|
||||
}
|
||||
|
||||
get paidTier(): GeminiUserTier | undefined {
|
||||
return this.wrapped.paidTier;
|
||||
}
|
||||
|
||||
private logApiRequest(
|
||||
contents: Content[],
|
||||
model: string,
|
||||
|
||||
@@ -140,6 +140,9 @@ async function processIntent(
|
||||
// based on the availability service state (which is updated before this).
|
||||
return true;
|
||||
|
||||
case 'retry_with_credits':
|
||||
return true;
|
||||
|
||||
case 'stop':
|
||||
// Do not switch model on stop. User wants to stay on current model (and stop).
|
||||
return false;
|
||||
|
||||
@@ -17,6 +17,7 @@ import type {
|
||||
export type FallbackIntent =
|
||||
| 'retry_always' // Retry with fallback model and stick to it for future requests.
|
||||
| 'retry_once' // Retry with fallback model for this request only.
|
||||
| 'retry_with_credits' // Retry the current request using Google One AI credits (and potentially future ones if strategy is 'always').
|
||||
| 'stop' // Switch to fallback for future requests, but stop the current request.
|
||||
| 'retry_later' // Stop the current request and do not fallback. Intend to try again later with the same model.
|
||||
| 'upgrade'; // Give user an option to upgrade the tier.
|
||||
|
||||
@@ -18,6 +18,7 @@ export * from './policy/policy-engine.js';
|
||||
export * from './policy/toml-loader.js';
|
||||
export * from './policy/config.js';
|
||||
export * from './policy/integrity.js';
|
||||
export * from './billing/index.js';
|
||||
export * from './confirmation-bus/types.js';
|
||||
export * from './confirmation-bus/message-bus.js';
|
||||
|
||||
@@ -76,6 +77,7 @@ export * from './utils/quotaErrorDetection.js';
|
||||
export * from './utils/userAccountManager.js';
|
||||
export * from './utils/authConsent.js';
|
||||
export * from './utils/googleQuotaErrors.js';
|
||||
export * from './utils/googleErrors.js';
|
||||
export * from './utils/fileUtils.js';
|
||||
export * from './utils/planUtils.js';
|
||||
export * from './utils/approvalModeUtils.js';
|
||||
@@ -98,6 +100,7 @@ export * from './utils/ignorePatterns.js';
|
||||
export * from './utils/partUtils.js';
|
||||
export * from './utils/promptIdContext.js';
|
||||
export * from './utils/thoughtUtils.js';
|
||||
export * from './utils/secure-browser-launcher.js';
|
||||
export * from './utils/debugLogger.js';
|
||||
export * from './utils/events.js';
|
||||
export * from './utils/extensionLoader.js';
|
||||
@@ -183,6 +186,8 @@ export { OAuthUtils } from './mcp/oauth-utils.js';
|
||||
|
||||
// Export telemetry functions
|
||||
export * from './telemetry/index.js';
|
||||
export * from './telemetry/billingEvents.js';
|
||||
export { logBillingEvent } from './telemetry/loggers.js';
|
||||
export { sessionId, createSessionId } from './utils/session.js';
|
||||
export * from './utils/compatibility.js';
|
||||
export * from './utils/browser.js';
|
||||
|
||||
@@ -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),
|
||||
isInteractive: vi.fn().mockReturnValue(true),
|
||||
getExperiments: vi.fn().mockReturnValue({ experimentIds: [] }),
|
||||
getContentGeneratorConfig: vi.fn().mockReturnValue({ authType: 'oauth' }),
|
||||
} as unknown as Config;
|
||||
|
||||
mockLogger = {
|
||||
|
||||
@@ -77,6 +77,7 @@ export type { TelemetryEvent } from './types.js';
|
||||
export { SpanStatusCode, ValueType } from '@opentelemetry/api';
|
||||
export { SemanticAttributes } from '@opentelemetry/semantic-conventions';
|
||||
export * from './uiTelemetry.js';
|
||||
export * from './billingEvents.js';
|
||||
export {
|
||||
MemoryMonitor,
|
||||
initializeMemoryMonitor,
|
||||
@@ -145,6 +146,9 @@ export {
|
||||
GenAiOperationName,
|
||||
GenAiProviderName,
|
||||
GenAiTokenType,
|
||||
// Billing metrics functions
|
||||
recordOverageOptionSelected,
|
||||
recordCreditPurchaseClick,
|
||||
} from './metrics.js';
|
||||
export { runInDevTraceSpan, type SpanMetadata } from './trace.js';
|
||||
export { startupProfiler, StartupProfiler } from './startupProfiler.js';
|
||||
|
||||
@@ -280,6 +280,7 @@ describe('loggers', () => {
|
||||
isInteractive: () => false,
|
||||
getExperiments: () => undefined,
|
||||
getExperimentsAsync: async () => undefined,
|
||||
getContentGeneratorConfig: () => undefined,
|
||||
} as unknown as Config;
|
||||
|
||||
it('should log a user prompt', () => {
|
||||
@@ -319,6 +320,7 @@ describe('loggers', () => {
|
||||
isInteractive: () => false,
|
||||
getExperiments: () => undefined,
|
||||
getExperimentsAsync: async () => undefined,
|
||||
getContentGeneratorConfig: () => undefined,
|
||||
} as unknown as Config;
|
||||
const event = new UserPromptEvent(
|
||||
11,
|
||||
@@ -356,7 +358,8 @@ describe('loggers', () => {
|
||||
isInteractive: () => false,
|
||||
getExperiments: () => undefined,
|
||||
getExperimentsAsync: async () => undefined,
|
||||
} as Config;
|
||||
getContentGeneratorConfig: () => undefined,
|
||||
} as unknown as Config;
|
||||
|
||||
const mockMetrics = {
|
||||
recordApiResponseMetrics: vi.fn(),
|
||||
@@ -558,7 +561,8 @@ describe('loggers', () => {
|
||||
isInteractive: () => false,
|
||||
getExperiments: () => undefined,
|
||||
getExperimentsAsync: async () => undefined,
|
||||
} as Config;
|
||||
getContentGeneratorConfig: () => undefined,
|
||||
} as unknown as Config;
|
||||
|
||||
const mockMetrics = {
|
||||
recordApiResponseMetrics: vi.fn(),
|
||||
@@ -996,6 +1000,7 @@ describe('loggers', () => {
|
||||
isInteractive: () => false,
|
||||
getExperiments: () => undefined,
|
||||
getExperimentsAsync: async () => undefined,
|
||||
getContentGeneratorConfig: () => undefined,
|
||||
} as unknown as Config;
|
||||
|
||||
it('should log flash fallback event', () => {
|
||||
@@ -1025,6 +1030,7 @@ describe('loggers', () => {
|
||||
isInteractive: () => false,
|
||||
getExperiments: () => undefined,
|
||||
getExperimentsAsync: async () => undefined,
|
||||
getContentGeneratorConfig: () => undefined,
|
||||
} as unknown as Config;
|
||||
|
||||
beforeEach(() => {
|
||||
@@ -1121,7 +1127,8 @@ describe('loggers', () => {
|
||||
isInteractive: () => false,
|
||||
getExperiments: () => undefined,
|
||||
getExperimentsAsync: async () => undefined,
|
||||
} as Config;
|
||||
getContentGeneratorConfig: () => undefined,
|
||||
} as unknown as Config;
|
||||
|
||||
const mockMetrics = {
|
||||
recordToolCallMetrics: vi.fn(),
|
||||
@@ -1741,7 +1748,8 @@ describe('loggers', () => {
|
||||
isInteractive: () => false,
|
||||
getExperiments: () => undefined,
|
||||
getExperimentsAsync: async () => undefined,
|
||||
} as Config;
|
||||
getContentGeneratorConfig: () => undefined,
|
||||
} as unknown as Config;
|
||||
|
||||
const mockMetrics = {
|
||||
recordFileOperationMetric: vi.fn(),
|
||||
@@ -1803,6 +1811,7 @@ describe('loggers', () => {
|
||||
isInteractive: () => false,
|
||||
getExperiments: () => undefined,
|
||||
getExperimentsAsync: async () => undefined,
|
||||
getContentGeneratorConfig: () => undefined,
|
||||
} as unknown as Config;
|
||||
|
||||
it('should log a tool output truncated event', () => {
|
||||
@@ -1842,6 +1851,7 @@ describe('loggers', () => {
|
||||
isInteractive: () => false,
|
||||
getExperiments: () => undefined,
|
||||
getExperimentsAsync: async () => undefined,
|
||||
getContentGeneratorConfig: () => undefined,
|
||||
} as unknown as Config;
|
||||
|
||||
beforeEach(() => {
|
||||
@@ -2099,6 +2109,7 @@ describe('loggers', () => {
|
||||
isInteractive: () => false,
|
||||
getExperiments: () => undefined,
|
||||
getExperimentsAsync: async () => undefined,
|
||||
getContentGeneratorConfig: () => undefined,
|
||||
} as unknown as Config;
|
||||
|
||||
beforeEach(() => {
|
||||
@@ -2146,6 +2157,7 @@ describe('loggers', () => {
|
||||
isInteractive: () => false,
|
||||
getExperiments: () => undefined,
|
||||
getExperimentsAsync: async () => undefined,
|
||||
getContentGeneratorConfig: () => undefined,
|
||||
} as unknown as Config;
|
||||
|
||||
beforeEach(() => {
|
||||
@@ -2193,6 +2205,7 @@ describe('loggers', () => {
|
||||
isInteractive: () => false,
|
||||
getExperiments: () => undefined,
|
||||
getExperimentsAsync: async () => undefined,
|
||||
getContentGeneratorConfig: () => undefined,
|
||||
} as unknown as Config;
|
||||
|
||||
beforeEach(() => {
|
||||
@@ -2231,6 +2244,7 @@ describe('loggers', () => {
|
||||
isInteractive: () => false,
|
||||
getExperiments: () => undefined,
|
||||
getExperimentsAsync: async () => undefined,
|
||||
getContentGeneratorConfig: () => undefined,
|
||||
} as unknown as Config;
|
||||
|
||||
beforeEach(() => {
|
||||
@@ -2284,6 +2298,7 @@ describe('loggers', () => {
|
||||
isInteractive: () => false,
|
||||
getExperiments: () => undefined,
|
||||
getExperimentsAsync: async () => undefined,
|
||||
getContentGeneratorConfig: () => undefined,
|
||||
} as unknown as Config;
|
||||
|
||||
beforeEach(() => {
|
||||
@@ -2322,6 +2337,7 @@ describe('loggers', () => {
|
||||
getExperiments: () => undefined,
|
||||
getExperimentsAsync: async () => undefined,
|
||||
getTelemetryLogPromptsEnabled: () => false,
|
||||
getContentGeneratorConfig: () => undefined,
|
||||
} as unknown as Config;
|
||||
|
||||
beforeEach(() => {
|
||||
|
||||
@@ -84,6 +84,7 @@ import type { UiEvent } from './uiTelemetry.js';
|
||||
import { uiTelemetryService } from './uiTelemetry.js';
|
||||
import { ClearcutLogger } from './clearcut-logger/clearcut-logger.js';
|
||||
import { debugLogger } from '../utils/debugLogger.js';
|
||||
import type { BillingTelemetryEvent } from './billingEvents.js';
|
||||
|
||||
export function logCliConfiguration(
|
||||
config: Config,
|
||||
@@ -827,3 +828,17 @@ export function logTokenStorageInitialization(
|
||||
recordTokenStorageInitialization(config, event);
|
||||
});
|
||||
}
|
||||
|
||||
export function logBillingEvent(
|
||||
config: Config,
|
||||
event: BillingTelemetryEvent,
|
||||
): void {
|
||||
bufferTelemetryEvent(() => {
|
||||
const logger = logs.getLogger(SERVICE_NAME);
|
||||
const logRecord: LogRecord = {
|
||||
body: event.toLogBody(),
|
||||
attributes: event.toOpenTelemetryAttributes(config),
|
||||
};
|
||||
logger.emit(logRecord);
|
||||
});
|
||||
}
|
||||
|
||||
@@ -41,6 +41,8 @@ const EVENT_HOOK_CALL_COUNT = 'gemini_cli.hook_call.count';
|
||||
const EVENT_HOOK_CALL_LATENCY = 'gemini_cli.hook_call.latency';
|
||||
const KEYCHAIN_AVAILABILITY_COUNT = 'gemini_cli.keychain.availability.count';
|
||||
const TOKEN_STORAGE_TYPE_COUNT = 'gemini_cli.token_storage.type.count';
|
||||
const OVERAGE_OPTION_COUNT = 'gemini_cli.overage_option.count';
|
||||
const CREDIT_PURCHASE_COUNT = 'gemini_cli.credit_purchase.count';
|
||||
|
||||
// Agent Metrics
|
||||
const AGENT_RUN_COUNT = 'gemini_cli.agent.run.count';
|
||||
@@ -259,6 +261,26 @@ const COUNTER_DEFINITIONS = {
|
||||
forced: boolean;
|
||||
},
|
||||
},
|
||||
[OVERAGE_OPTION_COUNT]: {
|
||||
description: 'Counts overage option selections.',
|
||||
valueType: ValueType.INT,
|
||||
assign: (c: Counter) => (overageOptionCounter = c),
|
||||
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
|
||||
attributes: {} as {
|
||||
selected_option: string;
|
||||
model: string;
|
||||
},
|
||||
},
|
||||
[CREDIT_PURCHASE_COUNT]: {
|
||||
description: 'Counts credit purchase link clicks.',
|
||||
valueType: ValueType.INT,
|
||||
assign: (c: Counter) => (creditPurchaseCounter = c),
|
||||
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
|
||||
attributes: {} as {
|
||||
source: string;
|
||||
model: string;
|
||||
},
|
||||
},
|
||||
} as const;
|
||||
|
||||
const HISTOGRAM_DEFINITIONS = {
|
||||
@@ -597,6 +619,8 @@ let hookCallCounter: Counter | undefined;
|
||||
let hookCallLatencyHistogram: Histogram | undefined;
|
||||
let keychainAvailabilityCounter: Counter | undefined;
|
||||
let tokenStorageTypeCounter: Counter | undefined;
|
||||
let overageOptionCounter: Counter | undefined;
|
||||
let creditPurchaseCounter: Counter | undefined;
|
||||
|
||||
// OpenTelemetry GenAI Semantic Convention Metrics
|
||||
let genAiClientTokenUsageHistogram: Histogram | undefined;
|
||||
@@ -1334,3 +1358,31 @@ export function recordTokenStorageInitialization(
|
||||
forced: event.forced,
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Records a metric for an overage option selection.
|
||||
*/
|
||||
export function recordOverageOptionSelected(
|
||||
config: Config,
|
||||
attributes: MetricDefinitions[typeof OVERAGE_OPTION_COUNT]['attributes'],
|
||||
): void {
|
||||
if (!overageOptionCounter || !isMetricsInitialized) return;
|
||||
overageOptionCounter.add(1, {
|
||||
...baseMetricDefinition.getCommonAttributes(config),
|
||||
...attributes,
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Records a metric for a credit purchase link click.
|
||||
*/
|
||||
export function recordCreditPurchaseClick(
|
||||
config: Config,
|
||||
attributes: MetricDefinitions[typeof CREDIT_PURCHASE_COUNT]['attributes'],
|
||||
): void {
|
||||
if (!creditPurchaseCounter || !isMetricsInitialized) return;
|
||||
creditPurchaseCounter.add(1, {
|
||||
...baseMetricDefinition.getCommonAttributes(config),
|
||||
...attributes,
|
||||
});
|
||||
}
|
||||
|
||||
@@ -32,6 +32,7 @@ function createMockConfig(logPromptsEnabled: boolean): Config {
|
||||
getModel: () => 'gemini-1.5-flash',
|
||||
isInteractive: () => true,
|
||||
getUserEmail: () => undefined,
|
||||
getContentGeneratorConfig: () => undefined,
|
||||
} as unknown as Config;
|
||||
}
|
||||
|
||||
|
||||
@@ -77,6 +77,7 @@ describe('Telemetry SDK', () => {
|
||||
isInteractive: () => false,
|
||||
getExperiments: () => undefined,
|
||||
getExperimentsAsync: async () => undefined,
|
||||
getContentGeneratorConfig: () => undefined,
|
||||
} as unknown as Config;
|
||||
});
|
||||
|
||||
|
||||
@@ -15,11 +15,13 @@ const installationManager = new InstallationManager();
|
||||
export function getCommonAttributes(config: Config): Attributes {
|
||||
const email = userAccountManager.getCachedGoogleAccount();
|
||||
const experiments = config.getExperiments();
|
||||
const authType = config.getContentGeneratorConfig()?.authType;
|
||||
return {
|
||||
'session.id': config.getSessionId(),
|
||||
'installation.id': installationManager.getInstallationId(),
|
||||
interactive: config.isInteractive(),
|
||||
...(email && { 'user.email': email }),
|
||||
...(authType && { auth_type: authType }),
|
||||
...(experiments &&
|
||||
experiments.experimentIds.length > 0 && {
|
||||
'experiments.ids': experiments.experimentIds,
|
||||
|
||||
@@ -16,8 +16,8 @@
|
||||
export interface ErrorInfo {
|
||||
'@type': 'type.googleapis.com/google.rpc.ErrorInfo';
|
||||
reason: string;
|
||||
domain: string;
|
||||
metadata: { [key: string]: string };
|
||||
domain?: string;
|
||||
metadata?: { [key: string]: string };
|
||||
}
|
||||
|
||||
export interface RetryInfo {
|
||||
|
||||
@@ -297,6 +297,26 @@ describe('classifyGoogleError', () => {
|
||||
expect(result).toBeInstanceOf(TerminalQuotaError);
|
||||
});
|
||||
|
||||
it('should return TerminalQuotaError for INSUFFICIENT_G1_CREDITS_BALANCE without domain', () => {
|
||||
const apiError: GoogleApiError = {
|
||||
code: 429,
|
||||
message: 'Resource has been exhausted (e.g. check quota).',
|
||||
details: [
|
||||
{
|
||||
'@type': 'type.googleapis.com/google.rpc.ErrorInfo',
|
||||
reason: 'INSUFFICIENT_G1_CREDITS_BALANCE',
|
||||
},
|
||||
],
|
||||
};
|
||||
vi.spyOn(errorParser, 'parseGoogleApiError').mockReturnValue(apiError);
|
||||
const result = classifyGoogleError(new Error());
|
||||
expect(result).toBeInstanceOf(TerminalQuotaError);
|
||||
expect((result as TerminalQuotaError).isInsufficientCredits).toBe(true);
|
||||
expect((result as TerminalQuotaError).reason).toBe(
|
||||
'INSUFFICIENT_G1_CREDITS_BALANCE',
|
||||
);
|
||||
});
|
||||
|
||||
it('should prioritize daily limit over retry info', () => {
|
||||
const apiError: GoogleApiError = {
|
||||
code: 429,
|
||||
|
||||
@@ -19,17 +19,24 @@ import { getErrorStatus, ModelNotFoundError } from './httpErrors.js';
|
||||
*/
|
||||
export class TerminalQuotaError extends Error {
|
||||
retryDelayMs?: number;
|
||||
reason?: string;
|
||||
|
||||
constructor(
|
||||
message: string,
|
||||
override readonly cause: GoogleApiError,
|
||||
retryDelaySeconds?: number,
|
||||
reason?: string,
|
||||
) {
|
||||
super(message);
|
||||
this.name = 'TerminalQuotaError';
|
||||
this.retryDelayMs = retryDelaySeconds
|
||||
? retryDelaySeconds * 1000
|
||||
: undefined;
|
||||
this.reason = reason;
|
||||
}
|
||||
|
||||
get isInsufficientCredits(): boolean {
|
||||
return this.reason === 'INSUFFICIENT_G1_CREDITS_BALANCE';
|
||||
}
|
||||
}
|
||||
|
||||
@@ -121,6 +128,7 @@ function classifyValidationRequiredError(
|
||||
}
|
||||
|
||||
if (
|
||||
!errorInfo.domain ||
|
||||
!CLOUDCODE_DOMAINS.includes(errorInfo.domain) ||
|
||||
errorInfo.reason !== 'VALIDATION_REQUIRED'
|
||||
) {
|
||||
@@ -293,6 +301,16 @@ export function classifyGoogleError(error: unknown): unknown {
|
||||
}
|
||||
|
||||
if (errorInfo) {
|
||||
// INSUFFICIENT_G1_CREDITS_BALANCE is always terminal, regardless of domain
|
||||
if (errorInfo.reason === 'INSUFFICIENT_G1_CREDITS_BALANCE') {
|
||||
return new TerminalQuotaError(
|
||||
`${googleApiError.message}`,
|
||||
googleApiError,
|
||||
delaySeconds,
|
||||
errorInfo.reason,
|
||||
);
|
||||
}
|
||||
|
||||
// New Cloud Code API quota handling
|
||||
if (errorInfo.domain) {
|
||||
const validDomains = [
|
||||
@@ -313,6 +331,7 @@ export function classifyGoogleError(error: unknown): unknown {
|
||||
`${googleApiError.message}`,
|
||||
googleApiError,
|
||||
delaySeconds,
|
||||
errorInfo.reason,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user