mirror of
https://github.com/tiennm99/ccs.git
synced 2026-07-16 04:18:05 +00:00
feat(cliproxy): enrich gemini quota metadata
This commit is contained in:
@@ -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)),
|
||||
};
|
||||
});
|
||||
}
|
||||
@@ -11,43 +11,28 @@ import { getAuthDir } from './config-generator';
|
||||
import { getProviderAccounts, getPausedDir } from './account-manager';
|
||||
import { sanitizeEmail, isTokenExpired } from './auth-utils';
|
||||
import { refreshGeminiToken } from './auth/gemini-token-refresh';
|
||||
import {
|
||||
buildGeminiCliBucketsFromParsedBuckets,
|
||||
type GeminiCliParsedBucket,
|
||||
} from './gemini-cli-quota-normalizer';
|
||||
import type { GeminiCliQuotaResult, GeminiCliBucket } from './quota-types';
|
||||
|
||||
/** Google Cloud Code API endpoints */
|
||||
const GEMINI_CLI_API_BASE = 'https://cloudcode-pa.googleapis.com';
|
||||
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_TRUNCATION_SUFFIX = '...[truncated]';
|
||||
|
||||
/**
|
||||
* Model groups for quota consolidation.
|
||||
* Update when Google releases new Gemini models to include them in quota display.
|
||||
*/
|
||||
const GEMINI_CLI_GROUPS: Record<
|
||||
string,
|
||||
{
|
||||
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'],
|
||||
},
|
||||
const GEMINI_CLI_G1_CREDIT_TYPE = 'GOOGLE_ONE_AI';
|
||||
const GEMINI_CLI_TIER_LABELS: Record<string, string> = {
|
||||
'free-tier': 'Free',
|
||||
'legacy-tier': 'Legacy',
|
||||
'standard-tier': 'Standard',
|
||||
'g1-pro-tier': 'Pro',
|
||||
'g1-ultra-tier': 'Ultra',
|
||||
};
|
||||
|
||||
/** Models to ignore in quota display (deprecated) */
|
||||
const IGNORED_MODEL_PREFIXES = ['gemini-2.0-flash'];
|
||||
|
||||
/** Auth data extracted from Gemini CLI auth file */
|
||||
interface GeminiCliAuthData {
|
||||
accessToken: string;
|
||||
@@ -75,12 +60,38 @@ interface GeminiCliQuotaResponse {
|
||||
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 {
|
||||
errorCode?: string;
|
||||
errorDetail?: string;
|
||||
message?: string;
|
||||
}
|
||||
|
||||
interface GeminiCliSupplementaryInfo {
|
||||
tierLabel: string | null;
|
||||
tierId: string | null;
|
||||
creditBalance: number | null;
|
||||
}
|
||||
|
||||
/**
|
||||
* Extract project ID from account field
|
||||
* Input: "user@example.com (cloudaicompanion-abc-123)"
|
||||
@@ -234,23 +245,112 @@ function readGeminiCliAuthData(accountId: string): GeminiCliAuthData | null {
|
||||
return null;
|
||||
}
|
||||
|
||||
/**
|
||||
* Find which group a model belongs to
|
||||
*/
|
||||
function findModelGroup(modelId: string): { groupId: string; label: string } | null {
|
||||
for (const [groupId, group] of Object.entries(GEMINI_CLI_GROUPS)) {
|
||||
if (group.models.includes(modelId)) {
|
||||
return { groupId, label: group.label };
|
||||
function normalizeStringValue(value: unknown): string | null {
|
||||
return typeof value === 'string' && value.trim().length > 0 ? value.trim() : null;
|
||||
}
|
||||
|
||||
function normalizeNumberValue(value: unknown): number | null {
|
||||
if (typeof value === 'number' && Number.isFinite(value)) {
|
||||
return value;
|
||||
}
|
||||
if (typeof value === 'string' && value.trim().length > 0) {
|
||||
const parsed = Number(value);
|
||||
if (Number.isFinite(parsed)) {
|
||||
return parsed;
|
||||
}
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if model should be ignored
|
||||
*/
|
||||
function shouldIgnoreModel(modelId: string): boolean {
|
||||
return IGNORED_MODEL_PREFIXES.some((prefix) => modelId.startsWith(prefix));
|
||||
function resolveGeminiCliTierId(payload: GeminiCliCodeAssistResponse | null): string | null {
|
||||
if (!payload) return null;
|
||||
const currentTier = payload.currentTier ?? payload.current_tier;
|
||||
const paidTier = payload.paidTier ?? payload.paid_tier;
|
||||
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(
|
||||
@@ -271,6 +371,9 @@ function buildGeminiCliFailureResult(
|
||||
success: false,
|
||||
buckets: [],
|
||||
projectId,
|
||||
tierLabel: null,
|
||||
tierId: null,
|
||||
creditBalance: null,
|
||||
lastUpdated: Date.now(),
|
||||
accountId,
|
||||
error: options.error,
|
||||
@@ -480,81 +583,38 @@ function buildGeminiCliHttpFailureResult(
|
||||
* Groups buckets by model series and token type
|
||||
*/
|
||||
function buildGeminiCliBuckets(rawBuckets: RawGeminiCliBucket[]): GeminiCliBucket[] {
|
||||
// Group buckets by groupId::tokenType
|
||||
const grouped = new Map<
|
||||
string,
|
||||
{
|
||||
label: string;
|
||||
tokenType: string | null;
|
||||
remainingFraction: number;
|
||||
resetTime: string | null;
|
||||
modelIds: string[];
|
||||
}
|
||||
>();
|
||||
const parsedBuckets = rawBuckets
|
||||
.map((bucket): GeminiCliParsedBucket | null => {
|
||||
const modelId = normalizeStringValue(bucket.model_id ?? bucket.modelId);
|
||||
if (!modelId) return null;
|
||||
|
||||
for (const bucket of rawBuckets) {
|
||||
const modelId = bucket.model_id || bucket.modelId || '';
|
||||
if (!modelId) continue;
|
||||
const tokenType = normalizeStringValue(bucket.token_type ?? bucket.tokenType);
|
||||
const remainingFractionRaw = normalizeNumberValue(
|
||||
bucket.remaining_fraction ?? bucket.remainingFraction
|
||||
);
|
||||
const remainingAmount = normalizeNumberValue(
|
||||
bucket.remaining_amount ?? bucket.remainingAmount
|
||||
);
|
||||
const resetTime = normalizeStringValue(bucket.reset_time ?? bucket.resetTime);
|
||||
|
||||
// Skip ignored models
|
||||
if (shouldIgnoreModel(modelId)) continue;
|
||||
|
||||
const tokenType = bucket.token_type ?? bucket.tokenType ?? null;
|
||||
// Clamp remainingFraction to [0, 1] range
|
||||
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;
|
||||
let fallbackFraction: number | null = null;
|
||||
if (remainingAmount !== null) {
|
||||
fallbackFraction = remainingAmount <= 0 ? 0 : null;
|
||||
} else if (resetTime) {
|
||||
fallbackFraction = 0;
|
||||
}
|
||||
existing.modelIds.push(modelId);
|
||||
} else {
|
||||
grouped.set(key, {
|
||||
label,
|
||||
|
||||
return {
|
||||
modelId,
|
||||
tokenType,
|
||||
remainingFraction,
|
||||
remainingFraction: remainingFractionRaw ?? fallbackFraction ?? 1,
|
||||
remainingAmount,
|
||||
resetTime,
|
||||
modelIds: [modelId],
|
||||
});
|
||||
}
|
||||
}
|
||||
};
|
||||
})
|
||||
.filter((bucket): bucket is GeminiCliParsedBucket => bucket !== null);
|
||||
|
||||
// Convert to array
|
||||
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;
|
||||
return buildGeminiCliBucketsFromParsedBuckets(parsedBuckets);
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -577,12 +637,11 @@ async function fetchWithAuthData(
|
||||
});
|
||||
}
|
||||
|
||||
const url = `${GEMINI_CLI_API_BASE}/${GEMINI_CLI_API_VERSION}:retrieveUserQuota`;
|
||||
const controller = new AbortController();
|
||||
const timeoutId = setTimeout(() => controller.abort(), 5000);
|
||||
|
||||
try {
|
||||
const response = await fetch(url, {
|
||||
const response = await fetch(GEMINI_CLI_QUOTA_URL, {
|
||||
method: 'POST',
|
||||
signal: controller.signal,
|
||||
headers: {
|
||||
@@ -609,6 +668,11 @@ async function fetchWithAuthData(
|
||||
const data = (await response.json()) as GeminiCliQuotaResponse;
|
||||
const rawBuckets = data.buckets || [];
|
||||
const buckets = buildGeminiCliBuckets(rawBuckets);
|
||||
const supplementary = await fetchGeminiCliSupplementary(
|
||||
authData.accessToken,
|
||||
authData.projectId,
|
||||
verbose
|
||||
);
|
||||
|
||||
if (verbose) console.error(`[i] Gemini CLI buckets found: ${buckets.length}`);
|
||||
|
||||
@@ -616,6 +680,9 @@ async function fetchWithAuthData(
|
||||
success: true,
|
||||
buckets,
|
||||
projectId: authData.projectId,
|
||||
tierLabel: supplementary.tierLabel,
|
||||
tierId: supplementary.tierId,
|
||||
creditBalance: supplementary.creditBalance,
|
||||
lastUpdated: Date.now(),
|
||||
accountId,
|
||||
};
|
||||
|
||||
@@ -176,6 +176,8 @@ export interface GeminiCliBucket {
|
||||
remainingFraction: number;
|
||||
/** Remaining quota as percentage (0-100) */
|
||||
remainingPercent: number;
|
||||
/** Remaining quota count when the upstream API provides it */
|
||||
remainingAmount?: number | null;
|
||||
/** ISO timestamp when quota resets, null if unknown */
|
||||
resetTime: string | null;
|
||||
/** Model IDs in this bucket */
|
||||
@@ -192,6 +194,12 @@ export interface GeminiCliQuotaResult extends QuotaErrorMetadata {
|
||||
buckets: GeminiCliBucket[];
|
||||
/** GCP project ID for this account */
|
||||
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 */
|
||||
lastUpdated: number;
|
||||
/** Error message if fetch failed */
|
||||
|
||||
@@ -638,15 +638,25 @@ function displayGeminiCliQuotaSection(
|
||||
if (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) {
|
||||
const bar = formatQuotaBar(bucket.remainingPercent);
|
||||
const tokenLabel = bucket.tokenType ? dim(` (${bucket.tokenType})`) : '';
|
||||
const amountLabel =
|
||||
bucket.remainingAmount !== null && bucket.remainingAmount !== undefined
|
||||
? dim(` ${bucket.remainingAmount.toLocaleString()} left`)
|
||||
: '';
|
||||
const resetLabel = bucket.resetTime
|
||||
? dim(` Resets ${formatResetTimeISO(bucket.resetTime)}`)
|
||||
: '';
|
||||
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('');
|
||||
|
||||
@@ -12,6 +12,7 @@ import { getCapturedFetchRequests, mockFetch, restoreFetch } from '../../mocks';
|
||||
|
||||
describe('Gemini CLI Quota Fetcher', () => {
|
||||
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';
|
||||
let tempHome: string;
|
||||
let originalCcsHome: string | undefined;
|
||||
@@ -159,15 +160,29 @@ describe('Gemini CLI Quota Fetcher', () => {
|
||||
|
||||
const flashBucket = buckets.find((b) => b.label === 'Gemini Flash Series');
|
||||
expect(flashBucket).toBeDefined();
|
||||
// Takes minimum remaining fraction (0.6)
|
||||
expect(flashBucket!.remainingFraction).toBe(0.6);
|
||||
expect(flashBucket!.remainingPercent).toBe(60);
|
||||
// Uses the preferred representative model when it exists
|
||||
expect(flashBucket!.remainingFraction).toBe(0.8);
|
||||
expect(flashBucket!.remainingPercent).toBe(80);
|
||||
|
||||
const proBucket = buckets.find((b) => b.label === 'Gemini Pro Series');
|
||||
expect(proBucket).toBeDefined();
|
||||
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', () => {
|
||||
const rawBuckets = [
|
||||
{ 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);
|
||||
});
|
||||
|
||||
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 buckets = buildGeminiCliBuckets(rawBuckets);
|
||||
|
||||
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', () => {
|
||||
@@ -268,7 +283,7 @@ describe('Gemini CLI Quota Fetcher', () => {
|
||||
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 = [
|
||||
{
|
||||
model_id: 'gemini-3-flash-preview',
|
||||
@@ -284,6 +299,26 @@ describe('Gemini CLI Quota Fetcher', () => {
|
||||
|
||||
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');
|
||||
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', () => {
|
||||
it('maps 401 responses to reauth-required metadata', async () => {
|
||||
writeActiveGeminiAccount('reauth@example.com');
|
||||
|
||||
Reference in New Issue
Block a user