diff --git a/config/base-claude.settings.json b/config/base-claude.settings.json new file mode 100644 index 00000000..739df901 --- /dev/null +++ b/config/base-claude.settings.json @@ -0,0 +1,10 @@ +{ + "env": { + "ANTHROPIC_BASE_URL": "http://127.0.0.1:8317/api/provider/claude", + "ANTHROPIC_AUTH_TOKEN": "ccs-internal-managed", + "ANTHROPIC_MODEL": "claude-sonnet-4-20250514", + "ANTHROPIC_DEFAULT_OPUS_MODEL": "claude-opus-4-20250514", + "ANTHROPIC_DEFAULT_SONNET_MODEL": "claude-sonnet-4-20250514", + "ANTHROPIC_DEFAULT_HAIKU_MODEL": "claude-haiku-3-5-20241022" + } +} diff --git a/package.json b/package.json index 96baede7..d5bbb33d 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "@kaitranntt/ccs", - "version": "7.30.1", + "version": "7.30.1-dev.1", "description": "Claude Code Switch - Instant profile switching between Claude Sonnet 4.5 and GLM 4.6", "keywords": [ "cli", diff --git a/src/cliproxy/auth-utils.ts b/src/cliproxy/auth-utils.ts new file mode 100644 index 00000000..77e371b1 --- /dev/null +++ b/src/cliproxy/auth-utils.ts @@ -0,0 +1,27 @@ +/** + * Shared Authentication Utilities + * + * Common functions for OAuth token handling across quota fetchers. + */ + +/** + * Sanitize email to match CLIProxyAPI auth file naming convention. + * Replaces @ and . with underscores for filesystem compatibility. + */ +export function sanitizeEmail(email: string): string { + return email.replace(/@/g, '_').replace(/\./g, '_'); +} + +/** + * Check if token is expired based on the expired timestamp. + * Returns false if timestamp is missing or invalid (fail-open for quota display). + */ +export function isTokenExpired(expiredStr?: string): boolean { + if (!expiredStr) return false; + try { + const expiredDate = new Date(expiredStr); + return expiredDate.getTime() < Date.now(); + } catch { + return false; + } +} diff --git a/src/cliproxy/quota-fetcher-codex.ts b/src/cliproxy/quota-fetcher-codex.ts new file mode 100644 index 00000000..d8fd53b6 --- /dev/null +++ b/src/cliproxy/quota-fetcher-codex.ts @@ -0,0 +1,356 @@ +/** + * 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 { sanitizeEmail, isTokenExpired } from './auth-utils'; +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 for API compatibility. + * Update when Codex CLI releases new versions to maintain compatibility. + */ +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; +} + +/** + * 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..f28c9b62 --- /dev/null +++ b/src/cliproxy/quota-fetcher-gemini-cli.ts @@ -0,0 +1,425 @@ +/** + * 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 { sanitizeEmail, isTokenExpired } from './auth-utils'; +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. + * 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-2.5-flash', 'gemini-2.5-flash-lite'], + }, + 'gemini-pro-series': { + label: 'Gemini Pro Series', + models: ['gemini-3-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 */ +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; +} + +/** + * 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-fetcher.ts b/src/cliproxy/quota-fetcher.ts index e9e7fd81..41203461 100644 --- a/src/cliproxy/quota-fetcher.ts +++ b/src/cliproxy/quota-fetcher.ts @@ -16,6 +16,7 @@ import { type AccountInfo, type AccountTier, } from './account-manager'; +import { sanitizeEmail, isTokenExpired } from './auth-utils'; /** Individual model quota info */ export interface ModelQuota { @@ -154,27 +155,6 @@ interface FetchAvailableModelsResponse { models?: Record; } -/** - * Sanitize email to match CLIProxyAPI auth file naming convention - * Replaces @ and . with underscores (matches Go sanitizeAntigravityFileName) - */ -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; - } -} - /** * Refresh access token using refresh_token via Google OAuth * This allows CCS to get fresh tokens independently of CLIProxyAPI diff --git a/src/cliproxy/quota-types.ts b/src/cliproxy/quota-types.ts new file mode 100644 index 00000000..d888bf6c --- /dev/null +++ b/src/cliproxy/quota-types.ts @@ -0,0 +1,84 @@ +/** + * Shared Quota Type Definitions + * + * Unified types for multi-provider quota system. + * Supports Antigravity, Codex, and Gemini CLI providers. + */ + +/** 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; +} diff --git a/src/commands/cliproxy-command.ts b/src/commands/cliproxy-command.ts index 3e94af8d..d9cb5670 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,75 @@ 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); + } 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> +): 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 +1078,6 @@ async function handleQuotaStatus(verbose = false): Promise { ]); } - console.log(''); console.log( table(rows, { head: ['', 'Account', 'Tier', 'Quota', 'Status'], @@ -969,25 +1085,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 +1268,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; } diff --git a/tests/unit/cliproxy/auth-utils.test.ts b/tests/unit/cliproxy/auth-utils.test.ts new file mode 100644 index 00000000..a6003b80 --- /dev/null +++ b/tests/unit/cliproxy/auth-utils.test.ts @@ -0,0 +1,88 @@ +/** + * Auth Utilities Unit Tests + * + * Tests for shared authentication utility functions + */ + +import { describe, it, expect } from 'bun:test'; +import { sanitizeEmail, isTokenExpired } from '../../../src/cliproxy/auth-utils'; + +describe('Auth Utilities', () => { + describe('sanitizeEmail', () => { + it('should replace @ with underscore', () => { + const result = sanitizeEmail('user@example.com'); + expect(result).not.toContain('@'); + expect(result).toContain('user_example'); + }); + + it('should replace . with underscore', () => { + const result = sanitizeEmail('user@example.com'); + expect(result).not.toContain('.'); + expect(result).toBe('user_example_com'); + }); + + it('should handle multiple dots', () => { + const result = sanitizeEmail('user.name@sub.example.com'); + expect(result).toBe('user_name_sub_example_com'); + }); + + it('should handle email without dots in domain', () => { + const result = sanitizeEmail('user@localhost'); + expect(result).toBe('user_localhost'); + }); + + it('should handle empty string', () => { + const result = sanitizeEmail(''); + expect(result).toBe(''); + }); + }); + + describe('isTokenExpired', () => { + it('should return false for undefined input', () => { + const result = isTokenExpired(undefined); + expect(result).toBe(false); + }); + + it('should return false for empty string', () => { + const result = isTokenExpired(''); + expect(result).toBe(false); + }); + + it('should return true for past date', () => { + const pastDate = new Date(Date.now() - 3600000).toISOString(); // 1 hour ago + const result = isTokenExpired(pastDate); + expect(result).toBe(true); + }); + + it('should return false for future date', () => { + const futureDate = new Date(Date.now() + 3600000).toISOString(); // 1 hour from now + const result = isTokenExpired(futureDate); + expect(result).toBe(false); + }); + + it('should return false for invalid date string', () => { + // new Date('invalid') returns Invalid Date, getTime() returns NaN + // NaN < Date.now() is false + const result = isTokenExpired('not-a-date'); + expect(result).toBe(false); + }); + + it('should handle ISO date strings', () => { + const pastISO = '2020-01-01T00:00:00.000Z'; + expect(isTokenExpired(pastISO)).toBe(true); + + const futureISO = '2030-01-01T00:00:00.000Z'; + expect(isTokenExpired(futureISO)).toBe(false); + }); + + it('should handle Unix timestamp strings', () => { + // JavaScript Date can parse numeric strings as timestamps + const pastTimestamp = String(Date.now() - 86400000); // Yesterday + // Note: Date parsing of pure numbers as strings is inconsistent + // This test documents the actual behavior + const result = isTokenExpired(pastTimestamp); + // The behavior depends on how Date parses the string + expect(typeof result).toBe('boolean'); + }); + }); +}); diff --git a/tests/unit/cliproxy/quota-fetcher-codex.test.ts b/tests/unit/cliproxy/quota-fetcher-codex.test.ts new file mode 100644 index 00000000..5c774f4b --- /dev/null +++ b/tests/unit/cliproxy/quota-fetcher-codex.test.ts @@ -0,0 +1,190 @@ +/** + * Codex Quota Fetcher Unit Tests + * + * Tests for Codex quota window parsing and transformation logic + */ + +import { describe, it, expect } from 'bun:test'; +import { buildCodexQuotaWindows } from '../../../src/cliproxy/quota-fetcher-codex'; + +describe('Codex Quota Fetcher', () => { + describe('buildCodexQuotaWindows', () => { + it('should parse snake_case API response', () => { + const response = { + rate_limit: { + primary_window: { + used_percent: 25, + reset_after_seconds: 3600, + }, + secondary_window: { + used_percent: 50, + reset_after_seconds: 86400, + }, + }, + }; + + const windows = buildCodexQuotaWindows(response); + + expect(windows).toHaveLength(2); + expect(windows[0].label).toBe('Primary'); + expect(windows[0].usedPercent).toBe(25); + expect(windows[0].remainingPercent).toBe(75); + expect(windows[0].resetAfterSeconds).toBe(3600); + expect(windows[1].label).toBe('Secondary'); + expect(windows[1].usedPercent).toBe(50); + }); + + it('should parse camelCase API response', () => { + const response = { + rateLimit: { + primaryWindow: { + usedPercent: 30, + resetAfterSeconds: 7200, + }, + }, + }; + + const windows = buildCodexQuotaWindows(response); + + expect(windows).toHaveLength(1); + expect(windows[0].usedPercent).toBe(30); + expect(windows[0].resetAfterSeconds).toBe(7200); + }); + + it('should handle code review rate limits', () => { + const response = { + code_review_rate_limit: { + primary_window: { + used_percent: 80, + reset_after_seconds: 1800, + }, + }, + }; + + const windows = buildCodexQuotaWindows(response); + + expect(windows).toHaveLength(1); + expect(windows[0].label).toBe('Code Review (Primary)'); + expect(windows[0].usedPercent).toBe(80); + }); + + it('should clamp usedPercent to 0-100 range', () => { + const response = { + rate_limit: { + primary_window: { + used_percent: 150, // Over 100 + reset_after_seconds: null, + }, + secondary_window: { + used_percent: -20, // Negative + reset_after_seconds: null, + }, + }, + }; + + const windows = buildCodexQuotaWindows(response); + + expect(windows[0].usedPercent).toBe(100); + expect(windows[0].remainingPercent).toBe(0); + expect(windows[1].usedPercent).toBe(0); + expect(windows[1].remainingPercent).toBe(100); + }); + + it('should handle null reset_after_seconds', () => { + const response = { + rate_limit: { + primary_window: { + used_percent: 10, + reset_after_seconds: null, + }, + }, + }; + + const windows = buildCodexQuotaWindows(response); + + expect(windows[0].resetAfterSeconds).toBeNull(); + expect(windows[0].resetAt).toBeNull(); + }); + + it('should calculate resetAt from positive seconds', () => { + const response = { + rate_limit: { + primary_window: { + used_percent: 10, + reset_after_seconds: 3600, // 1 hour + }, + }, + }; + + const before = Date.now(); + const windows = buildCodexQuotaWindows(response); + const after = Date.now(); + + expect(windows[0].resetAt).not.toBeNull(); + const resetTime = new Date(windows[0].resetAt!).getTime(); + expect(resetTime).toBeGreaterThanOrEqual(before + 3600000); + expect(resetTime).toBeLessThanOrEqual(after + 3600000); + }); + + it('should not calculate resetAt for zero or negative seconds', () => { + const response = { + rate_limit: { + primary_window: { + used_percent: 10, + reset_after_seconds: 0, + }, + secondary_window: { + used_percent: 20, + reset_after_seconds: -100, + }, + }, + }; + + const windows = buildCodexQuotaWindows(response); + + expect(windows[0].resetAt).toBeNull(); + expect(windows[1].resetAt).toBeNull(); + }); + + it('should return empty array for empty response', () => { + const windows = buildCodexQuotaWindows({}); + expect(windows).toHaveLength(0); + }); + + it('should return empty array for missing rate limit', () => { + const response = { + plan_type: 'plus', + }; + + const windows = buildCodexQuotaWindows(response); + expect(windows).toHaveLength(0); + }); + + it('should handle missing window data gracefully', () => { + const response = { + rate_limit: { + primary_window: undefined, + secondary_window: null, + }, + }; + + const windows = buildCodexQuotaWindows(response as never); + expect(windows).toHaveLength(0); + }); + + it('should default usedPercent to 0 when missing', () => { + const response = { + rate_limit: { + primary_window: { + reset_after_seconds: 3600, + }, + }, + }; + + const windows = buildCodexQuotaWindows(response); + + expect(windows[0].usedPercent).toBe(0); + expect(windows[0].remainingPercent).toBe(100); + }); + }); +}); diff --git a/tests/unit/cliproxy/quota-fetcher-gemini-cli.test.ts b/tests/unit/cliproxy/quota-fetcher-gemini-cli.test.ts new file mode 100644 index 00000000..2a9e6b6f --- /dev/null +++ b/tests/unit/cliproxy/quota-fetcher-gemini-cli.test.ts @@ -0,0 +1,196 @@ +/** + * Gemini CLI Quota Fetcher Unit Tests + * + * Tests for Gemini CLI bucket parsing and transformation logic + */ + +import { describe, it, expect } from 'bun:test'; +import { + buildGeminiCliBuckets, + resolveGeminiCliProjectId, +} from '../../../src/cliproxy/quota-fetcher-gemini-cli'; + +describe('Gemini CLI Quota Fetcher', () => { + describe('resolveGeminiCliProjectId', () => { + it('should extract project ID from account field', () => { + const account = 'user@example.com (cloudaicompanion-abc-123)'; + const projectId = resolveGeminiCliProjectId(account); + expect(projectId).toBe('cloudaicompanion-abc-123'); + }); + + it('should return last parenthetical when multiple exist', () => { + const account = 'user (org) (cloudaicompanion-xyz-789)'; + const projectId = resolveGeminiCliProjectId(account); + expect(projectId).toBe('cloudaicompanion-xyz-789'); + }); + + it('should return null for account without parentheses', () => { + const account = 'user@example.com'; + const projectId = resolveGeminiCliProjectId(account); + expect(projectId).toBeNull(); + }); + + it('should return null for empty string', () => { + const projectId = resolveGeminiCliProjectId(''); + expect(projectId).toBeNull(); + }); + + it('should handle nested parentheses', () => { + const account = 'user@example.com (project-id)'; + const projectId = resolveGeminiCliProjectId(account); + expect(projectId).toBe('project-id'); + }); + }); + + describe('buildGeminiCliBuckets', () => { + it('should group models by series', () => { + const rawBuckets = [ + { model_id: 'gemini-3-flash-preview', remaining_fraction: 0.8 }, + { model_id: 'gemini-2.5-flash', remaining_fraction: 0.6 }, + { model_id: 'gemini-3-pro-preview', remaining_fraction: 0.9 }, + ]; + + const buckets = buildGeminiCliBuckets(rawBuckets); + + // Should have 2 groups: Flash Series and Pro Series + expect(buckets.length).toBeGreaterThanOrEqual(2); + + 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); + + const proBucket = buckets.find((b) => b.label === 'Gemini Pro Series'); + expect(proBucket).toBeDefined(); + expect(proBucket!.remainingFraction).toBe(0.9); + }); + + it('should handle camelCase API response', () => { + const rawBuckets = [ + { modelId: 'gemini-3-flash-preview', remainingFraction: 0.75 }, + ]; + + const buckets = buildGeminiCliBuckets(rawBuckets); + + expect(buckets).toHaveLength(1); + expect(buckets[0].remainingFraction).toBe(0.75); + }); + + it('should clamp remainingFraction to 0-1 range', () => { + const rawBuckets = [ + { model_id: 'gemini-3-flash-preview', remaining_fraction: 1.5 }, + { model_id: 'gemini-3-pro-preview', remaining_fraction: -0.2 }, + ]; + + const buckets = buildGeminiCliBuckets(rawBuckets); + + const flashBucket = buckets.find((b) => b.label === 'Gemini Flash Series'); + expect(flashBucket!.remainingFraction).toBe(1); + expect(flashBucket!.remainingPercent).toBe(100); + + const proBucket = buckets.find((b) => b.label === 'Gemini Pro Series'); + expect(proBucket!.remainingFraction).toBe(0); + expect(proBucket!.remainingPercent).toBe(0); + }); + + it('should group by token type', () => { + const rawBuckets = [ + { model_id: 'gemini-3-flash-preview', token_type: 'input', remaining_fraction: 0.8 }, + { model_id: 'gemini-3-flash-preview', token_type: 'output', remaining_fraction: 0.5 }, + ]; + + const buckets = buildGeminiCliBuckets(rawBuckets); + + // Should have separate buckets for input and output + expect(buckets.length).toBe(2); + const inputBucket = buckets.find((b) => b.tokenType === 'input'); + const outputBucket = buckets.find((b) => b.tokenType === 'output'); + expect(inputBucket).toBeDefined(); + expect(outputBucket).toBeDefined(); + expect(inputBucket!.remainingFraction).toBe(0.8); + expect(outputBucket!.remainingFraction).toBe(0.5); + }); + + it('should ignore deprecated models', () => { + const rawBuckets = [ + { model_id: 'gemini-2.0-flash-deprecated', remaining_fraction: 0.1 }, + { model_id: 'gemini-3-flash-preview', remaining_fraction: 0.9 }, + ]; + + const buckets = buildGeminiCliBuckets(rawBuckets); + + // Only gemini-3-flash-preview should be included + expect(buckets).toHaveLength(1); + expect(buckets[0].remainingFraction).toBe(0.9); + }); + + it('should categorize unknown models as "other"', () => { + 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'); + }); + + it('should handle empty buckets array', () => { + const buckets = buildGeminiCliBuckets([]); + expect(buckets).toHaveLength(0); + }); + + it('should skip buckets with empty model_id', () => { + const rawBuckets = [ + { model_id: '', remaining_fraction: 0.5 }, + { model_id: 'gemini-3-flash-preview', remaining_fraction: 0.8 }, + ]; + + const buckets = buildGeminiCliBuckets(rawBuckets); + + expect(buckets).toHaveLength(1); + expect(buckets[0].remainingFraction).toBe(0.8); + }); + + it('should keep earliest reset time when merging', () => { + const rawBuckets = [ + { + model_id: 'gemini-3-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', // Earlier + }, + ]; + + const buckets = buildGeminiCliBuckets(rawBuckets); + + const flashBucket = buckets.find((b) => b.label === 'Gemini Flash Series'); + expect(flashBucket!.resetTime).toBe('2026-01-30T10:00:00Z'); + }); + + it('should default remainingFraction to 1 when missing', () => { + const rawBuckets = [{ model_id: 'gemini-3-flash-preview' }]; + + const buckets = buildGeminiCliBuckets(rawBuckets); + + expect(buckets[0].remainingFraction).toBe(1); + expect(buckets[0].remainingPercent).toBe(100); + }); + + it('should collect modelIds in bucket', () => { + const rawBuckets = [ + { model_id: 'gemini-3-flash-preview', remaining_fraction: 0.8 }, + { model_id: 'gemini-2.5-flash', remaining_fraction: 0.6 }, + ]; + + const buckets = buildGeminiCliBuckets(rawBuckets); + + const flashBucket = buckets.find((b) => b.label === 'Gemini Flash Series'); + expect(flashBucket!.modelIds).toContain('gemini-3-flash-preview'); + expect(flashBucket!.modelIds).toContain('gemini-2.5-flash'); + }); + }); +});