diff --git a/src/cliproxy/account-manager.ts b/src/cliproxy/account-manager.ts index a65947c9..fc42bc12 100644 --- a/src/cliproxy/account-manager.ts +++ b/src/cliproxy/account-manager.ts @@ -14,6 +14,9 @@ import { CLIProxyProvider } from './types'; import { getCliproxyDir, getAuthDir } from './config-generator'; import { PROVIDER_TYPE_VALUES } from './auth/auth-types'; +/** Account tier for quota management */ +export type AccountTier = 'free' | 'pro' | 'ultra' | 'unknown'; + /** Account information */ export interface AccountInfo { /** Account identifier (email or custom name) */ @@ -32,6 +35,12 @@ export interface AccountInfo { createdAt: string; /** Last usage time */ lastUsedAt?: string; + /** User-paused state (skip in quota rotation) */ + paused?: boolean; + /** ISO timestamp when paused */ + pausedAt?: string; + /** Account tier: free, pro, ultra */ + tier?: AccountTier; } /** Provider accounts configuration */ @@ -330,6 +339,76 @@ export function setDefaultAccount(provider: CLIProxyProvider, accountId: string) return true; } +/** + * Pause an account (skip in quota rotation) + */ +export function pauseAccount(provider: CLIProxyProvider, accountId: string): boolean { + const registry = loadAccountsRegistry(); + const providerAccounts = registry.providers[provider]; + + if (!providerAccounts?.accounts[accountId]) { + return false; + } + + providerAccounts.accounts[accountId].paused = true; + providerAccounts.accounts[accountId].pausedAt = new Date().toISOString(); + saveAccountsRegistry(registry); + return true; +} + +/** + * Resume a paused account + */ +export function resumeAccount(provider: CLIProxyProvider, accountId: string): boolean { + const registry = loadAccountsRegistry(); + const providerAccounts = registry.providers[provider]; + + if (!providerAccounts?.accounts[accountId]) { + return false; + } + + providerAccounts.accounts[accountId].paused = false; + providerAccounts.accounts[accountId].pausedAt = undefined; + saveAccountsRegistry(registry); + return true; +} + +/** + * Check if an account is paused + */ +export function isAccountPaused(provider: CLIProxyProvider, accountId: string): boolean { + const accounts = getProviderAccounts(provider); + const account = accounts.find((a) => a.id === accountId); + return account?.paused ?? false; +} + +/** + * Update account tier + */ +export function setAccountTier( + provider: CLIProxyProvider, + accountId: string, + tier: AccountTier +): boolean { + const registry = loadAccountsRegistry(); + const providerAccounts = registry.providers[provider]; + + if (!providerAccounts?.accounts[accountId]) { + return false; + } + + providerAccounts.accounts[accountId].tier = tier; + saveAccountsRegistry(registry); + return true; +} + +/** + * Get non-paused accounts for a provider + */ +export function getActiveAccounts(provider: CLIProxyProvider): AccountInfo[] { + return getProviderAccounts(provider).filter((a) => !a.paused); +} + /** * Remove an account */ diff --git a/src/cliproxy/quota-fetcher.ts b/src/cliproxy/quota-fetcher.ts index d2daa236..6748601d 100644 --- a/src/cliproxy/quota-fetcher.ts +++ b/src/cliproxy/quota-fetcher.ts @@ -9,7 +9,12 @@ import * as fs from 'node:fs'; import * as path from 'node:path'; import { getAuthDir } from './config-generator'; import { CLIProxyProvider } from './types'; -import { getProviderAccounts, type AccountInfo } from './account-manager'; +import { + getProviderAccounts, + setAccountTier, + type AccountInfo, + type AccountTier, +} from './account-manager'; /** Individual model quota info */ export interface ModelQuota { @@ -45,6 +50,8 @@ export interface QuotaResult { accountId?: string; /** GCP project ID for this account */ projectId?: string; + /** Detected account tier based on model access */ + tier?: AccountTier; } /** Google Cloud Code API endpoints */ @@ -104,6 +111,38 @@ interface TokenRefreshResponse { error_description?: string; } +/** + * Detect account tier from quota API model list. + * Ultra accounts have access to experimental/preview models. + * Pro accounts have standard model access. + * Free/unknown accounts have limited model access. + */ +export function detectTier(quotaResult: QuotaResult): AccountTier { + if (!quotaResult.success || quotaResult.models.length === 0) { + return 'unknown'; + } + + const modelNames = quotaResult.models.map((m) => m.name.toLowerCase()); + + // Ultra indicators: experimental, preview, ultra in name + const ultraIndicators = ['ultra', 'experimental', 'preview', '2.5-pro', '3-ultra']; + const hasUltra = modelNames.some((name) => + ultraIndicators.some((indicator) => name.includes(indicator)) + ); + + if (hasUltra) return 'ultra'; + + // Pro indicators: gemini-2.0, gemini-3, pro models + const proIndicators = ['gemini-2', 'gemini-3', '-pro']; + const hasPro = modelNames.some((name) => + proIndicators.some((indicator) => name.includes(indicator)) + ); + + if (hasPro) return 'pro'; + + return 'free'; +} + /** loadCodeAssist response */ interface LoadCodeAssistResponse { cloudaicompanionProject?: string | { id?: string }; @@ -522,10 +561,26 @@ export async function fetchAccountQuota( if (!result.success && result.error?.includes('expired') && authData.refreshToken) { const refreshResult = await refreshAccessToken(authData.refreshToken); if (refreshResult.accessToken) { - return fetchAvailableModels(refreshResult.accessToken, projectId); + const retryResult = await fetchAvailableModels(refreshResult.accessToken, projectId); + // Detect and persist tier for retry result + if (retryResult.success) { + const tier = detectTier(retryResult); + retryResult.tier = tier; + retryResult.accountId = accountId; + setAccountTier(provider, accountId, tier); + } + return retryResult; } } + // Detect and persist tier for successful result + if (result.success) { + const tier = detectTier(result); + result.tier = tier; + result.accountId = accountId; + setAccountTier(provider, accountId, tier); + } + return result; } diff --git a/src/cliproxy/quota-manager.ts b/src/cliproxy/quota-manager.ts new file mode 100644 index 00000000..69a60005 --- /dev/null +++ b/src/cliproxy/quota-manager.ts @@ -0,0 +1,379 @@ +/** + * Quota Manager for Hybrid Auto+Manual Account Selection + * + * Provides pre-flight quota checking with caching, tier-based failover, + * and cooldown tracking for exhausted accounts. + * + * Key features: + * - 30-second in-memory cache for quota results + * - Tier-priority failover (ultra > pro by default) + * - Cooldown tracking for exhausted accounts + * - Respects paused accounts from manual config + * - Graceful degradation on API failures + */ + +import { CLIProxyProvider } from './types'; +import { QuotaResult, fetchAccountQuota } from './quota-fetcher'; +import { + getDefaultAccount, + getProviderAccounts, + isAccountPaused, + setDefaultAccount, + touchAccount, + type AccountInfo, +} from './account-manager'; +import { loadOrCreateUnifiedConfig } from '../config/unified-config-loader'; + +// ============================================================================ +// QUOTA CACHE (30-second TTL) +// ============================================================================ + +interface CacheEntry { + result: QuotaResult; + timestamp: number; +} + +const CACHE_TTL_MS = 30_000; // 30 seconds +const quotaCache = new Map(); + +function getCacheKey(provider: CLIProxyProvider, accountId: string): string { + return `${provider}:${accountId}`; +} + +/** + * Get cached quota result if still valid + */ +export function getCachedQuota(provider: CLIProxyProvider, accountId: string): QuotaResult | null { + const key = getCacheKey(provider, accountId); + const entry = quotaCache.get(key); + + if (!entry) return null; + + if (Date.now() - entry.timestamp > CACHE_TTL_MS) { + quotaCache.delete(key); + return null; + } + + return entry.result; +} + +/** + * Cache quota result + */ +export function setCachedQuota( + provider: CLIProxyProvider, + accountId: string, + result: QuotaResult +): void { + const key = getCacheKey(provider, accountId); + quotaCache.set(key, { result, timestamp: Date.now() }); +} + +/** + * Clear all cached quota results + */ +export function clearQuotaCache(): void { + quotaCache.clear(); +} + +// ============================================================================ +// COOLDOWN TRACKING +// ============================================================================ + +interface CooldownEntry { + until: number; // timestamp when cooldown expires +} + +const cooldownMap = new Map(); + +/** + * Check if account is on cooldown + */ +export function isOnCooldown(provider: CLIProxyProvider, accountId: string): boolean { + const key = getCacheKey(provider, accountId); + const entry = cooldownMap.get(key); + + if (!entry) return false; + + if (Date.now() > entry.until) { + cooldownMap.delete(key); + return false; + } + + return true; +} + +/** + * Apply cooldown to an exhausted account + */ +export function applyCooldown( + provider: CLIProxyProvider, + accountId: string, + minutes: number +): void { + const key = getCacheKey(provider, accountId); + cooldownMap.set(key, { until: Date.now() + minutes * 60 * 1000 }); +} + +/** + * Clear cooldown for an account + */ +export function clearCooldown(provider: CLIProxyProvider, accountId: string): void { + const key = getCacheKey(provider, accountId); + cooldownMap.delete(key); +} + +// ============================================================================ +// PRE-FLIGHT CHECK +// ============================================================================ + +/** + * Result of pre-flight quota check + */ +export interface PreflightResult { + /** Whether to proceed with session */ + proceed: boolean; + /** Account to use (may differ from original default) */ + accountId: string; + /** If switched, the original account ID */ + switchedFrom?: string; + /** Reason for switch or failure */ + reason?: string; + /** Average quota percentage of selected account */ + quotaPercent?: number; +} + +/** + * Calculate average quota percentage from models + */ +function calculateAverageQuota(quota: QuotaResult): number { + if (!quota.success || quota.models.length === 0) { + return 100; // Assume OK if no data + } + const total = quota.models.reduce((sum, m) => sum + m.percentage, 0); + return total / quota.models.length; +} + +/** + * Find healthy account with remaining quota + * Respects tier priority and skips paused/cooldown accounts + */ +export async function findHealthyAccount( + provider: CLIProxyProvider, + exclude: string[] +): Promise<{ id: string; tier: string; lastQuota: number } | null> { + const config = loadOrCreateUnifiedConfig(); + const tierPriority = config.quota_management?.auto?.tier_priority ?? ['ultra', 'pro']; + const threshold = config.quota_management?.auto?.exhaustion_threshold ?? 5; + + const accounts = getProviderAccounts(provider); + + // Filter available accounts + const available = accounts.filter( + (a) => + !exclude.includes(a.id) && !isAccountPaused(provider, a.id) && !isOnCooldown(provider, a.id) + ); + + if (available.length === 0) return null; + + // Fetch quota for each available account (with caching) + const withQuotas = await Promise.all( + available.map(async (account) => { + let quota = getCachedQuota(provider, account.id); + if (!quota) { + try { + quota = await fetchAccountQuota(provider, account.id); + setCachedQuota(provider, account.id, quota); + } catch { + quota = { success: false, models: [], lastUpdated: Date.now() }; + } + } + + const avgQuota = calculateAverageQuota(quota); + + return { + id: account.id, + tier: account.tier || 'pro', + lastQuota: avgQuota, + }; + }) + ); + + // Filter by threshold + const healthy = withQuotas.filter((a) => a.lastQuota >= threshold); + if (healthy.length === 0) return null; + + // Sort by tier priority then quota descending + healthy.sort((a, b) => { + const tierA = tierPriority.indexOf(a.tier); + const tierB = tierPriority.indexOf(b.tier); + const tierOrderA = tierA === -1 ? 999 : tierA; + const tierOrderB = tierB === -1 ? 999 : tierB; + + if (tierOrderA !== tierOrderB) return tierOrderA - tierOrderB; + return b.lastQuota - a.lastQuota; + }); + + return healthy[0]; +} + +/** + * Find and switch to a healthy account + */ +async function findAndSwitch( + provider: CLIProxyProvider, + excludeAccountId: string, + reason: string +): Promise { + const alternative = await findHealthyAccount(provider, [excludeAccountId]); + + if (!alternative) { + // No alternatives: use original anyway (graceful degradation) + return { + proceed: true, + accountId: excludeAccountId, + reason: `${reason}, no alternatives available`, + }; + } + + // Switch default + setDefaultAccount(provider, alternative.id); + touchAccount(provider, alternative.id); + + return { + proceed: true, + accountId: alternative.id, + switchedFrom: excludeAccountId, + reason, + quotaPercent: alternative.lastQuota, + }; +} + +/** + * Perform pre-flight quota check before session start + * + * Checks if default account has sufficient quota, auto-switches if needed. + * Respects paused accounts, tier priority, and cooldown settings. + * + * @param provider - CLIProxy provider (only 'agy' supports quota) + * @returns PreflightResult with account to use and any switch info + */ +export async function preflightCheck(provider: CLIProxyProvider): Promise { + // Only Antigravity supports quota checking + if (provider !== 'agy') { + const defaultAccount = getDefaultAccount(provider); + return { proceed: true, accountId: defaultAccount?.id || '' }; + } + + const config = loadOrCreateUnifiedConfig(); + const quotaConfig = config.quota_management; + + // Skip if preflight disabled or mode is manual + if (!quotaConfig?.auto?.preflight_check || quotaConfig?.mode === 'manual') { + const defaultAccount = getDefaultAccount(provider); + return { proceed: true, accountId: defaultAccount?.id || '' }; + } + + const defaultAccount = getDefaultAccount(provider); + if (!defaultAccount) { + return { proceed: false, accountId: '', reason: 'No accounts configured' }; + } + + // Check forced_default override (manual mode) + const forcedDefault = quotaConfig.manual?.forced_default; + if (forcedDefault) { + const forcedAccount = getProviderAccounts(provider).find((a) => a.id === forcedDefault); + if (forcedAccount) { + return { proceed: true, accountId: forcedAccount.id, reason: 'Forced default override' }; + } + } + + // Check if default is paused + if (isAccountPaused(provider, defaultAccount.id)) { + return await findAndSwitch(provider, defaultAccount.id, 'Default account is paused'); + } + + // Check cooldown + if (isOnCooldown(provider, defaultAccount.id)) { + return await findAndSwitch(provider, defaultAccount.id, 'Default account on cooldown'); + } + + // Check quota (with cache) + let quota = getCachedQuota(provider, defaultAccount.id); + if (!quota) { + try { + quota = await fetchAccountQuota(provider, defaultAccount.id); + setCachedQuota(provider, defaultAccount.id, quota); + } catch { + // API failure: proceed anyway (graceful degradation) + return { + proceed: true, + accountId: defaultAccount.id, + reason: 'Quota check failed, proceeding', + }; + } + } + + // Calculate average quota + const avgQuota = calculateAverageQuota(quota); + const threshold = quotaConfig.auto?.exhaustion_threshold ?? 5; + + if (avgQuota < threshold) { + // Apply cooldown to exhausted account + applyCooldown(provider, defaultAccount.id, quotaConfig.auto?.cooldown_minutes ?? 5); + return await findAndSwitch( + provider, + defaultAccount.id, + `Quota exhausted (${avgQuota.toFixed(1)}%)` + ); + } + + return { + proceed: true, + accountId: defaultAccount.id, + quotaPercent: avgQuota, + }; +} + +/** + * Get quota status for all accounts of a provider + * Used by CLI status command + */ +export async function getQuotaStatus(provider: CLIProxyProvider): Promise<{ + accounts: Array<{ + account: AccountInfo; + quota: number; + paused: boolean; + onCooldown: boolean; + isDefault: boolean; + }>; +}> { + const accounts = getProviderAccounts(provider); + const defaultAccount = getDefaultAccount(provider); + + const results = await Promise.all( + accounts.map(async (account) => { + let quota = getCachedQuota(provider, account.id); + if (!quota && provider === 'agy') { + try { + quota = await fetchAccountQuota(provider, account.id); + setCachedQuota(provider, account.id, quota); + } catch { + quota = { success: false, models: [], lastUpdated: Date.now() }; + } + } + + const avgQuota = quota ? calculateAverageQuota(quota) : 100; + + return { + account, + quota: avgQuota, + paused: isAccountPaused(provider, account.id), + onCooldown: isOnCooldown(provider, account.id), + isDefault: defaultAccount?.id === account.id, + }; + }) + ); + + return { accounts: results }; +} diff --git a/src/config/unified-config-types.ts b/src/config/unified-config-types.ts index 6a20a4d4..ac17e17d 100644 --- a/src/config/unified-config-types.ts +++ b/src/config/unified-config-types.ts @@ -16,8 +16,9 @@ * Version 4 = Copilot API integration (GitHub Copilot proxy) * Version 5 = Remote proxy configuration (connect to remote CLIProxyAPI) * Version 6 = Customizable auth tokens (API key and management secret) + * Version 7 = Quota management for hybrid auto+manual account control */ -export const UNIFIED_CONFIG_VERSION = 6; +export const UNIFIED_CONFIG_VERSION = 7; /** * Account configuration (formerly in profiles.json). @@ -342,12 +343,93 @@ export interface WebSearchConfig { customMcp?: unknown[]; } +// ============================================================================ +// QUOTA MANAGEMENT CONFIGURATION (v7+) +// ============================================================================ + +/** + * Auto quota management configuration. + * Controls automatic failover behavior. + */ +export interface AutoQuotaConfig { + /** Enable pre-flight quota check before requests (default: true) */ + preflight_check: boolean; + /** Quota percentage below which account is "exhausted" (default: 5) */ + exhaustion_threshold: number; + /** Tier priority for failover, highest to lowest (default: ['ultra', 'pro']) */ + tier_priority: string[]; + /** Minutes to skip exhausted account before retry (default: 5) */ + cooldown_minutes: number; +} + +/** + * Manual quota management configuration. + * User-controlled overrides for account selection. + */ +export interface ManualQuotaConfig { + /** User-paused accounts (stored in accounts.json) */ + paused_accounts: string[]; + /** Force use of specific account (overrides auto-selection) */ + forced_default: string | null; + /** Lock to specific tier only */ + tier_lock: string | null; +} + +/** + * Quota management mode. + * - auto: Fully automatic failover based on quota + * - manual: User controls everything, no auto-switching + * - hybrid: Auto-failover with user overrides (default) + */ +export type QuotaManagementMode = 'auto' | 'manual' | 'hybrid'; + +/** + * Quota management configuration section. + * Controls hybrid auto+manual account selection for multi-account setups. + */ +export interface QuotaManagementConfig { + /** Management mode (default: hybrid) */ + mode: QuotaManagementMode; + /** Auto mode settings */ + auto: AutoQuotaConfig; + /** Manual mode settings */ + manual: ManualQuotaConfig; +} + +/** + * Default auto quota configuration. + */ +export const DEFAULT_AUTO_QUOTA_CONFIG: AutoQuotaConfig = { + preflight_check: true, + exhaustion_threshold: 5, + tier_priority: ['ultra', 'pro'], + cooldown_minutes: 5, +}; + +/** + * Default manual quota configuration. + */ +export const DEFAULT_MANUAL_QUOTA_CONFIG: ManualQuotaConfig = { + paused_accounts: [], + forced_default: null, + tier_lock: null, +}; + +/** + * Default quota management configuration. + */ +export const DEFAULT_QUOTA_MANAGEMENT_CONFIG: QuotaManagementConfig = { + mode: 'hybrid', + auto: { ...DEFAULT_AUTO_QUOTA_CONFIG }, + manual: { ...DEFAULT_MANUAL_QUOTA_CONFIG }, +}; + /** * Main unified configuration structure. * Stored in ~/.ccs/config.yaml */ export interface UnifiedConfig { - /** Config version (5 for remote proxy support) */ + /** Config version (7 for quota management) */ version: number; /** Default profile name to use when none specified */ default?: string; @@ -367,6 +449,8 @@ export interface UnifiedConfig { copilot?: CopilotConfig; /** CLIProxy server configuration for remote/local mode */ cliproxy_server?: CliproxyServerConfig; + /** Quota management configuration (v7+) */ + quota_management?: QuotaManagementConfig; } /** @@ -455,6 +539,7 @@ export function createEmptyUnifiedConfig(): UnifiedConfig { }, copilot: { ...DEFAULT_COPILOT_CONFIG }, cliproxy_server: { ...DEFAULT_CLIPROXY_SERVER_CONFIG }, + quota_management: { ...DEFAULT_QUOTA_MANAGEMENT_CONFIG }, }; }