feat(cliproxy): enrich gemini quota metadata

This commit is contained in:
Tam Nhu Tran
2026-04-03 19:43:17 -04:00
parent d78f35b582
commit a762d8dda0
5 changed files with 511 additions and 118 deletions
+169
View File
@@ -0,0 +1,169 @@
import type { GeminiCliBucket } from './quota-types';
export interface GeminiCliParsedBucket {
modelId: string;
tokenType: string | null;
remainingFraction: number | null;
remainingAmount: number | null;
resetTime: string | null;
}
interface GeminiCliQuotaGroupDefinition {
id: string;
label: string;
preferredModelId?: string;
modelIds: string[];
}
const GEMINI_CLI_QUOTA_GROUPS: GeminiCliQuotaGroupDefinition[] = [
{
id: 'gemini-flash-lite-series',
label: 'Gemini Flash Lite Series',
preferredModelId: 'gemini-2.5-flash-lite',
modelIds: ['gemini-2.5-flash-lite'],
},
{
id: 'gemini-flash-series',
label: 'Gemini Flash Series',
preferredModelId: 'gemini-3-flash-preview',
modelIds: ['gemini-3-flash-preview', 'gemini-3.1-flash-preview', 'gemini-2.5-flash'],
},
{
id: 'gemini-pro-series',
label: 'Gemini Pro Series',
preferredModelId: 'gemini-3.1-pro-preview',
modelIds: ['gemini-3.1-pro-preview', 'gemini-3-pro-preview', 'gemini-2.5-pro'],
},
];
const GEMINI_CLI_GROUP_ORDER = new Map(
GEMINI_CLI_QUOTA_GROUPS.map((group, index) => [group.id, index] as const)
);
const GEMINI_CLI_GROUP_LOOKUP = new Map(
GEMINI_CLI_QUOTA_GROUPS.flatMap((group) =>
group.modelIds.map((modelId) => [modelId, group] as const)
)
);
const GEMINI_CLI_IGNORED_MODEL_PREFIXES = ['gemini-2.0-flash'];
type GeminiCliQuotaBucketGroup = {
id: string;
label: string;
tokenType: string | null;
modelIds: string[];
preferredModelId?: string;
preferredBucket?: GeminiCliParsedBucket;
fallbackRemainingFraction: number | null;
fallbackRemainingAmount: number | null;
fallbackResetTime: string | null;
};
function isIgnoredGeminiCliModel(modelId: string): boolean {
return GEMINI_CLI_IGNORED_MODEL_PREFIXES.some((prefix) => modelId.startsWith(prefix));
}
function pickEarlierResetTime(current: string | null, next: string | null): string | null {
if (!current) return next;
if (!next) return current;
const currentTime = new Date(current).getTime();
const nextTime = new Date(next).getTime();
if (Number.isNaN(currentTime)) return next;
if (Number.isNaN(nextTime)) return current;
return currentTime <= nextTime ? current : next;
}
function minNullableNumber(current: number | null, next: number | null): number | null {
if (current === null) return next;
if (next === null) return current;
return Math.min(current, next);
}
function clampQuotaFraction(value: number | null): number {
if (value === null) return 0;
return Math.max(0, Math.min(1, value));
}
function getGroupOrder(bucket: GeminiCliQuotaBucketGroup): number {
return GEMINI_CLI_GROUP_ORDER.get(bucket.id) ?? Number.MAX_SAFE_INTEGER;
}
export function buildGeminiCliBucketsFromParsedBuckets(
buckets: GeminiCliParsedBucket[]
): GeminiCliBucket[] {
if (buckets.length === 0) return [];
const grouped = new Map<string, GeminiCliQuotaBucketGroup>();
for (const bucket of buckets) {
if (isIgnoredGeminiCliModel(bucket.modelId)) continue;
const group = GEMINI_CLI_GROUP_LOOKUP.get(bucket.modelId);
const groupId = group?.id ?? bucket.modelId;
const label = group?.label ?? bucket.modelId;
const tokenKey = bucket.tokenType ?? '';
const cacheKey = `${groupId}::${tokenKey || 'combined'}`;
const existing = grouped.get(cacheKey);
if (!existing) {
const preferredModelId = group?.preferredModelId;
grouped.set(cacheKey, {
id: groupId,
label,
tokenType: bucket.tokenType,
modelIds: [bucket.modelId],
preferredModelId,
preferredBucket:
preferredModelId && bucket.modelId === preferredModelId ? bucket : undefined,
fallbackRemainingFraction: bucket.remainingFraction,
fallbackRemainingAmount: bucket.remainingAmount,
fallbackResetTime: bucket.resetTime,
});
continue;
}
existing.modelIds.push(bucket.modelId);
existing.fallbackRemainingFraction = minNullableNumber(
existing.fallbackRemainingFraction,
bucket.remainingFraction
);
existing.fallbackRemainingAmount = minNullableNumber(
existing.fallbackRemainingAmount,
bucket.remainingAmount
);
existing.fallbackResetTime = pickEarlierResetTime(existing.fallbackResetTime, bucket.resetTime);
if (existing.preferredModelId && bucket.modelId === existing.preferredModelId) {
existing.preferredBucket = bucket;
}
}
return Array.from(grouped.values())
.sort((a, b) => {
const orderDiff = getGroupOrder(a) - getGroupOrder(b);
if (orderDiff !== 0) return orderDiff;
const tokenDiff = (a.tokenType ?? '').localeCompare(b.tokenType ?? '');
if (tokenDiff !== 0) return tokenDiff;
return a.label.localeCompare(b.label);
})
.map((bucket) => {
const preferred = bucket.preferredBucket;
const remainingFraction = clampQuotaFraction(
preferred?.remainingFraction ?? bucket.fallbackRemainingFraction
);
const remainingAmount = preferred?.remainingAmount ?? bucket.fallbackRemainingAmount;
const resetTime = preferred?.resetTime ?? bucket.fallbackResetTime;
return {
id: `${bucket.id}::${bucket.tokenType || 'combined'}`,
label: bucket.label,
tokenType: bucket.tokenType,
remainingFraction,
remainingPercent: Math.round(remainingFraction * 100),
remainingAmount,
resetTime,
modelIds: Array.from(new Set(bucket.modelIds)),
};
});
}
+178 -111
View File
@@ -11,43 +11,28 @@ import { getAuthDir } from './config-generator';
import { getProviderAccounts, getPausedDir } from './account-manager'; import { getProviderAccounts, getPausedDir } from './account-manager';
import { sanitizeEmail, isTokenExpired } from './auth-utils'; import { sanitizeEmail, isTokenExpired } from './auth-utils';
import { refreshGeminiToken } from './auth/gemini-token-refresh'; import { refreshGeminiToken } from './auth/gemini-token-refresh';
import {
buildGeminiCliBucketsFromParsedBuckets,
type GeminiCliParsedBucket,
} from './gemini-cli-quota-normalizer';
import type { GeminiCliQuotaResult, GeminiCliBucket } from './quota-types'; import type { GeminiCliQuotaResult, GeminiCliBucket } from './quota-types';
/** Google Cloud Code API endpoints */ /** Google Cloud Code API endpoints */
const GEMINI_CLI_API_BASE = 'https://cloudcode-pa.googleapis.com'; const GEMINI_CLI_API_BASE = 'https://cloudcode-pa.googleapis.com';
const GEMINI_CLI_API_VERSION = 'v1internal'; const GEMINI_CLI_API_VERSION = 'v1internal';
const GEMINI_CLI_QUOTA_URL = `${GEMINI_CLI_API_BASE}/${GEMINI_CLI_API_VERSION}:retrieveUserQuota`;
const GEMINI_CLI_CODE_ASSIST_URL = `${GEMINI_CLI_API_BASE}/${GEMINI_CLI_API_VERSION}:loadCodeAssist`;
const GEMINI_CLI_ERROR_DETAIL_MAX_LENGTH = 320; const GEMINI_CLI_ERROR_DETAIL_MAX_LENGTH = 320;
const GEMINI_CLI_ERROR_DETAIL_TRUNCATION_SUFFIX = '...[truncated]'; const GEMINI_CLI_ERROR_DETAIL_TRUNCATION_SUFFIX = '...[truncated]';
const GEMINI_CLI_G1_CREDIT_TYPE = 'GOOGLE_ONE_AI';
/** const GEMINI_CLI_TIER_LABELS: Record<string, string> = {
* Model groups for quota consolidation. 'free-tier': 'Free',
* Update when Google releases new Gemini models to include them in quota display. 'legacy-tier': 'Legacy',
*/ 'standard-tier': 'Standard',
const GEMINI_CLI_GROUPS: Record< 'g1-pro-tier': 'Pro',
string, 'g1-ultra-tier': 'Ultra',
{
label: string;
models: string[];
}
> = {
'gemini-flash-series': {
label: 'Gemini Flash Series',
models: [
'gemini-3-flash-preview',
'gemini-3.1-flash-preview',
'gemini-2.5-flash',
'gemini-2.5-flash-lite',
],
},
'gemini-pro-series': {
label: 'Gemini Pro Series',
models: ['gemini-3-pro-preview', 'gemini-3.1-pro-preview', 'gemini-2.5-pro'],
},
}; };
/** Models to ignore in quota display (deprecated) */
const IGNORED_MODEL_PREFIXES = ['gemini-2.0-flash'];
/** Auth data extracted from Gemini CLI auth file */ /** Auth data extracted from Gemini CLI auth file */
interface GeminiCliAuthData { interface GeminiCliAuthData {
accessToken: string; accessToken: string;
@@ -75,12 +60,38 @@ interface GeminiCliQuotaResponse {
buckets?: RawGeminiCliBucket[]; buckets?: RawGeminiCliBucket[];
} }
interface GeminiCliCredits {
creditType?: string;
credit_type?: string;
creditAmount?: string | number;
credit_amount?: string | number;
}
interface GeminiCliUserTier {
id?: string;
availableCredits?: GeminiCliCredits[];
available_credits?: GeminiCliCredits[];
}
interface GeminiCliCodeAssistResponse {
currentTier?: GeminiCliUserTier | null;
current_tier?: GeminiCliUserTier | null;
paidTier?: GeminiCliUserTier | null;
paid_tier?: GeminiCliUserTier | null;
}
interface ParsedGeminiCliErrorBody { interface ParsedGeminiCliErrorBody {
errorCode?: string; errorCode?: string;
errorDetail?: string; errorDetail?: string;
message?: string; message?: string;
} }
interface GeminiCliSupplementaryInfo {
tierLabel: string | null;
tierId: string | null;
creditBalance: number | null;
}
/** /**
* Extract project ID from account field * Extract project ID from account field
* Input: "user@example.com (cloudaicompanion-abc-123)" * Input: "user@example.com (cloudaicompanion-abc-123)"
@@ -234,23 +245,112 @@ function readGeminiCliAuthData(accountId: string): GeminiCliAuthData | null {
return null; return null;
} }
/** function normalizeStringValue(value: unknown): string | null {
* Find which group a model belongs to return typeof value === 'string' && value.trim().length > 0 ? value.trim() : null;
*/ }
function findModelGroup(modelId: string): { groupId: string; label: string } | null {
for (const [groupId, group] of Object.entries(GEMINI_CLI_GROUPS)) { function normalizeNumberValue(value: unknown): number | null {
if (group.models.includes(modelId)) { if (typeof value === 'number' && Number.isFinite(value)) {
return { groupId, label: group.label }; return value;
}
if (typeof value === 'string' && value.trim().length > 0) {
const parsed = Number(value);
if (Number.isFinite(parsed)) {
return parsed;
} }
} }
return null; return null;
} }
/** function resolveGeminiCliTierId(payload: GeminiCliCodeAssistResponse | null): string | null {
* Check if model should be ignored if (!payload) return null;
*/ const currentTier = payload.currentTier ?? payload.current_tier;
function shouldIgnoreModel(modelId: string): boolean { const paidTier = payload.paidTier ?? payload.paid_tier;
return IGNORED_MODEL_PREFIXES.some((prefix) => modelId.startsWith(prefix)); const rawId = normalizeStringValue(paidTier?.id) ?? normalizeStringValue(currentTier?.id);
return rawId ? rawId.toLowerCase() : null;
}
function resolveGeminiCliTierLabel(payload: GeminiCliCodeAssistResponse | null): string | null {
const tierId = resolveGeminiCliTierId(payload);
if (!tierId) return null;
return GEMINI_CLI_TIER_LABELS[tierId] ?? tierId;
}
function resolveGeminiCliCreditBalance(payload: GeminiCliCodeAssistResponse | null): number | null {
if (!payload) return null;
const paidTier = payload.paidTier ?? payload.paid_tier;
const currentTier = payload.currentTier ?? payload.current_tier;
const tier = paidTier ?? currentTier;
if (!tier) return null;
const credits = tier.availableCredits ?? tier.available_credits ?? [];
let total = 0;
let found = false;
for (const credit of credits) {
const creditType = normalizeStringValue(credit.creditType ?? credit.credit_type);
if (creditType !== GEMINI_CLI_G1_CREDIT_TYPE) continue;
const amount = normalizeNumberValue(credit.creditAmount ?? credit.credit_amount);
if (amount !== null) {
total += amount;
found = true;
}
}
return found ? total : null;
}
async function fetchGeminiCliSupplementary(
accessToken: string,
projectId: string,
verbose: boolean
): Promise<GeminiCliSupplementaryInfo> {
const controller = new AbortController();
const timeoutId = setTimeout(() => controller.abort(), 5000);
try {
const response = await fetch(GEMINI_CLI_CODE_ASSIST_URL, {
method: 'POST',
signal: controller.signal,
headers: {
Authorization: `Bearer ${accessToken}`,
'Content-Type': 'application/json',
},
body: JSON.stringify({
cloudaicompanionProject: projectId,
metadata: {
ideType: 'IDE_UNSPECIFIED',
platform: 'PLATFORM_UNSPECIFIED',
pluginType: 'GEMINI',
duetProject: projectId,
},
}),
});
clearTimeout(timeoutId);
if (!response.ok) {
if (verbose) {
console.error(`[i] Gemini CLI supplementary metadata unavailable: HTTP ${response.status}`);
}
return { tierLabel: null, tierId: null, creditBalance: null };
}
const payload = (await response.json()) as GeminiCliCodeAssistResponse;
return {
tierLabel: resolveGeminiCliTierLabel(payload),
tierId: resolveGeminiCliTierId(payload),
creditBalance: resolveGeminiCliCreditBalance(payload),
};
} catch (error) {
clearTimeout(timeoutId);
if (verbose) {
const message = error instanceof Error ? error.message : 'Unknown error';
console.error(`[i] Gemini CLI supplementary metadata skipped: ${message}`);
}
return { tierLabel: null, tierId: null, creditBalance: null };
}
} }
function buildGeminiCliFailureResult( function buildGeminiCliFailureResult(
@@ -271,6 +371,9 @@ function buildGeminiCliFailureResult(
success: false, success: false,
buckets: [], buckets: [],
projectId, projectId,
tierLabel: null,
tierId: null,
creditBalance: null,
lastUpdated: Date.now(), lastUpdated: Date.now(),
accountId, accountId,
error: options.error, error: options.error,
@@ -480,81 +583,38 @@ function buildGeminiCliHttpFailureResult(
* Groups buckets by model series and token type * Groups buckets by model series and token type
*/ */
function buildGeminiCliBuckets(rawBuckets: RawGeminiCliBucket[]): GeminiCliBucket[] { function buildGeminiCliBuckets(rawBuckets: RawGeminiCliBucket[]): GeminiCliBucket[] {
// Group buckets by groupId::tokenType const parsedBuckets = rawBuckets
const grouped = new Map< .map((bucket): GeminiCliParsedBucket | null => {
string, const modelId = normalizeStringValue(bucket.model_id ?? bucket.modelId);
{ if (!modelId) return null;
label: string;
tokenType: string | null;
remainingFraction: number;
resetTime: string | null;
modelIds: string[];
}
>();
for (const bucket of rawBuckets) { const tokenType = normalizeStringValue(bucket.token_type ?? bucket.tokenType);
const modelId = bucket.model_id || bucket.modelId || ''; const remainingFractionRaw = normalizeNumberValue(
if (!modelId) continue; bucket.remaining_fraction ?? bucket.remainingFraction
);
const remainingAmount = normalizeNumberValue(
bucket.remaining_amount ?? bucket.remainingAmount
);
const resetTime = normalizeStringValue(bucket.reset_time ?? bucket.resetTime);
// Skip ignored models let fallbackFraction: number | null = null;
if (shouldIgnoreModel(modelId)) continue; if (remainingAmount !== null) {
fallbackFraction = remainingAmount <= 0 ? 0 : null;
const tokenType = bucket.token_type ?? bucket.tokenType ?? null; } else if (resetTime) {
// Clamp remainingFraction to [0, 1] range fallbackFraction = 0;
const rawRemainingFraction = bucket.remaining_fraction ?? bucket.remainingFraction ?? 1;
const remainingFraction = Math.max(0, Math.min(1, rawRemainingFraction));
const resetTime = bucket.reset_time ?? bucket.resetTime ?? null;
// Find group for this model
const group = findModelGroup(modelId);
const groupId = group?.groupId || 'other';
const label = group?.label || 'Other Models';
// Create compound key for grouping
const key = `${groupId}::${tokenType || 'combined'}`;
const existing = grouped.get(key);
if (existing) {
// Merge: take the minimum remaining fraction (most limiting)
existing.remainingFraction = Math.min(existing.remainingFraction, remainingFraction);
// Keep earliest reset time if available
if (resetTime && (!existing.resetTime || resetTime < existing.resetTime)) {
existing.resetTime = resetTime;
} }
existing.modelIds.push(modelId);
} else { return {
grouped.set(key, { modelId,
label,
tokenType, tokenType,
remainingFraction, remainingFraction: remainingFractionRaw ?? fallbackFraction ?? 1,
remainingAmount,
resetTime, resetTime,
modelIds: [modelId], };
}); })
} .filter((bucket): bucket is GeminiCliParsedBucket => bucket !== null);
}
// Convert to array return buildGeminiCliBucketsFromParsedBuckets(parsedBuckets);
const buckets: GeminiCliBucket[] = [];
for (const [key, data] of grouped.entries()) {
buckets.push({
id: key,
label: data.label,
tokenType: data.tokenType,
remainingFraction: data.remainingFraction,
remainingPercent: Math.round(data.remainingFraction * 100),
resetTime: data.resetTime,
modelIds: data.modelIds,
});
}
// Sort by label then token type
buckets.sort((a, b) => {
const labelCompare = a.label.localeCompare(b.label);
if (labelCompare !== 0) return labelCompare;
return (a.tokenType || '').localeCompare(b.tokenType || '');
});
return buckets;
} }
/** /**
@@ -577,12 +637,11 @@ async function fetchWithAuthData(
}); });
} }
const url = `${GEMINI_CLI_API_BASE}/${GEMINI_CLI_API_VERSION}:retrieveUserQuota`;
const controller = new AbortController(); const controller = new AbortController();
const timeoutId = setTimeout(() => controller.abort(), 5000); const timeoutId = setTimeout(() => controller.abort(), 5000);
try { try {
const response = await fetch(url, { const response = await fetch(GEMINI_CLI_QUOTA_URL, {
method: 'POST', method: 'POST',
signal: controller.signal, signal: controller.signal,
headers: { headers: {
@@ -609,6 +668,11 @@ async function fetchWithAuthData(
const data = (await response.json()) as GeminiCliQuotaResponse; const data = (await response.json()) as GeminiCliQuotaResponse;
const rawBuckets = data.buckets || []; const rawBuckets = data.buckets || [];
const buckets = buildGeminiCliBuckets(rawBuckets); const buckets = buildGeminiCliBuckets(rawBuckets);
const supplementary = await fetchGeminiCliSupplementary(
authData.accessToken,
authData.projectId,
verbose
);
if (verbose) console.error(`[i] Gemini CLI buckets found: ${buckets.length}`); if (verbose) console.error(`[i] Gemini CLI buckets found: ${buckets.length}`);
@@ -616,6 +680,9 @@ async function fetchWithAuthData(
success: true, success: true,
buckets, buckets,
projectId: authData.projectId, projectId: authData.projectId,
tierLabel: supplementary.tierLabel,
tierId: supplementary.tierId,
creditBalance: supplementary.creditBalance,
lastUpdated: Date.now(), lastUpdated: Date.now(),
accountId, accountId,
}; };
+8
View File
@@ -176,6 +176,8 @@ export interface GeminiCliBucket {
remainingFraction: number; remainingFraction: number;
/** Remaining quota as percentage (0-100) */ /** Remaining quota as percentage (0-100) */
remainingPercent: number; remainingPercent: number;
/** Remaining quota count when the upstream API provides it */
remainingAmount?: number | null;
/** ISO timestamp when quota resets, null if unknown */ /** ISO timestamp when quota resets, null if unknown */
resetTime: string | null; resetTime: string | null;
/** Model IDs in this bucket */ /** Model IDs in this bucket */
@@ -192,6 +194,12 @@ export interface GeminiCliQuotaResult extends QuotaErrorMetadata {
buckets: GeminiCliBucket[]; buckets: GeminiCliBucket[];
/** GCP project ID for this account */ /** GCP project ID for this account */
projectId: string | null; projectId: string | null;
/** Human-readable Gemini tier label when available */
tierLabel?: string | null;
/** Stable Gemini tier identifier when available */
tierId?: string | null;
/** Available Google One AI credits when reported by the API */
creditBalance?: number | null;
/** Timestamp of fetch */ /** Timestamp of fetch */
lastUpdated: number; lastUpdated: number;
/** Error message if fetch failed */ /** Error message if fetch failed */
+11 -1
View File
@@ -638,15 +638,25 @@ function displayGeminiCliQuotaSection(
if (quota.projectId) { if (quota.projectId) {
console.log(` Project: ${dim(quota.projectId)}`); console.log(` Project: ${dim(quota.projectId)}`);
} }
if (quota.tierLabel) {
console.log(` Tier: ${dim(quota.tierLabel)}`);
}
if (quota.creditBalance !== null && quota.creditBalance !== undefined) {
console.log(` Credits: ${dim(quota.creditBalance.toLocaleString())}`);
}
for (const bucket of quota.buckets) { for (const bucket of quota.buckets) {
const bar = formatQuotaBar(bucket.remainingPercent); const bar = formatQuotaBar(bucket.remainingPercent);
const tokenLabel = bucket.tokenType ? dim(` (${bucket.tokenType})`) : ''; const tokenLabel = bucket.tokenType ? dim(` (${bucket.tokenType})`) : '';
const amountLabel =
bucket.remainingAmount !== null && bucket.remainingAmount !== undefined
? dim(` ${bucket.remainingAmount.toLocaleString()} left`)
: '';
const resetLabel = bucket.resetTime const resetLabel = bucket.resetTime
? dim(` Resets ${formatResetTimeISO(bucket.resetTime)}`) ? dim(` Resets ${formatResetTimeISO(bucket.resetTime)}`)
: ''; : '';
console.log( console.log(
` ${bucket.label.padEnd(24)} ${bar} ${bucket.remainingPercent.toFixed(0)}%${tokenLabel}${resetLabel}` ` ${bucket.label.padEnd(24)} ${bar} ${bucket.remainingPercent.toFixed(0)}%${tokenLabel}${amountLabel}${resetLabel}`
); );
} }
console.log(''); console.log('');
@@ -12,6 +12,7 @@ import { getCapturedFetchRequests, mockFetch, restoreFetch } from '../../mocks';
describe('Gemini CLI Quota Fetcher', () => { describe('Gemini CLI Quota Fetcher', () => {
const GEMINI_QUOTA_URL = 'https://cloudcode-pa.googleapis.com/v1internal:retrieveUserQuota'; const GEMINI_QUOTA_URL = 'https://cloudcode-pa.googleapis.com/v1internal:retrieveUserQuota';
const GEMINI_CODE_ASSIST_URL = 'https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist';
const GOOGLE_TOKEN_URL = 'https://oauth2.googleapis.com/token'; const GOOGLE_TOKEN_URL = 'https://oauth2.googleapis.com/token';
let tempHome: string; let tempHome: string;
let originalCcsHome: string | undefined; let originalCcsHome: string | undefined;
@@ -159,15 +160,29 @@ describe('Gemini CLI Quota Fetcher', () => {
const flashBucket = buckets.find((b) => b.label === 'Gemini Flash Series'); const flashBucket = buckets.find((b) => b.label === 'Gemini Flash Series');
expect(flashBucket).toBeDefined(); expect(flashBucket).toBeDefined();
// Takes minimum remaining fraction (0.6) // Uses the preferred representative model when it exists
expect(flashBucket!.remainingFraction).toBe(0.6); expect(flashBucket!.remainingFraction).toBe(0.8);
expect(flashBucket!.remainingPercent).toBe(60); expect(flashBucket!.remainingPercent).toBe(80);
const proBucket = buckets.find((b) => b.label === 'Gemini Pro Series'); const proBucket = buckets.find((b) => b.label === 'Gemini Pro Series');
expect(proBucket).toBeDefined(); expect(proBucket).toBeDefined();
expect(proBucket!.remainingFraction).toBe(0.9); expect(proBucket!.remainingFraction).toBe(0.9);
}); });
it('should split Flash Lite into its own bucket', () => {
const rawBuckets = [
{ model_id: 'gemini-2.5-flash-lite', remaining_fraction: 1 },
{ model_id: 'gemini-2.5-flash', remaining_fraction: 0.7 },
];
const buckets = buildGeminiCliBuckets(rawBuckets);
expect(buckets.map((bucket) => bucket.label)).toEqual([
'Gemini Flash Lite Series',
'Gemini Flash Series',
]);
});
it('should recognize Gemini 3.1 preview IDs during the rollout', () => { it('should recognize Gemini 3.1 preview IDs during the rollout', () => {
const rawBuckets = [ const rawBuckets = [
{ model_id: 'gemini-3.1-flash-preview', remaining_fraction: 0.7 }, { model_id: 'gemini-3.1-flash-preview', remaining_fraction: 0.7 },
@@ -242,13 +257,13 @@ describe('Gemini CLI Quota Fetcher', () => {
expect(buckets[0].remainingFraction).toBe(0.9); expect(buckets[0].remainingFraction).toBe(0.9);
}); });
it('should categorize unknown models as "other"', () => { it('should preserve unknown model IDs instead of collapsing them', () => {
const rawBuckets = [{ model_id: 'unknown-model-xyz', remaining_fraction: 0.7 }]; const rawBuckets = [{ model_id: 'unknown-model-xyz', remaining_fraction: 0.7 }];
const buckets = buildGeminiCliBuckets(rawBuckets); const buckets = buildGeminiCliBuckets(rawBuckets);
expect(buckets).toHaveLength(1); expect(buckets).toHaveLength(1);
expect(buckets[0].label).toBe('Other Models'); expect(buckets[0].label).toBe('unknown-model-xyz');
}); });
it('should handle empty buckets array', () => { it('should handle empty buckets array', () => {
@@ -268,7 +283,7 @@ describe('Gemini CLI Quota Fetcher', () => {
expect(buckets[0].remainingFraction).toBe(0.8); expect(buckets[0].remainingFraction).toBe(0.8);
}); });
it('should keep earliest reset time when merging', () => { it('should keep the representative model reset time when it exists', () => {
const rawBuckets = [ const rawBuckets = [
{ {
model_id: 'gemini-3-flash-preview', model_id: 'gemini-3-flash-preview',
@@ -284,6 +299,26 @@ describe('Gemini CLI Quota Fetcher', () => {
const buckets = buildGeminiCliBuckets(rawBuckets); const buckets = buildGeminiCliBuckets(rawBuckets);
const flashBucket = buckets.find((b) => b.label === 'Gemini Flash Series');
expect(flashBucket!.resetTime).toBe('2026-01-30T12:00:00Z');
});
it('should keep earliest reset time when the representative model is missing', () => {
const rawBuckets = [
{
model_id: 'gemini-3.1-flash-preview',
remaining_fraction: 0.8,
reset_time: '2026-01-30T12:00:00Z',
},
{
model_id: 'gemini-2.5-flash',
remaining_fraction: 0.6,
reset_time: '2026-01-30T10:00:00Z',
},
];
const buckets = buildGeminiCliBuckets(rawBuckets);
const flashBucket = buckets.find((b) => b.label === 'Gemini Flash Series'); const flashBucket = buckets.find((b) => b.label === 'Gemini Flash Series');
expect(flashBucket!.resetTime).toBe('2026-01-30T10:00:00Z'); expect(flashBucket!.resetTime).toBe('2026-01-30T10:00:00Z');
}); });
@@ -311,6 +346,110 @@ describe('Gemini CLI Quota Fetcher', () => {
}); });
}); });
describe('fetchGeminiCliQuota success metadata', () => {
it('merges tier and credit metadata into successful quota responses', async () => {
writeActiveGeminiAccount('success@example.com');
mockFetch([
{
url: GEMINI_QUOTA_URL,
method: 'POST',
status: 200,
response: {
buckets: [
{
model_id: 'gemini-2.5-flash-lite',
remaining_fraction: 1,
remaining_amount: 100,
reset_time: '2026-01-30T09:00:00Z',
},
{
model_id: 'gemini-3-flash-preview',
remaining_fraction: 0.82,
remaining_amount: 82,
reset_time: '2026-01-30T14:00:00Z',
},
{
model_id: 'gemini-2.5-flash',
remaining_fraction: 0.4,
remaining_amount: 40,
reset_time: '2026-01-30T10:00:00Z',
},
{
model_id: 'gemini-3.1-pro-preview',
remaining_fraction: 0.91,
remaining_amount: 91,
reset_time: '2026-01-30T15:00:00Z',
},
],
},
},
{
url: GEMINI_CODE_ASSIST_URL,
method: 'POST',
status: 200,
response: {
paidTier: {
id: 'g1-pro-tier',
availableCredits: [{ creditType: 'GOOGLE_ONE_AI', creditAmount: 12 }],
},
},
},
]);
const result = await fetchGeminiCliQuota('success@example.com');
expect(result.success).toBe(true);
expect(result.tierLabel).toBe('Pro');
expect(result.tierId).toBe('g1-pro-tier');
expect(result.creditBalance).toBe(12);
expect(result.buckets.map((bucket) => bucket.label)).toEqual([
'Gemini Flash Lite Series',
'Gemini Flash Series',
'Gemini Pro Series',
]);
expect(result.buckets[0].remainingAmount).toBe(100);
const flashBucket = result.buckets.find((bucket) => bucket.label === 'Gemini Flash Series');
expect(flashBucket?.remainingPercent).toBe(82);
expect(flashBucket?.remainingAmount).toBe(82);
expect(flashBucket?.resetTime).toBe('2026-01-30T14:00:00Z');
const requestUrls = getCapturedFetchRequests().map((request) => request.url);
expect(requestUrls).toContain(GEMINI_QUOTA_URL);
expect(requestUrls).toContain(GEMINI_CODE_ASSIST_URL);
});
it('keeps base quota success when supplementary metadata fails', async () => {
writeActiveGeminiAccount('supplementary-failure@example.com');
mockFetch([
{
url: GEMINI_QUOTA_URL,
method: 'POST',
status: 200,
response: {
buckets: [{ model_id: 'gemini-3-flash-preview', remaining_fraction: 0.75 }],
},
},
{
url: GEMINI_CODE_ASSIST_URL,
method: 'POST',
status: 503,
response: { error: { message: 'Service unavailable' } },
},
]);
const result = await fetchGeminiCliQuota('supplementary-failure@example.com');
expect(result.success).toBe(true);
expect(result.tierLabel).toBeNull();
expect(result.tierId).toBeNull();
expect(result.creditBalance).toBeNull();
expect(result.buckets[0].remainingPercent).toBe(75);
});
});
describe('fetchGeminiCliQuota failure metadata', () => { describe('fetchGeminiCliQuota failure metadata', () => {
it('maps 401 responses to reauth-required metadata', async () => { it('maps 401 responses to reauth-required metadata', async () => {
writeActiveGeminiAccount('reauth@example.com'); writeActiveGeminiAccount('reauth@example.com');