mirror of
https://github.com/tiennm99/ccs.git
synced 2026-08-23 16:23:35 +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 { 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,
|
||||||
};
|
};
|
||||||
|
|||||||
@@ -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 */
|
||||||
|
|||||||
@@ -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');
|
||||||
|
|||||||
Reference in New Issue
Block a user