diff --git a/src/cliproxy/quota-fetcher-codex.ts b/src/cliproxy/quota-fetcher-codex.ts new file mode 100644 index 00000000..27f66f92 --- /dev/null +++ b/src/cliproxy/quota-fetcher-codex.ts @@ -0,0 +1,372 @@ +/** + * Quota Fetcher for Codex (ChatGPT) Accounts + * + * Fetches quota information from ChatGPT backend API. + * Used for displaying rate limit windows and reset times. + */ + +import * as fs from 'node:fs'; +import * as path from 'node:path'; +import { getAuthDir } from './config-generator'; +import { getProviderAccounts, getPausedDir } from './account-manager'; +import type { CodexQuotaResult, CodexQuotaWindow } from './quota-types'; + +/** ChatGPT backend API base URL */ +const CODEX_API_BASE = 'https://chatgpt.com/backend-api'; + +/** User agent matching Codex CLI */ +const USER_AGENT = 'codex_cli_rs/0.76.0 (Debian 13.0.0; x86_64) WindowsTerminal'; + +/** Auth data extracted from Codex auth file */ +interface CodexAuthData { + accessToken: string; + accountId: string; // ChatGPT-Account-Id header + isExpired: boolean; + expiresAt: string | null; +} + +/** Raw API response structure */ +interface CodexUsageResponse { + plan_type?: string; + planType?: string; + rate_limit?: CodexRateLimitWindow; + rateLimit?: CodexRateLimitWindow; + code_review_rate_limit?: CodexRateLimitWindow; + codeReviewRateLimit?: CodexRateLimitWindow; +} + +/** Rate limit window from API */ +interface CodexRateLimitWindow { + primary_window?: CodexWindowData; + primaryWindow?: CodexWindowData; + secondary_window?: CodexWindowData; + secondaryWindow?: CodexWindowData; +} + +/** Individual window data */ +interface CodexWindowData { + used_percent?: number; + usedPercent?: number; + reset_after_seconds?: number | null; + resetAfterSeconds?: number | null; +} + +/** + * Sanitize email to match CLIProxyAPI auth file naming convention + */ +function sanitizeEmail(email: string): string { + return email.replace(/@/g, '_').replace(/\./g, '_'); +} + +/** + * Check if token is expired based on the expired timestamp + */ +function isTokenExpired(expiredStr?: string): boolean { + if (!expiredStr) return false; + try { + const expiredDate = new Date(expiredStr); + return expiredDate.getTime() < Date.now(); + } catch { + return false; + } +} + +/** + * Read auth data from Codex auth file + */ +function readCodexAuthData(accountId: string): CodexAuthData | null { + const authDirs = [getAuthDir(), getPausedDir()]; + const sanitizedId = sanitizeEmail(accountId); + const expectedFile = `codex-${sanitizedId}.json`; + + for (const authDir of authDirs) { + if (!fs.existsSync(authDir)) continue; + + const filePath = path.join(authDir, expectedFile); + if (fs.existsSync(filePath)) { + try { + const content = fs.readFileSync(filePath, 'utf-8'); + const data = JSON.parse(content); + if (!data.access_token) continue; + + return { + accessToken: data.access_token, + accountId: data.account_id || data.accountId || '', + isExpired: isTokenExpired(data.expired), + expiresAt: data.expired || null, + }; + } catch { + continue; + } + } + + // Fallback: scan directory for matching email in file content + const files = fs.readdirSync(authDir); + for (const file of files) { + if (file.startsWith('codex-') && file.endsWith('.json')) { + const candidatePath = path.join(authDir, file); + try { + const content = fs.readFileSync(candidatePath, 'utf-8'); + const data = JSON.parse(content); + if (data.email === accountId && data.access_token) { + return { + accessToken: data.access_token, + accountId: data.account_id || data.accountId || '', + isExpired: isTokenExpired(data.expired), + expiresAt: data.expired || null, + }; + } + } catch { + continue; + } + } + } + } + + return null; +} + +/** + * Build CodexQuotaWindow array from API response + * Handles both snake_case and camelCase field names + */ +function buildCodexQuotaWindows(payload: CodexUsageResponse): CodexQuotaWindow[] { + const windows: CodexQuotaWindow[] = []; + + // Get rate limit object (handles both cases) + const rateLimit = payload.rate_limit || payload.rateLimit; + const codeReviewRateLimit = payload.code_review_rate_limit || payload.codeReviewRateLimit; + + // Helper to extract window data + const addWindow = (label: string, windowData: CodexWindowData | undefined): void => { + if (!windowData) return; + + // Clamp usedPercent to [0, 100] range + const rawUsedPercent = windowData.used_percent ?? windowData.usedPercent ?? 0; + const usedPercent = Math.max(0, Math.min(100, rawUsedPercent)); + const resetAfterSeconds = + windowData.reset_after_seconds ?? windowData.resetAfterSeconds ?? null; + + // Calculate reset timestamp if we have seconds + let resetAt: string | null = null; + if (resetAfterSeconds !== null && resetAfterSeconds > 0) { + resetAt = new Date(Date.now() + resetAfterSeconds * 1000).toISOString(); + } + + windows.push({ + label, + usedPercent, + remainingPercent: Math.max(0, 100 - usedPercent), + resetAfterSeconds, + resetAt, + }); + }; + + // Add main rate limit windows + if (rateLimit) { + addWindow('Primary', rateLimit.primary_window || rateLimit.primaryWindow); + addWindow('Secondary', rateLimit.secondary_window || rateLimit.secondaryWindow); + } + + // Add code review rate limit windows + if (codeReviewRateLimit) { + addWindow( + 'Code Review (Primary)', + codeReviewRateLimit.primary_window || codeReviewRateLimit.primaryWindow + ); + addWindow( + 'Code Review (Secondary)', + codeReviewRateLimit.secondary_window || codeReviewRateLimit.secondaryWindow + ); + } + + return windows; +} + +/** + * Fetch quota for a single Codex account + * + * @param accountId - Account identifier (email) + * @param verbose - Show detailed diagnostics + * @returns Quota result with windows and percentages + */ +export async function fetchCodexQuota( + accountId: string, + verbose = false +): Promise { + if (verbose) console.error(`[i] Fetching Codex quota for ${accountId}...`); + + const authData = readCodexAuthData(accountId); + if (!authData) { + const error = 'Auth file not found for Codex account'; + if (verbose) console.error(`[!] Error: ${error}`); + return { + success: false, + windows: [], + planType: null, + lastUpdated: Date.now(), + error, + accountId, + }; + } + + if (authData.isExpired) { + const error = 'Token expired - re-authenticate with ccs cliproxy auth codex'; + if (verbose) console.error(`[!] Error: ${error}`); + return { + success: false, + windows: [], + planType: null, + lastUpdated: Date.now(), + error, + accountId, + }; + } + + if (!authData.accountId) { + const error = 'Missing ChatGPT-Account-Id in auth file'; + if (verbose) console.error(`[!] Error: ${error}`); + return { + success: false, + windows: [], + planType: null, + lastUpdated: Date.now(), + error, + accountId, + }; + } + + const url = `${CODEX_API_BASE}/wham/usage`; + const controller = new AbortController(); + const timeoutId = setTimeout(() => controller.abort(), 5000); + + try { + const response = await fetch(url, { + method: 'GET', + signal: controller.signal, + headers: { + Authorization: `Bearer ${authData.accessToken}`, + 'ChatGPT-Account-Id': authData.accountId, + 'User-Agent': USER_AGENT, + }, + }); + + clearTimeout(timeoutId); + + if (verbose) console.error(`[i] Codex API status: ${response.status}`); + + if (response.status === 401) { + return { + success: false, + windows: [], + planType: null, + lastUpdated: Date.now(), + error: 'Token expired or invalid', + accountId, + }; + } + + if (response.status === 403) { + return { + success: false, + windows: [], + planType: null, + lastUpdated: Date.now(), + error: 'Quota access not available (free plan may not have quota API access)', + accountId, + }; + } + + if (response.status === 429) { + return { + success: false, + windows: [], + planType: null, + lastUpdated: Date.now(), + error: 'Rate limited - try again later', + accountId, + }; + } + + if (!response.ok) { + return { + success: false, + windows: [], + planType: null, + lastUpdated: Date.now(), + error: `API error: ${response.status}`, + accountId, + }; + } + + const data = (await response.json()) as CodexUsageResponse; + const windows = buildCodexQuotaWindows(data); + + // Extract plan type + const planTypeRaw = data.plan_type || data.planType; + let planType: 'free' | 'plus' | 'team' | null = null; + if (planTypeRaw) { + const normalized = planTypeRaw.toLowerCase(); + if (normalized === 'free') planType = 'free'; + else if (normalized === 'plus') planType = 'plus'; + else if (normalized === 'team') planType = 'team'; + } + + if (verbose) console.error(`[i] Codex windows found: ${windows.length}`); + + return { + success: true, + windows, + planType, + lastUpdated: Date.now(), + accountId, + }; + } catch (err) { + clearTimeout(timeoutId); + const errorMsg = + err instanceof Error && err.name === 'AbortError' + ? 'Request timeout' + : err instanceof Error + ? err.message + : 'Unknown error'; + + if (verbose) console.error(`[!] Codex quota error: ${errorMsg}`); + + return { + success: false, + windows: [], + planType: null, + lastUpdated: Date.now(), + error: errorMsg, + accountId, + }; + } +} + +/** + * Fetch quota for all Codex accounts + * + * @param verbose - Show detailed diagnostics + * @returns Array of account quotas + */ +export async function fetchAllCodexQuotas( + verbose = false +): Promise<{ account: string; quota: CodexQuotaResult }[]> { + const accounts = getProviderAccounts('codex'); + + if (accounts.length === 0) { + return []; + } + + const results = await Promise.all( + accounts.map(async (account) => ({ + account: account.id, + quota: await fetchCodexQuota(account.id, verbose), + })) + ); + + return results; +} + +// Export for testing +export { readCodexAuthData, buildCodexQuotaWindows }; diff --git a/src/cliproxy/quota-fetcher-gemini-cli.ts b/src/cliproxy/quota-fetcher-gemini-cli.ts new file mode 100644 index 00000000..e013e23e --- /dev/null +++ b/src/cliproxy/quota-fetcher-gemini-cli.ts @@ -0,0 +1,444 @@ +/** + * Quota Fetcher for Gemini CLI Accounts + * + * Fetches quota information from Google Cloud Code internal API. + * Used for displaying bucket-based quotas grouped by model series. + */ + +import * as fs from 'node:fs'; +import * as path from 'node:path'; +import { getAuthDir } from './config-generator'; +import { getProviderAccounts, getPausedDir } from './account-manager'; +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'; + +/** Model groups for quota consolidation */ +const GEMINI_CLI_GROUPS: Record< + string, + { + label: string; + models: string[]; + preferred: string; + } +> = { + 'gemini-flash-series': { + label: 'Gemini Flash Series', + models: ['gemini-3-flash-preview', 'gemini-2.5-flash', 'gemini-2.5-flash-lite'], + preferred: 'gemini-3-flash-preview', + }, + 'gemini-pro-series': { + label: 'Gemini Pro Series', + models: ['gemini-3-pro-preview', 'gemini-2.5-pro'], + preferred: 'gemini-3-pro-preview', + }, +}; + +/** 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; + projectId: string | null; + isExpired: boolean; + expiresAt: string | null; +} + +/** Raw bucket from API response */ +interface RawGeminiCliBucket { + model_id?: string; + modelId?: string; + token_type?: string | null; + tokenType?: string | null; + remaining_fraction?: number; + remainingFraction?: number; + remaining_amount?: number; + remainingAmount?: number; + reset_time?: string | null; + resetTime?: string | null; +} + +/** Raw API response structure */ +interface GeminiCliQuotaResponse { + buckets?: RawGeminiCliBucket[]; +} + +/** + * Extract project ID from account field + * Input: "user@example.com (cloudaicompanion-abc-123)" + * Output: "cloudaicompanion-abc-123" + */ +function resolveGeminiCliProjectId(accountField: string): string | null { + const regex = /\(([^()]+)\)/g; + let match: RegExpExecArray | null; + let lastMatch: string | null = null; + while ((match = regex.exec(accountField)) !== null) { + lastMatch = match[1]; + } + return lastMatch; +} + +/** + * Sanitize email to match CLIProxyAPI auth file naming convention + */ +function sanitizeEmail(email: string): string { + return email.replace(/@/g, '_').replace(/\./g, '_'); +} + +/** + * Check if token is expired based on the expired timestamp + */ +function isTokenExpired(expiredStr?: string): boolean { + if (!expiredStr) return false; + try { + const expiredDate = new Date(expiredStr); + return expiredDate.getTime() < Date.now(); + } catch { + return false; + } +} + +/** + * Read auth data from Gemini CLI auth file + */ +function readGeminiCliAuthData(accountId: string): GeminiCliAuthData | null { + const authDirs = [getAuthDir(), getPausedDir()]; + const sanitizedId = sanitizeEmail(accountId); + const expectedFile = `gemini-${sanitizedId}.json`; + + for (const authDir of authDirs) { + if (!fs.existsSync(authDir)) continue; + + const filePath = path.join(authDir, expectedFile); + if (fs.existsSync(filePath)) { + try { + const content = fs.readFileSync(filePath, 'utf-8'); + const data = JSON.parse(content); + if (!data.access_token) continue; + + // Extract project ID from account field + const accountField = data.account || ''; + const projectId = resolveGeminiCliProjectId(accountField); + + return { + accessToken: data.access_token, + projectId, + isExpired: isTokenExpired(data.expired), + expiresAt: data.expired || null, + }; + } catch { + continue; + } + } + + // Fallback: scan directory for matching email in file content + const files = fs.readdirSync(authDir); + for (const file of files) { + if (file.startsWith('gemini-') && file.endsWith('.json')) { + const candidatePath = path.join(authDir, file); + try { + const content = fs.readFileSync(candidatePath, 'utf-8'); + const data = JSON.parse(content); + if (data.email === accountId && data.access_token) { + const accountField = data.account || ''; + const projectId = resolveGeminiCliProjectId(accountField); + + return { + accessToken: data.access_token, + projectId, + isExpired: isTokenExpired(data.expired), + expiresAt: data.expired || null, + }; + } + } catch { + continue; + } + } + } + } + + 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 }; + } + } + return null; +} + +/** + * Check if model should be ignored + */ +function shouldIgnoreModel(modelId: string): boolean { + return IGNORED_MODEL_PREFIXES.some((prefix) => modelId.startsWith(prefix)); +} + +/** + * Build GeminiCliBucket array from API response + * 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[]; + } + >(); + + for (const bucket of rawBuckets) { + const modelId = bucket.model_id || bucket.modelId || ''; + if (!modelId) continue; + + // 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; + } + existing.modelIds.push(modelId); + } else { + grouped.set(key, { + label, + tokenType, + remainingFraction, + resetTime, + modelIds: [modelId], + }); + } + } + + // 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; +} + +/** + * Fetch quota for a single Gemini CLI account + * + * @param accountId - Account identifier (email) + * @param verbose - Show detailed diagnostics + * @returns Quota result with buckets and percentages + */ +export async function fetchGeminiCliQuota( + accountId: string, + verbose = false +): Promise { + if (verbose) console.error(`[i] Fetching Gemini CLI quota for ${accountId}...`); + + const authData = readGeminiCliAuthData(accountId); + if (!authData) { + const error = 'Auth file not found for Gemini account'; + if (verbose) console.error(`[!] Error: ${error}`); + return { + success: false, + buckets: [], + projectId: null, + lastUpdated: Date.now(), + error, + accountId, + }; + } + + if (authData.isExpired) { + const error = 'Token expired - re-authenticate with ccs cliproxy auth gemini'; + if (verbose) console.error(`[!] Error: ${error}`); + return { + success: false, + buckets: [], + projectId: null, + lastUpdated: Date.now(), + error, + accountId, + }; + } + + if (!authData.projectId) { + const error = 'Cannot resolve project ID from auth file'; + if (verbose) console.error(`[!] Error: ${error}`); + return { + success: false, + buckets: [], + projectId: null, + lastUpdated: Date.now(), + error, + accountId, + }; + } + + 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, { + method: 'POST', + signal: controller.signal, + headers: { + Authorization: `Bearer ${authData.accessToken}`, + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ project: authData.projectId }), + }); + + clearTimeout(timeoutId); + + if (verbose) console.error(`[i] Gemini CLI API status: ${response.status}`); + + if (response.status === 401) { + return { + success: false, + buckets: [], + projectId: authData.projectId, + lastUpdated: Date.now(), + error: 'Token expired or invalid', + accountId, + }; + } + + if (response.status === 403) { + return { + success: false, + buckets: [], + projectId: authData.projectId, + lastUpdated: Date.now(), + error: 'Quota access forbidden for this account', + accountId, + }; + } + + if (response.status === 429) { + return { + success: false, + buckets: [], + projectId: authData.projectId, + lastUpdated: Date.now(), + error: 'Rate limited - try again later', + accountId, + }; + } + + if (!response.ok) { + return { + success: false, + buckets: [], + projectId: authData.projectId, + lastUpdated: Date.now(), + error: `API error: ${response.status}`, + accountId, + }; + } + + const data = (await response.json()) as GeminiCliQuotaResponse; + const rawBuckets = data.buckets || []; + const buckets = buildGeminiCliBuckets(rawBuckets); + + if (verbose) console.error(`[i] Gemini CLI buckets found: ${buckets.length}`); + + return { + success: true, + buckets, + projectId: authData.projectId, + lastUpdated: Date.now(), + accountId, + }; + } catch (err) { + clearTimeout(timeoutId); + const errorMsg = + err instanceof Error && err.name === 'AbortError' + ? 'Request timeout' + : err instanceof Error + ? err.message + : 'Unknown error'; + + if (verbose) console.error(`[!] Gemini CLI quota error: ${errorMsg}`); + + return { + success: false, + buckets: [], + projectId: authData.projectId, + lastUpdated: Date.now(), + error: errorMsg, + accountId, + }; + } +} + +/** + * Fetch quota for all Gemini CLI accounts + * + * @param verbose - Show detailed diagnostics + * @returns Array of account quotas + */ +export async function fetchAllGeminiCliQuotas( + verbose = false +): Promise<{ account: string; quota: GeminiCliQuotaResult }[]> { + const accounts = getProviderAccounts('gemini'); + + if (accounts.length === 0) { + return []; + } + + const results = await Promise.all( + accounts.map(async (account) => ({ + account: account.id, + quota: await fetchGeminiCliQuota(account.id, verbose), + })) + ); + + return results; +} + +// Export for testing +export { resolveGeminiCliProjectId, buildGeminiCliBuckets }; diff --git a/src/cliproxy/quota-types.ts b/src/cliproxy/quota-types.ts new file mode 100644 index 00000000..fa4c2277 --- /dev/null +++ b/src/cliproxy/quota-types.ts @@ -0,0 +1,109 @@ +/** + * Shared Quota Type Definitions + * + * Unified types for multi-provider quota system. + * Supports Antigravity, Codex, and Gemini CLI providers. + */ + +import type { QuotaResult as AntigravityQuotaResult } from './quota-fetcher'; + +/** Supported quota providers */ +export type QuotaProvider = 'agy' | 'codex' | 'gemini'; + +// Re-export Antigravity types for unified access +export type { QuotaResult as AntigravityQuotaResult } from './quota-fetcher'; + +/** + * Codex quota window (primary, secondary, code review) + */ +export interface CodexQuotaWindow { + /** Window label: "Primary", "Secondary", "Code Review (Primary)", "Code Review (Secondary)" */ + label: string; + /** Percentage used (0-100) */ + usedPercent: number; + /** Percentage remaining (100 - usedPercent) */ + remainingPercent: number; + /** Seconds until quota resets, null if unknown */ + resetAfterSeconds: number | null; + /** ISO timestamp when quota resets, null if unknown */ + resetAt: string | null; +} + +/** + * Codex quota fetch result + */ +export interface CodexQuotaResult { + /** Whether fetch succeeded */ + success: boolean; + /** Quota windows (primary, secondary, code review) */ + windows: CodexQuotaWindow[]; + /** Plan type: free, plus, team, or null if unknown */ + planType: 'free' | 'plus' | 'team' | null; + /** Timestamp of fetch */ + lastUpdated: number; + /** Error message if fetch failed */ + error?: string; + /** Account ID (email) this quota belongs to */ + accountId?: string; +} + +/** + * Gemini CLI quota bucket (grouped by model series and token type) + */ +export interface GeminiCliBucket { + /** Unique bucket identifier (e.g., "gemini-flash-series::input") */ + id: string; + /** Display label (e.g., "Gemini Flash Series") */ + label: string; + /** Token type: "input", "output", or null if combined */ + tokenType: string | null; + /** Remaining quota as fraction (0-1) */ + remainingFraction: number; + /** Remaining quota as percentage (0-100) */ + remainingPercent: number; + /** ISO timestamp when quota resets, null if unknown */ + resetTime: string | null; + /** Model IDs in this bucket */ + modelIds: string[]; +} + +/** + * Gemini CLI quota fetch result + */ +export interface GeminiCliQuotaResult { + /** Whether fetch succeeded */ + success: boolean; + /** Quota buckets grouped by model series */ + buckets: GeminiCliBucket[]; + /** GCP project ID for this account */ + projectId: string | null; + /** Timestamp of fetch */ + lastUpdated: number; + /** Error message if fetch failed */ + error?: string; + /** Account ID (email) this quota belongs to */ + accountId?: string; +} + +/** + * Unified quota result wrapper for CLI/Dashboard + * Contains provider-specific data in nested fields + */ +export interface UnifiedQuotaResult { + /** Provider this result belongs to */ + provider: QuotaProvider; + /** Whether fetch succeeded */ + success: boolean; + /** Timestamp of fetch */ + lastUpdated: number; + /** Error message if fetch failed */ + error?: string; + /** Account ID (email) this quota belongs to */ + accountId?: string; + /** Antigravity-specific quota data */ + antigravity?: AntigravityQuotaResult; + /** Codex-specific quota data */ + codex?: CodexQuotaResult; + /** Gemini CLI-specific quota data */ + geminiCli?: GeminiCliQuotaResult; +} diff --git a/src/commands/cliproxy-command.ts b/src/commands/cliproxy-command.ts index 3e94af8d..e15474f7 100644 --- a/src/commands/cliproxy-command.ts +++ b/src/commands/cliproxy-command.ts @@ -28,6 +28,9 @@ import { findAccountByQuery, } from '../cliproxy/account-manager'; import { fetchAllProviderQuotas } from '../cliproxy/quota-fetcher'; +import { fetchAllCodexQuotas } from '../cliproxy/quota-fetcher-codex'; +import { fetchAllGeminiCliQuotas } from '../cliproxy/quota-fetcher-gemini-cli'; +import type { CodexQuotaResult, GeminiCliQuotaResult } from '../cliproxy/quota-types'; import { isOnCooldown } from '../cliproxy/quota-manager'; import { DEFAULT_BACKEND, getFallbackVersion, BACKEND_CONFIG } from '../cliproxy/platform-detector'; import { CLIPROXY_PROFILES, CLIProxyProfileName } from '../auth/profile-detector'; @@ -71,6 +74,66 @@ import { handleSync } from './cliproxy-sync-handler'; // ARGUMENT PARSING // ============================================================================ +/** + * Parse --provider flag from args for quota command + * Returns the provider filter value and remaining args + * Accepts: agy, codex, gemini, gemini-cli, all + */ +function parseProviderArg(args: string[]): { + provider: 'agy' | 'codex' | 'gemini' | 'all'; + remainingArgs: string[]; +} { + const providerIdx = args.indexOf('--provider'); + if (providerIdx === -1) { + // Also check for --provider=value format + const providerEqualsIdx = args.findIndex((a) => a.startsWith('--provider=')); + if (providerEqualsIdx !== -1) { + const value = args[providerEqualsIdx].split('=')[1]?.toLowerCase() || ''; + const remainingArgs = [...args]; + remainingArgs.splice(providerEqualsIdx, 1); + // Handle empty value + if (!value) { + console.error( + warn('--provider requires a value. Valid options: agy, codex, gemini, gemini-cli, all') + ); + return { provider: 'all', remainingArgs }; + } + // Normalize gemini-cli to gemini + const normalized = value === 'gemini-cli' ? 'gemini' : value; + if ( + normalized !== 'agy' && + normalized !== 'codex' && + normalized !== 'gemini' && + normalized !== 'all' + ) { + console.error( + warn(`Invalid provider '${value}'. Valid options: agy, codex, gemini, gemini-cli, all`) + ); + return { provider: 'all', remainingArgs }; + } + return { provider: normalized as 'agy' | 'codex' | 'gemini' | 'all', remainingArgs }; + } + return { provider: 'all', remainingArgs: args }; + } + const value = args[providerIdx + 1]?.toLowerCase() || 'all'; + const remainingArgs = [...args]; + remainingArgs.splice(providerIdx, 2); + // Normalize gemini-cli to gemini + const normalized = value === 'gemini-cli' ? 'gemini' : value; + if ( + normalized !== 'agy' && + normalized !== 'codex' && + normalized !== 'gemini' && + normalized !== 'all' + ) { + console.error( + warn(`Invalid provider '${value}'. Valid options: agy, codex, gemini, gemini-cli, all`) + ); + return { provider: 'all', remainingArgs }; + } + return { provider: normalized as 'agy' | 'codex' | 'gemini' | 'all', remainingArgs }; +} + /** * Parse --backend flag from args * Returns the backend value and remaining args without --backend flag @@ -635,7 +698,8 @@ async function showHelp(): Promise { ['default ', 'Set default account for rotation'], ['pause ', 'Pause account (skip in rotation)'], ['resume ', 'Resume paused account'], - ['quota', 'Show quota status for all accounts'], + ['quota', 'Show quota status for all providers'], + ['quota --provider ', 'Filter by provider (agy|codex|gemini)'], ], ], [ @@ -911,22 +975,76 @@ async function handleResumeAccount(args: string[]): Promise { } } -async function handleQuotaStatus(verbose = false): Promise { +async function handleQuotaStatus( + verbose = false, + providerFilter: 'agy' | 'codex' | 'gemini' | 'all' = 'all' +): Promise { await initUI(); console.log(header('Quota Status')); console.log(''); + const shouldFetch = { + agy: providerFilter === 'all' || providerFilter === 'agy', + codex: providerFilter === 'all' || providerFilter === 'codex', + gemini: providerFilter === 'all' || providerFilter === 'gemini', + }; + + console.log(dim('Fetching quotas...')); + + // Parallel fetch from all requested providers + const [agyResults, codexResults, geminiResults] = await Promise.all([ + shouldFetch.agy ? fetchAllProviderQuotas('agy', verbose) : null, + shouldFetch.codex ? fetchAllCodexQuotas(verbose) : null, + shouldFetch.gemini ? fetchAllGeminiCliQuotas(verbose) : null, + ]); + + console.log(''); + + // Display Antigravity section + if (agyResults && agyResults.accounts.length > 0) { + displayAntigravityQuotaSection(agyResults, verbose); + } else if (shouldFetch.agy) { + console.log(subheader('Antigravity (0 accounts)')); + console.log(info('No Antigravity accounts configured')); + console.log(` Run: ${color('ccs agy --auth', 'command')} to authenticate`); + console.log(''); + } + + // Display Codex section + if (codexResults && codexResults.length > 0) { + displayCodexQuotaSection(codexResults); + } else if (shouldFetch.codex) { + console.log(subheader('Codex (0 accounts)')); + console.log(info('No Codex accounts configured')); + console.log(` Run: ${color('ccs codex --auth', 'command')} to authenticate`); + console.log(''); + } + + // Display Gemini CLI section + if (geminiResults && geminiResults.length > 0) { + displayGeminiCliQuotaSection(geminiResults); + } else if (shouldFetch.gemini) { + console.log(subheader('Gemini CLI (0 accounts)')); + console.log(info('No Gemini CLI accounts configured')); + console.log(` Run: ${color('ccs gemini --auth', 'command')} to authenticate`); + console.log(''); + } +} + +/** + * Display Antigravity quota section + */ +function displayAntigravityQuotaSection( + quotaResult: Awaited>, + _verbose: boolean +): void { const provider: CLIProxyProvider = 'agy'; const accounts = getProviderAccounts(provider); - if (accounts.length === 0) { - console.log(info('No Antigravity accounts configured')); - console.log(` Run: ${color('ccs agy --auth', 'command')} to authenticate`); - return; - } - - console.log(dim('Fetching quotas...')); - const quotaResult = await fetchAllProviderQuotas(provider, verbose); + console.log( + subheader(`Antigravity (${accounts.length} account${accounts.length !== 1 ? 's' : ''})`) + ); + console.log(''); // Build table rows const rows: string[][] = []; @@ -961,7 +1079,6 @@ async function handleQuotaStatus(verbose = false): Promise { ]); } - console.log(''); console.log( table(rows, { head: ['', 'Account', 'Tier', 'Quota', 'Status'], @@ -969,25 +1086,120 @@ async function handleQuotaStatus(verbose = false): Promise { }) ); console.log(''); - console.log(info(`Default account marked with ${color('*', 'success')}`)); +} + +/** + * Display Codex quota section + */ +function displayCodexQuotaSection(results: { account: string; quota: CodexQuotaResult }[]): void { + console.log(subheader(`Codex (${results.length} account${results.length !== 1 ? 's' : ''})`)); console.log(''); - // Show summary of paused/cooldown accounts - const pausedCount = accounts.filter((a) => a.paused).length; - const cooldownCount = accounts.filter((a) => isOnCooldown(provider, a.id)).length; - if (pausedCount > 0) { - console.log( - warn(`${pausedCount} account(s) paused - use 'ccs cliproxy resume ' to re-enable`) - ); - } - if (cooldownCount > 0) { - console.log(info(`${cooldownCount} account(s) on cooldown (exhausted recently)`)); - } - if (pausedCount > 0 || cooldownCount > 0) { + for (const { account, quota } of results) { + const accountInfo = findAccountByQuery('codex', account); + const defaultMark = accountInfo?.isDefault ? color(' (default)', 'info') : ''; + + if (!quota.success) { + console.log(` ${fail(account)}${defaultMark}`); + console.log(` ${color(quota.error || 'Failed to fetch quota', 'error')}`); + console.log(''); + continue; + } + + // Calculate overall quota health + const avgQuota = + quota.windows.length > 0 + ? quota.windows.reduce((sum, w) => sum + w.remainingPercent, 0) / quota.windows.length + : 0; + const statusIcon = avgQuota > 50 ? ok('') : avgQuota > 10 ? warn('') : fail(''); + const planBadge = quota.planType ? color(` [${quota.planType}]`, 'info') : ''; + + console.log(` ${statusIcon}${account}${defaultMark}${planBadge}`); + + // Show windows + for (const window of quota.windows) { + const bar = formatQuotaBar(window.remainingPercent); + const resetLabel = window.resetAfterSeconds + ? dim(` Resets ${formatResetTime(window.resetAfterSeconds)}`) + : ''; + console.log( + ` ${window.label.padEnd(24)} ${bar} ${window.remainingPercent.toFixed(0)}%${resetLabel}` + ); + } console.log(''); } } +/** + * Display Gemini CLI quota section + */ +function displayGeminiCliQuotaSection( + results: { account: string; quota: GeminiCliQuotaResult }[] +): void { + console.log( + subheader(`Gemini CLI (${results.length} account${results.length !== 1 ? 's' : ''})`) + ); + console.log(''); + + for (const { account, quota } of results) { + const accountInfo = findAccountByQuery('gemini', account); + const defaultMark = accountInfo?.isDefault ? color(' (default)', 'info') : ''; + + if (!quota.success) { + console.log(` ${fail(account)}${defaultMark}`); + console.log(` ${color(quota.error || 'Failed to fetch quota', 'error')}`); + console.log(''); + continue; + } + + // Calculate overall quota health + const avgQuota = + quota.buckets.length > 0 + ? quota.buckets.reduce((sum, b) => sum + b.remainingPercent, 0) / quota.buckets.length + : 0; + const statusIcon = avgQuota > 50 ? ok('') : avgQuota > 10 ? warn('') : fail(''); + + console.log(` ${statusIcon}${account}${defaultMark}`); + if (quota.projectId) { + console.log(` Project: ${dim(quota.projectId)}`); + } + + // Show buckets + for (const bucket of quota.buckets) { + const bar = formatQuotaBar(bucket.remainingPercent); + const tokenLabel = bucket.tokenType ? dim(` (${bucket.tokenType})`) : ''; + const resetLabel = bucket.resetTime + ? dim(` Resets ${formatResetTimeISO(bucket.resetTime)}`) + : ''; + console.log( + ` ${bucket.label.padEnd(24)} ${bar} ${bucket.remainingPercent.toFixed(0)}%${tokenLabel}${resetLabel}` + ); + } + console.log(''); + } +} + +/** + * Format reset time from seconds + */ +function formatResetTime(seconds: number): string { + if (seconds <= 0) return 'now'; + if (seconds < 60) return `in ${seconds}s`; + if (seconds < 3600) return `in ${Math.round(seconds / 60)}m`; + return `in ${Math.round(seconds / 3600)}h`; +} + +/** + * Format reset time from ISO timestamp + */ +function formatResetTimeISO(isoTime: string): string { + if (!isoTime) return 'unknown'; + const resetDate = new Date(isoTime); + if (isNaN(resetDate.getTime())) return 'unknown'; + const seconds = Math.max(0, Math.round((resetDate.getTime() - Date.now()) / 1000)); + return formatResetTime(seconds); +} + // ============================================================================ // MAIN ROUTER // ============================================================================ @@ -1057,7 +1269,8 @@ export async function handleCliproxyCommand(args: string[]): Promise { } if (command === 'quota') { - await handleQuotaStatus(verbose); + const { provider: providerFilter } = parseProviderArg(remainingArgs.slice(1)); + await handleQuotaStatus(verbose, providerFilter); return; }