From 5ae19a5b72c8060008db2b84aa4a9c71f4ed4fd1 Mon Sep 17 00:00:00 2001 From: Badri Narayanan S Date: Sun, 18 Jan 2026 03:48:43 +0530 Subject: [PATCH] feat: add configurable account selection strategies Refactor account selection into a strategy pattern with three options: - Sticky: cache-optimized, stays on same account until rate-limited - Round-robin: load-balanced, rotates every request - Hybrid (default): smart distribution using health scores, token buckets, and LRU The hybrid strategy uses multiple signals for optimal account selection: health tracking for reliability, client-side token buckets for rate limiting, and LRU freshness to prefer rested accounts. Includes WebUI settings for strategy selection and unit tests. Co-Authored-By: Claude --- CLAUDE.md | 50 +- README.md | 37 +- package.json | 3 +- public/js/components/server-config.js | 97 +++ public/js/config/constants.js | 22 +- public/js/translations/en.js | 24 + public/js/translations/id.js | 25 + public/js/translations/pt.js | 26 + public/js/translations/tr.js | 26 + public/js/translations/zh.js | 25 + public/views/settings.html | 157 ++++ src/account-manager/index.js | 154 ++-- src/account-manager/selection.js | 201 ----- .../strategies/base-strategy.js | 104 +++ .../strategies/hybrid-strategy.js | 195 +++++ src/account-manager/strategies/index.js | 85 ++ .../strategies/round-robin-strategy.js | 76 ++ .../strategies/sticky-strategy.js | 138 +++ .../strategies/trackers/health-tracker.js | 162 ++++ .../strategies/trackers/index.js | 8 + .../trackers/token-bucket-tracker.js | 121 +++ src/cloudcode/message-handler.js | 200 ++++- src/cloudcode/streaming-handler.js | 196 ++++- src/config.js | 21 +- src/constants.js | 20 + src/errors.js | 37 +- src/index.js | 42 +- src/server.js | 14 +- src/webui/index.js | 12 +- tests/run-all.cjs | 1 + tests/test-strategies.cjs | 795 ++++++++++++++++++ 31 files changed, 2721 insertions(+), 353 deletions(-) delete mode 100644 src/account-manager/selection.js create mode 100644 src/account-manager/strategies/base-strategy.js create mode 100644 src/account-manager/strategies/hybrid-strategy.js create mode 100644 src/account-manager/strategies/index.js create mode 100644 src/account-manager/strategies/round-robin-strategy.js create mode 100644 src/account-manager/strategies/sticky-strategy.js create mode 100644 src/account-manager/strategies/trackers/health-tracker.js create mode 100644 src/account-manager/strategies/trackers/index.js create mode 100644 src/account-manager/strategies/trackers/token-bucket-tracker.js create mode 100644 tests/test-strategies.cjs diff --git a/CLAUDE.md b/CLAUDE.md index 8879e9d..96eb2d9 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -17,6 +17,11 @@ npm install # Start server (runs on port 8080) npm start +# Start with specific account selection strategy +npm start -- --strategy=sticky # Cache-optimized (stays on same account) +npm start -- --strategy=round-robin # Load-balanced (rotates every request) +npm start -- --strategy=hybrid # Smart distribution (default) + # Start with model fallback enabled (falls back to alternate model when quota exhausted) npm start -- --fallback @@ -50,6 +55,9 @@ npm run test:images # Image processing npm run test:caching # Prompt caching npm run test:crossmodel # Cross-model thinking signatures npm run test:oauth # OAuth no-browser mode + +# Run strategy unit tests (no server required) +node tests/test-strategies.cjs ``` ## Architecture @@ -83,9 +91,18 @@ src/ ├── account-manager/ # Multi-account pool management │ ├── index.js # AccountManager class facade │ ├── storage.js # Config file I/O and persistence -│ ├── selection.js # Account picking (round-robin, sticky) │ ├── rate-limits.js # Rate limit tracking and state -│ └── credentials.js # OAuth token and project handling +│ ├── credentials.js # OAuth token and project handling +│ └── strategies/ # Account selection strategies +│ ├── index.js # Strategy factory (createStrategy) +│ ├── base-strategy.js # Abstract base class +│ ├── sticky-strategy.js # Cache-optimized sticky selection +│ ├── round-robin-strategy.js # Load-balanced rotation +│ ├── hybrid-strategy.js # Smart multi-signal distribution +│ └── trackers/ # State trackers for hybrid strategy +│ ├── index.js # Re-exports trackers +│ ├── health-tracker.js # Account health scores +│ └── token-bucket-tracker.js # Client-side rate limiting │ ├── auth/ # Authentication │ ├── oauth.js # Google OAuth with PKCE @@ -161,7 +178,8 @@ public/ - **src/webui/index.js**: WebUI backend handling API routes (`/api/*`) for config, accounts, and logs - **src/cloudcode/**: Cloud Code API client with retry/failover logic, streaming and non-streaming support - `model-api.js`: Model listing, quota retrieval (`getModelQuotas()`), and subscription tier detection (`getSubscriptionTier()`) -- **src/account-manager/**: Multi-account pool with sticky selection, rate limit handling, and automatic cooldown +- **src/account-manager/**: Multi-account pool with configurable selection strategies, rate limit handling, and automatic cooldown + - Strategies: `sticky` (cache-optimized), `round-robin` (load-balanced), `hybrid` (smart distribution) - **src/auth/**: Authentication including Google OAuth, token extraction, database access, and auto-rebuild of native modules - **src/format/**: Format conversion between Anthropic and Google Generative AI formats - **src/constants.js**: API endpoints, model mappings, fallback config, OAuth config, and all configuration values @@ -170,12 +188,36 @@ public/ - **src/errors.js**: Custom error classes (`RateLimitError`, `AuthError`, `ApiError`, etc.) **Multi-Account Load Balancing:** -- Sticky account selection for prompt caching (stays on same account across turns) +- Configurable selection strategy via `--strategy` flag or WebUI +- Three strategies available: + - **Sticky** (`--strategy=sticky`): Best for prompt caching, stays on same account + - **Round-Robin** (`--strategy=round-robin`): Maximum throughput, rotates every request + - **Hybrid** (`--strategy=hybrid`, default): Smart selection using health + tokens + LRU - Model-specific rate limiting via `account.modelRateLimits[modelId]` - Automatic switch only when rate-limited for > 2 minutes on the current model - Session ID derived from first user message hash for cache continuity - Account state persisted to `~/.config/antigravity-proxy/accounts.json` +**Account Selection Strategies:** + +1. **Sticky Strategy** (best for caching): + - Stays on current account until rate-limited or unavailable + - Waits up to 2 minutes for short rate limits before switching + - Maintains prompt cache continuity across requests + +2. **Round-Robin Strategy** (best for throughput): + - Rotates to next account on every request + - Skips rate-limited/disabled accounts + - Maximizes concurrent request distribution + +3. **Hybrid Strategy** (default, smart distribution): + - Uses health scores, token buckets, and LRU for selection + - Scoring formula: `score = (Health × 2) + ((Tokens / MaxTokens × 100) × 5) + (LRU × 0.1)` + - Health scores: Track success/failure patterns with passive recovery + - Token buckets: Client-side rate limiting (50 tokens, 6 per minute regeneration) + - LRU freshness: Prefer accounts that have rested longer + - Configuration in `src/config.js` under `accountSelection` + **Account Data Model:** Each account object in `accounts.json` contains: - **Basic Info**: `email`, `source` (oauth/manual/database), `enabled`, `lastUsed` diff --git a/README.md b/README.md index 337923d..9bc2825 100644 --- a/README.md +++ b/README.md @@ -272,13 +272,37 @@ Gemini models include full thinking support with `thoughtSignature` handling for ## Multi-Account Load Balancing -When you add multiple accounts, the proxy automatically: +When you add multiple accounts, the proxy intelligently distributes requests across them using configurable selection strategies. -- **Sticky account selection**: Stays on the same account to maximize prompt cache hits -- **Smart rate limit handling**: Waits for short rate limits (≤2 min), switches accounts for longer ones -- **Automatic cooldown**: Rate-limited accounts become available after reset time expires -- **Invalid account detection**: Accounts needing re-authentication are marked and skipped -- **Prompt caching support**: Stable session IDs enable cache hits across conversation turns +### Account Selection Strategies + +Choose a strategy based on your needs: + +| Strategy | Best For | Description | +| --- | --- | --- | +| **Hybrid** (Default) | Most users | Smart selection combining health score, token bucket rate limiting, and LRU freshness | +| **Sticky** | Prompt caching | Stays on the same account to maximize cache hits, switches only when rate-limited | +| **Round-Robin** | Even distribution | Cycles through accounts sequentially for balanced load | + +**Configure via CLI:** + +```bash +antigravity-claude-proxy start --strategy=hybrid # Default: smart distribution +antigravity-claude-proxy start --strategy=sticky # Cache-optimized +antigravity-claude-proxy start --strategy=round-robin # Load-balanced +``` + +**Or via WebUI:** Settings → Server → Account Selection Strategy + +### How It Works + +- **Health Score Tracking**: Accounts earn points for successful requests and lose points for failures/rate-limits +- **Token Bucket Rate Limiting**: Client-side throttling with regenerating tokens (50 max, 6/minute) +- **Automatic Cooldown**: Rate-limited accounts recover automatically after reset time expires +- **Invalid Account Detection**: Accounts needing re-authentication are marked and skipped +- **Prompt Caching Support**: Session IDs derived from conversation enable cache hits across turns + +### Monitoring Check account status, subscription tiers, and quota anytime: @@ -395,6 +419,7 @@ npm run test:streaming # Streaming SSE events npm run test:interleaved # Interleaved thinking npm run test:images # Image processing npm run test:caching # Prompt caching +npm run test:strategies # Account selection strategies ``` --- diff --git a/package.json b/package.json index dec5b4f..d32b26a 100644 --- a/package.json +++ b/package.json @@ -34,7 +34,8 @@ "test:crossmodel": "node tests/test-cross-model-thinking.cjs", "test:oauth": "node tests/test-oauth-no-browser.cjs", "test:emptyretry": "node tests/test-empty-response-retry.cjs", - "test:sanitizer": "node tests/test-schema-sanitizer.cjs" + "test:sanitizer": "node tests/test-schema-sanitizer.cjs", + "test:strategies": "node tests/test-strategies.cjs" }, "keywords": [ "claude", diff --git a/public/js/components/server-config.js b/public/js/components/server-config.js index e5ca4ed..6b139e2 100644 --- a/public/js/components/server-config.js +++ b/public/js/components/server-config.js @@ -248,5 +248,102 @@ window.Components.serverConfig = () => ({ const { MAX_WAIT_MIN, MAX_WAIT_MAX } = window.AppConstants.VALIDATION; this.saveConfigField('maxWaitBeforeErrorMs', value, 'Max Wait Threshold', (v) => window.Validators.validateTimeout(v, MAX_WAIT_MIN, MAX_WAIT_MAX)); + }, + + toggleRateLimitDedupWindowMs(value) { + const { RATE_LIMIT_DEDUP_MIN, RATE_LIMIT_DEDUP_MAX } = window.AppConstants.VALIDATION; + this.saveConfigField('rateLimitDedupWindowMs', value, 'Rate Limit Dedup Window', + (v) => window.Validators.validateTimeout(v, RATE_LIMIT_DEDUP_MIN, RATE_LIMIT_DEDUP_MAX)); + }, + + toggleMaxConsecutiveFailures(value) { + const { MAX_CONSECUTIVE_FAILURES_MIN, MAX_CONSECUTIVE_FAILURES_MAX } = window.AppConstants.VALIDATION; + this.saveConfigField('maxConsecutiveFailures', value, 'Max Consecutive Failures', + (v) => window.Validators.validateRange(v, MAX_CONSECUTIVE_FAILURES_MIN, MAX_CONSECUTIVE_FAILURES_MAX, 'Max Consecutive Failures')); + }, + + toggleExtendedCooldownMs(value) { + const { EXTENDED_COOLDOWN_MIN, EXTENDED_COOLDOWN_MAX } = window.AppConstants.VALIDATION; + this.saveConfigField('extendedCooldownMs', value, 'Extended Cooldown', + (v) => window.Validators.validateTimeout(v, EXTENDED_COOLDOWN_MIN, EXTENDED_COOLDOWN_MAX)); + }, + + toggleCapacityRetryDelayMs(value) { + const { CAPACITY_RETRY_DELAY_MIN, CAPACITY_RETRY_DELAY_MAX } = window.AppConstants.VALIDATION; + this.saveConfigField('capacityRetryDelayMs', value, 'Capacity Retry Delay', + (v) => window.Validators.validateTimeout(v, CAPACITY_RETRY_DELAY_MIN, CAPACITY_RETRY_DELAY_MAX)); + }, + + toggleMaxCapacityRetries(value) { + const { MAX_CAPACITY_RETRIES_MIN, MAX_CAPACITY_RETRIES_MAX } = window.AppConstants.VALIDATION; + this.saveConfigField('maxCapacityRetries', value, 'Max Capacity Retries', + (v) => window.Validators.validateRange(v, MAX_CAPACITY_RETRIES_MIN, MAX_CAPACITY_RETRIES_MAX, 'Max Capacity Retries')); + }, + + // Toggle Account Selection Strategy + async toggleStrategy(strategy) { + const store = Alpine.store('global'); + const validStrategies = ['sticky', 'round-robin', 'hybrid']; + + if (!validStrategies.includes(strategy)) { + store.showToast(store.t('invalidStrategy'), 'error'); + return; + } + + // Optimistic update + const previousValue = this.serverConfig.accountSelection?.strategy || 'hybrid'; + if (!this.serverConfig.accountSelection) { + this.serverConfig.accountSelection = {}; + } + this.serverConfig.accountSelection.strategy = strategy; + + try { + const { response, newPassword } = await window.utils.request('/api/config', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ accountSelection: { strategy } }) + }, store.webuiPassword); + + if (newPassword) store.webuiPassword = newPassword; + + const data = await response.json(); + if (data.status === 'ok') { + const strategyLabel = this.getStrategyLabel(strategy); + store.showToast(store.t('strategyUpdated', { strategy: strategyLabel }), 'success'); + await this.fetchServerConfig(); // Confirm server state + } else { + throw new Error(data.error || store.t('failedToUpdateStrategy')); + } + } catch (e) { + // Rollback on error + if (!this.serverConfig.accountSelection) { + this.serverConfig.accountSelection = {}; + } + this.serverConfig.accountSelection.strategy = previousValue; + store.showToast(store.t('failedToUpdateStrategy') + ': ' + e.message, 'error'); + } + }, + + // Get display label for a strategy + getStrategyLabel(strategy) { + const store = Alpine.store('global'); + const labels = { + 'sticky': store.t('strategyStickyLabel'), + 'round-robin': store.t('strategyRoundRobinLabel'), + 'hybrid': store.t('strategyHybridLabel') + }; + return labels[strategy] || strategy; + }, + + // Get description for current strategy + currentStrategyDescription() { + const store = Alpine.store('global'); + const strategy = this.serverConfig.accountSelection?.strategy || 'hybrid'; + const descriptions = { + 'sticky': store.t('strategyStickyDesc'), + 'round-robin': store.t('strategyRoundRobinDesc'), + 'hybrid': store.t('strategyHybridDesc') + }; + return descriptions[strategy] || ''; } }); diff --git a/public/js/config/constants.js b/public/js/config/constants.js index a5a1b6c..7074faf 100644 --- a/public/js/config/constants.js +++ b/public/js/config/constants.js @@ -67,7 +67,27 @@ window.AppConstants.VALIDATION = { // Max wait threshold (1 - 30 minutes) MAX_WAIT_MIN: 60000, - MAX_WAIT_MAX: 1800000 + MAX_WAIT_MAX: 1800000, + + // Rate limit dedup window (1 - 30 seconds) + RATE_LIMIT_DEDUP_MIN: 1000, + RATE_LIMIT_DEDUP_MAX: 30000, + + // Consecutive failures (1 - 10) + MAX_CONSECUTIVE_FAILURES_MIN: 1, + MAX_CONSECUTIVE_FAILURES_MAX: 10, + + // Extended cooldown (10 seconds - 5 minutes) + EXTENDED_COOLDOWN_MIN: 10000, + EXTENDED_COOLDOWN_MAX: 300000, + + // Capacity retry delay (500ms - 10 seconds) + CAPACITY_RETRY_DELAY_MIN: 500, + CAPACITY_RETRY_DELAY_MAX: 10000, + + // Capacity retries (1 - 10) + MAX_CAPACITY_RETRIES_MIN: 1, + MAX_CAPACITY_RETRIES_MAX: 10 }; /** diff --git a/public/js/translations/en.js b/public/js/translations/en.js index 61942ca..5157122 100644 --- a/public/js/translations/en.js +++ b/public/js/translations/en.js @@ -237,6 +237,18 @@ window.translations.en = { defaultCooldownDesc: "Fallback cooldown when API doesn't provide a reset time.", maxWaitThreshold: "Max Wait Before Error", maxWaitDesc: "If all accounts are rate-limited longer than this, error immediately instead of waiting.", + // Error Handling Tuning + errorHandlingTuning: "Error Handling Tuning", + rateLimitDedupWindow: "Rate Limit Dedup Window", + rateLimitDedupWindowDesc: "Prevents concurrent retry storms when multiple requests hit rate limits simultaneously.", + maxConsecutiveFailures: "Max Consecutive Failures", + maxConsecutiveFailuresDesc: "Number of consecutive failures before applying extended cooldown to an account.", + extendedCooldown: "Extended Cooldown", + extendedCooldownDesc: "Cooldown duration applied after max consecutive failures reached.", + capacityRetryDelay: "Capacity Retry Delay", + capacityRetryDelayDesc: "Delay before retrying when model capacity is exhausted (not quota).", + maxCapacityRetries: "Max Capacity Retries", + maxCapacityRetriesDesc: "Maximum retries for capacity exhaustion before switching accounts.", saveConfigServer: "Save Configuration", serverRestartAlert: "Changes saved to {path}. Restart server to apply some settings.", changePassword: "Change WebUI Password", @@ -318,6 +330,18 @@ window.translations.en = { failedToUpdateModelConfig: "Failed to update model config", fieldUpdated: "{displayName} updated to {value}", failedToUpdateField: "Failed to update {displayName}", + // Account Selection Strategy + accountSelectionStrategy: "Account Selection Strategy", + selectionStrategy: "Selection Strategy", + strategyStickyLabel: "Sticky (Cache Optimized)", + strategyRoundRobinLabel: "Round Robin (Load Balanced)", + strategyHybridLabel: "Hybrid (Smart Distribution)", + strategyStickyDesc: "Stays on same account until rate-limited. Best for prompt caching.", + strategyRoundRobinDesc: "Rotates to next account on every request. Maximum throughput.", + strategyHybridDesc: "Smart selection based on health, tokens, and freshness.", + strategyUpdated: "Strategy updated to: {strategy}", + failedToUpdateStrategy: "Failed to update strategy", + invalidStrategy: "Invalid strategy selected", // Validation Messages mustBeValidNumber: "{fieldName} must be a valid number", mustBeAtLeast: "{fieldName} must be at least {min}", diff --git a/public/js/translations/id.js b/public/js/translations/id.js index 6e53bb1..917bc06 100644 --- a/public/js/translations/id.js +++ b/public/js/translations/id.js @@ -270,6 +270,18 @@ window.translations.id = { defaultCooldownDesc: "Cooldown bawaan jika API tidak memberikan waktu reset.", maxWaitThreshold: "Batas Tunggu Maksimal", maxWaitDesc: "Jika semua akun terkena rate limit lebih lama dari ini, langsung gagal.", + // Error Handling Tuning + errorHandlingTuning: "Penyetelan Penanganan Error", + rateLimitDedupWindow: "Jendela Deduplikasi Rate Limit", + rateLimitDedupWindowDesc: "Mencegah badai retry ketika beberapa permintaan terkena rate limit bersamaan.", + maxConsecutiveFailures: "Maks. Kegagalan Berturut-turut", + maxConsecutiveFailuresDesc: "Jumlah kegagalan berturut-turut sebelum menerapkan cooldown diperpanjang.", + extendedCooldown: "Cooldown Diperpanjang", + extendedCooldownDesc: "Durasi cooldown setelah mencapai maks. kegagalan berturut-turut.", + capacityRetryDelay: "Jeda Retry Kapasitas", + capacityRetryDelayDesc: "Jeda sebelum retry saat kapasitas model habis (bukan kuota).", + maxCapacityRetries: "Maks. Retry Kapasitas", + maxCapacityRetriesDesc: "Maksimum retry untuk kehabisan kapasitas sebelum ganti akun.", saveConfigServer: "Simpan Konfigurasi", serverRestartAlert: "Tersimpan ke {path}. Restart server untuk menerapkan.", @@ -368,4 +380,17 @@ window.translations.id = { mustBeAtMost: "{fieldName} maksimal {max}", cannotBeEmpty: "{fieldName} tidak boleh kosong", mustBeTrueOrFalse: "Nilai harus true atau false", + + // Account Selection Strategy translations + accountSelectionStrategy: "Strategi Pemilihan Akun", + selectionStrategy: "Strategi Pemilihan", + strategyStickyLabel: "Tetap (Optimisasi Cache)", + strategyRoundRobinLabel: "Bergilir (Load Balanced)", + strategyHybridLabel: "Hibrida (Distribusi Cerdas)", + strategyStickyDesc: "Tetap di akun yang sama hingga terkena rate limit. Terbaik untuk cache prompt.", + strategyRoundRobinDesc: "Berputar ke akun berikutnya setiap permintaan. Throughput maksimum.", + strategyHybridDesc: "Pemilihan cerdas berdasarkan kesehatan, token, dan kesegaran.", + strategyUpdated: "Strategi diubah ke: {strategy}", + failedToUpdateStrategy: "Gagal memperbarui strategi", + invalidStrategy: "Strategi tidak valid dipilih", }; diff --git a/public/js/translations/pt.js b/public/js/translations/pt.js index 7c6fb87..468d35c 100644 --- a/public/js/translations/pt.js +++ b/public/js/translations/pt.js @@ -212,8 +212,21 @@ window.translations.pt = { persistTokenDesc: "Salvar sessões OAuth no disco para reinicializações mais rápidas", rateLimiting: "Limitação de Taxa de Conta & Timeouts", defaultCooldown: "Tempo de Resfriamento Padrão", + defaultCooldownDesc: "Resfriamento de fallback quando a API não fornece tempo de reset.", maxWaitThreshold: "Limiar Máximo de Espera (Sticky)", maxWaitDesc: "Tempo máximo para aguardar uma conta sticky resetar antes de trocar.", + // Ajuste de Tratamento de Erros + errorHandlingTuning: "Ajuste de Tratamento de Erros", + rateLimitDedupWindow: "Janela de Deduplicação de Rate Limit", + rateLimitDedupWindowDesc: "Previne tempestades de retry quando múltiplas requisições atingem rate limits simultaneamente.", + maxConsecutiveFailures: "Máx. Falhas Consecutivas", + maxConsecutiveFailuresDesc: "Número de falhas consecutivas antes de aplicar resfriamento estendido.", + extendedCooldown: "Resfriamento Estendido", + extendedCooldownDesc: "Duração do resfriamento aplicado após atingir máx. de falhas consecutivas.", + capacityRetryDelay: "Atraso de Retry de Capacidade", + capacityRetryDelayDesc: "Atraso antes de tentar novamente quando capacidade do modelo está esgotada (não quota).", + maxCapacityRetries: "Máx. Retries de Capacidade", + maxCapacityRetriesDesc: "Máximo de retries para esgotamento de capacidade antes de trocar conta.", saveConfigServer: "Salvar Configuração", serverRestartAlert: "Alterações salvas em {path}. Reinicie o servidor para aplicar algumas configurações.", changePassword: "Alterar Senha da WebUI", @@ -258,4 +271,17 @@ window.translations.pt = { gemini1mDesc: "Adiciona sufixo [1m] aos modelos Gemini para suporte a janela de contexto de 1M.", gemini1mWarning: "⚠ Contexto grande pode reduzir o desempenho do Gemini-3-Pro.", clickToSet: "Clique para configurar...", + + // Account Selection Strategy translations + accountSelectionStrategy: "Estratégia de Seleção de Conta", + selectionStrategy: "Estratégia de Seleção", + strategyStickyLabel: "Fixo (Otimizado para Cache)", + strategyRoundRobinLabel: "Rodízio (Balanceamento de Carga)", + strategyHybridLabel: "Híbrido (Distribuição Inteligente)", + strategyStickyDesc: "Permanece na mesma conta até atingir limite. Melhor para cache de prompts.", + strategyRoundRobinDesc: "Alterna para próxima conta a cada requisição. Máximo throughput.", + strategyHybridDesc: "Seleção inteligente baseada em saúde, tokens e frescor.", + strategyUpdated: "Estratégia atualizada para: {strategy}", + failedToUpdateStrategy: "Falha ao atualizar estratégia", + invalidStrategy: "Estratégia inválida selecionada", }; diff --git a/public/js/translations/tr.js b/public/js/translations/tr.js index 746b083..ca68977 100644 --- a/public/js/translations/tr.js +++ b/public/js/translations/tr.js @@ -216,8 +216,21 @@ window.translations.tr = { persistTokenDesc: "Daha hızlı yeniden başlatmalar için OAuth oturumlarını diske kaydet", rateLimiting: "Hesap Hız Sınırlama ve Zaman Aşımları", defaultCooldown: "Varsayılan Soğuma Süresi", + defaultCooldownDesc: "API sıfırlama zamanı sağlamadığında yedek soğuma süresi.", maxWaitThreshold: "Maksimum Bekleme Eşiği (Yapışkan)", maxWaitDesc: "Yapışkan bir hesabın değiştirmeden önce sıfırlanması için beklenecek maksimum süre.", + // Hata İşleme Ayarları + errorHandlingTuning: "Hata İşleme Ayarları", + rateLimitDedupWindow: "Hız Sınırı Tekilleştirme Penceresi", + rateLimitDedupWindowDesc: "Birden fazla istek aynı anda hız sınırına ulaştığında yeniden deneme fırtınasını önler.", + maxConsecutiveFailures: "Maks. Ardışık Başarısızlık", + maxConsecutiveFailuresDesc: "Uzatılmış soğuma uygulamadan önce ardışık başarısızlık sayısı.", + extendedCooldown: "Uzatılmış Soğuma", + extendedCooldownDesc: "Maks. ardışık başarısızlık sonrası uygulanan soğuma süresi.", + capacityRetryDelay: "Kapasite Yeniden Deneme Gecikmesi", + capacityRetryDelayDesc: "Model kapasitesi tükendiğinde (kota değil) yeniden denemeden önceki gecikme.", + maxCapacityRetries: "Maks. Kapasite Yeniden Denemesi", + maxCapacityRetriesDesc: "Hesap değiştirmeden önce kapasite tükenmesi için maksimum yeniden deneme.", saveConfigServer: "Yapılandırmayı Kaydet", serverRestartAlert: "Değişiklikler {path} konumuna kaydedildi. Bazı ayarları uygulamak için sunucuyu yeniden başlatın.", changePassword: "WebUI Parolasını Değiştir", @@ -313,4 +326,17 @@ window.translations.tr = { // TODO: Missing translation - Server config (exists in EN but missing here) // defaultCooldownDesc: "Fallback cooldown when API doesn't provide a reset time.", + + // Account Selection Strategy translations + accountSelectionStrategy: "Hesap Seçim Stratejisi", + selectionStrategy: "Seçim Stratejisi", + strategyStickyLabel: "Sabit (Önbellek Optimizasyonu)", + strategyRoundRobinLabel: "Döngüsel (Yük Dengeleme)", + strategyHybridLabel: "Hibrit (Akıllı Dağıtım)", + strategyStickyDesc: "Hız sınırına ulaşılana kadar aynı hesapta kalır. Önbellek için en iyisi.", + strategyRoundRobinDesc: "Her istekte bir sonraki hesaba geçer. Maksimum verimlilik.", + strategyHybridDesc: "Sağlık, token ve tazeliğe dayalı akıllı seçim.", + strategyUpdated: "Strateji şu şekilde güncellendi: {strategy}", + failedToUpdateStrategy: "Strateji güncellenemedi", + invalidStrategy: "Geçersiz strateji seçildi", }; diff --git a/public/js/translations/zh.js b/public/js/translations/zh.js index e441c22..a1e7ef2 100644 --- a/public/js/translations/zh.js +++ b/public/js/translations/zh.js @@ -237,6 +237,18 @@ window.translations.zh = { defaultCooldownDesc: "当 API 未提供重置时间时的备用冷却时间。", maxWaitThreshold: "最大等待阈值", maxWaitDesc: "如果所有账号的限流时间超过此阈值，立即返回错误而非等待。", + // 错误处理调优 + errorHandlingTuning: "错误处理调优", + rateLimitDedupWindow: "限流去重窗口", + rateLimitDedupWindowDesc: "当多个请求同时触发限流时，防止并发重试风暴。", + maxConsecutiveFailures: "最大连续失败次数", + maxConsecutiveFailuresDesc: "触发扩展冷却前允许的连续失败次数。", + extendedCooldown: "扩展冷却时间", + extendedCooldownDesc: "达到最大连续失败后应用的冷却时长。", + capacityRetryDelay: "容量重试延迟", + capacityRetryDelayDesc: "模型容量耗尽（非配额）时重试前的延迟。", + maxCapacityRetries: "最大容量重试次数", + maxCapacityRetriesDesc: "容量耗尽时在切换账号前的最大重试次数。", saveConfigServer: "保存配置", serverRestartAlert: "配置已保存至 {path}。部分更改可能需要重启服务器。", changePassword: "修改 WebUI 密码", @@ -329,4 +341,17 @@ window.translations.zh = { // mustBeAtMost: "{fieldName} must be at most {max}", // cannotBeEmpty: "{fieldName} cannot be empty", // mustBeTrueOrFalse: "Value must be true or false", + + // Account Selection Strategy translations + accountSelectionStrategy: "账户选择策略", + selectionStrategy: "选择策略", + strategyStickyLabel: "固定 (缓存优化)", + strategyRoundRobinLabel: "轮询 (负载均衡)", + strategyHybridLabel: "混合 (智能分配)", + strategyStickyDesc: "保持使用同一账户直到被限速。最适合提示词缓存。", + strategyRoundRobinDesc: "每次请求轮换到下一个账户。最大吞吐量。", + strategyHybridDesc: "基于健康度、令牌和新鲜度的智能选择。", + strategyUpdated: "策略已更新为: {strategy}", + failedToUpdateStrategy: "更新策略失败", + invalidStrategy: "选择了无效的策略", }; diff --git a/public/views/settings.html b/public/views/settings.html index 169a1bf..c14f5de 100644 --- a/public/views/settings.html +++ b/public/views/settings.html @@ -936,6 +936,35 @@ + +
+
+ Account Selection Strategy +
+
+ +
+
+
+ Selection Strategy + How accounts are selected for requests +
+ +
+
+
+
If all accounts are rate-limited longer than this, error immediately.

+ + +
+
+ Error Handling Tuning +
+ +
+ +
+ + +
+

Prevents concurrent retry storms.

+
+ +
+ +
+ + +
+

Failures before extended cooldown.

+
+ +
+ +
+ + +
+

Applied after max consecutive failures.

+
+ +
+ +
+ + +
+

Delay for capacity (not quota) issues.

+
+ +
+ +
+ + +
+

Retries before switching accounts.

+
+
diff --git a/src/account-manager/index.js b/src/account-manager/index.js index 8db296d..8231344 100644 --- a/src/account-manager/index.js +++ b/src/account-manager/index.js @@ -1,6 +1,6 @@ /** * Account Manager - * Manages multiple Antigravity accounts with sticky selection, + * Manages multiple Antigravity accounts with configurable selection strategies, * automatic failover, and smart cooldown for rate-limited accounts. */ @@ -23,13 +23,9 @@ import { clearProjectCache as clearProject, clearTokenCache as clearToken } from './credentials.js'; -import { - pickNext as selectNext, - getCurrentStickyAccount as getSticky, - shouldWaitForCurrentAccount as shouldWait, - pickStickyAccount as selectSticky -} from './selection.js'; +import { createStrategy, getStrategyLabel, DEFAULT_STRATEGY } from './strategies/index.js'; import { logger } from '../utils/logger.js'; +import { config } from '../config.js'; export class AccountManager { #accounts = []; @@ -37,19 +33,26 @@ export class AccountManager { #configPath; #settings = {}; #initialized = false; + #strategy = null; + #strategyName = DEFAULT_STRATEGY; // Per-account caches #tokenCache = new Map(); // email -> { token, extractedAt } #projectCache = new Map(); // email -> projectId - constructor(configPath = ACCOUNT_CONFIG_PATH) { + constructor(configPath = ACCOUNT_CONFIG_PATH, strategyName = null) { this.#configPath = configPath; + // Strategy name can be set at construction or later via initialize + if (strategyName) { + this.#strategyName = strategyName; + } } /** * Initialize the account manager by loading config + * @param {string} [strategyOverride] - Override strategy name (from CLI flag or env var) */ - async initialize() { + async initialize(strategyOverride = null) { if (this.#initialized) return; const { accounts, settings, activeIndex } = await loadAccounts(this.#configPath); @@ -66,6 +69,16 @@ export class AccountManager { this.#tokenCache = tokenCache; } + // Determine strategy: CLI override > env var > config file > default + const configStrategy = config?.accountSelection?.strategy; + const envStrategy = process.env.ACCOUNT_STRATEGY; + this.#strategyName = strategyOverride || envStrategy || configStrategy || this.#strategyName; + + // Create the strategy instance + const strategyConfig = config?.accountSelection || {}; + this.#strategy = createStrategy(this.#strategyName, strategyConfig); + logger.info(`[AccountManager] Using ${getStrategyLabel(this.#strategyName)} selection strategy`); + // Clear any expired rate limits this.clearExpiredLimits(); @@ -138,51 +151,88 @@ export class AccountManager { } /** - * Pick the next available account (fallback when current is unavailable). - * Sets activeIndex to the selected account's index. - * @param {string} [modelId] - Optional model ID - * @returns {Object|null} The next available account or null if none available - */ - pickNext(modelId = null) { - const { account, newIndex } = selectNext(this.#accounts, this.#currentIndex, () => this.saveToDisk(), modelId); - this.#currentIndex = newIndex; - return account; - } - - /** - * Get the current account without advancing the index (sticky selection). - * Used for cache continuity - sticks to the same account until rate-limited. - * @param {string} [modelId] - Optional model ID - * @returns {Object|null} The current account or null if unavailable/rate-limited - */ - getCurrentStickyAccount(modelId = null) { - const { account, newIndex } = getSticky(this.#accounts, this.#currentIndex, () => this.saveToDisk(), modelId); - this.#currentIndex = newIndex; - return account; - } - - /** - * Check if we should wait for the current account's rate limit to reset. - * Used for sticky account selection - wait if rate limit is short (≤ threshold). - * @param {string} [modelId] - Optional model ID - * @returns {{shouldWait: boolean, waitMs: number, account: Object|null}} - */ - shouldWaitForCurrentAccount(modelId = null) { - return shouldWait(this.#accounts, this.#currentIndex, modelId); - } - - /** - * Pick an account with sticky selection preference. - * Prefers the current account for cache continuity, only switches when: - * - Current account is rate-limited for > 2 minutes - * - Current account is invalid - * @param {string} [modelId] - Optional model ID + * Select an account using the configured strategy. + * This is the main method to use for account selection. + * @param {string} [modelId] - Model ID for the request + * @param {Object} [options] - Additional options + * @param {string} [options.sessionId] - Session ID for cache continuity * @returns {{account: Object|null, waitMs: number}} Account to use and optional wait time */ - pickStickyAccount(modelId = null) { - const { account, waitMs, newIndex } = selectSticky(this.#accounts, this.#currentIndex, () => this.saveToDisk(), modelId); - this.#currentIndex = newIndex; - return { account, waitMs }; + selectAccount(modelId = null, options = {}) { + if (!this.#strategy) { + throw new Error('AccountManager not initialized. Call initialize() first.'); + } + + const result = this.#strategy.selectAccount(this.#accounts, modelId, { + currentIndex: this.#currentIndex, + onSave: () => this.saveToDisk(), + ...options + }); + + this.#currentIndex = result.index; + return { account: result.account, waitMs: result.waitMs || 0 }; + } + + /** + * Notify the strategy of a successful request + * @param {Object} account - The account that was used + * @param {string} modelId - The model ID that was used + */ + notifySuccess(account, modelId) { + if (this.#strategy) { + this.#strategy.onSuccess(account, modelId); + } + } + + /** + * Notify the strategy of a rate limit + * @param {Object} account - The account that was rate-limited + * @param {string} modelId - The model ID that was rate-limited + */ + notifyRateLimit(account, modelId) { + if (this.#strategy) { + this.#strategy.onRateLimit(account, modelId); + } + } + + /** + * Notify the strategy of a failure + * @param {Object} account - The account that failed + * @param {string} modelId - The model ID that failed + */ + notifyFailure(account, modelId) { + if (this.#strategy) { + this.#strategy.onFailure(account, modelId); + } + } + + /** + * Get the current strategy name + * @returns {string} Strategy name + */ + getStrategyName() { + return this.#strategyName; + } + + /** + * Get the strategy display label + * @returns {string} Strategy display label + */ + getStrategyLabel() { + return getStrategyLabel(this.#strategyName); + } + + /** + * Get the health tracker from the current strategy (if available) + * Used by handlers for consecutive failure tracking + * Only available when using hybrid strategy + * @returns {Object|null} Health tracker instance or null if not available + */ + getHealthTracker() { + if (this.#strategy && typeof this.#strategy.getHealthTracker === 'function') { + return this.#strategy.getHealthTracker(); + } + return null; } /** diff --git a/src/account-manager/selection.js b/src/account-manager/selection.js deleted file mode 100644 index ad41307..0000000 --- a/src/account-manager/selection.js +++ /dev/null @@ -1,201 +0,0 @@ -/** - * Account Selection - * - * Handles account picking logic (round-robin, sticky) for cache continuity. - * All rate limit checks are model-specific. - */ - -import { MAX_WAIT_BEFORE_ERROR_MS } from '../constants.js'; -import { formatDuration } from '../utils/helpers.js'; -import { logger } from '../utils/logger.js'; -import { clearExpiredLimits, getAvailableAccounts } from './rate-limits.js'; - -/** - * Check if an account is usable for a specific model - * @param {Object} account - Account object - * @param {string} modelId - Model ID to check - * @returns {boolean} True if account is usable - */ -function isAccountUsable(account, modelId) { - if (!account || account.isInvalid) return false; - - // WebUI: Skip disabled accounts - if (account.enabled === false) return false; - - if (modelId && account.modelRateLimits && account.modelRateLimits[modelId]) { - const limit = account.modelRateLimits[modelId]; - if (limit.isRateLimited && limit.resetTime > Date.now()) { - return false; - } - } - - return true; -} - -/** - * Pick the next available account (fallback when current is unavailable). - * - * @param {Array} accounts - Array of account objects - * @param {number} currentIndex - Current account index - * @param {Function} onSave - Callback to save changes - * @param {string} [modelId] - Model ID to check rate limits for - * @returns {{account: Object|null, newIndex: number}} The next available account and new index - */ -export function pickNext(accounts, currentIndex, onSave, modelId = null) { - clearExpiredLimits(accounts); - - const available = getAvailableAccounts(accounts, modelId); - if (available.length === 0) { - return { account: null, newIndex: currentIndex }; - } - - // Clamp index to valid range - let index = currentIndex; - if (index >= accounts.length) { - index = 0; - } - - // Find next available account starting from index AFTER current - for (let i = 1; i <= accounts.length; i++) { - const idx = (index + i) % accounts.length; - const account = accounts[idx]; - - if (isAccountUsable(account, modelId)) { - account.lastUsed = Date.now(); - - const position = idx + 1; - const total = accounts.length; - logger.info(`[AccountManager] Using account: ${account.email} (${position}/${total})`); - - // Trigger save (don't await to avoid blocking) - if (onSave) onSave(); - - return { account, newIndex: idx }; - } - } - - return { account: null, newIndex: currentIndex }; -} - -/** - * Get the current account without advancing the index (sticky selection). - * - * @param {Array} accounts - Array of account objects - * @param {number} currentIndex - Current account index - * @param {Function} onSave - Callback to save changes - * @param {string} [modelId] - Model ID to check rate limits for - * @returns {{account: Object|null, newIndex: number}} The current account and index - */ -export function getCurrentStickyAccount(accounts, currentIndex, onSave, modelId = null) { - clearExpiredLimits(accounts); - - if (accounts.length === 0) { - return { account: null, newIndex: currentIndex }; - } - - // Clamp index to valid range - let index = currentIndex; - if (index >= accounts.length) { - index = 0; - } - - // Get current account directly (activeIndex = current account) - const account = accounts[index]; - - if (isAccountUsable(account, modelId)) { - account.lastUsed = Date.now(); - // Trigger save (don't await to avoid blocking) - if (onSave) onSave(); - return { account, newIndex: index }; - } - - return { account: null, newIndex: index }; -} - -/** - * Check if we should wait for the current account's rate limit to reset. - * - * @param {Array} accounts - Array of account objects - * @param {number} currentIndex - Current account index - * @param {string} [modelId] - Model ID to check rate limits for - * @returns {{shouldWait: boolean, waitMs: number, account: Object|null}} - */ -export function shouldWaitForCurrentAccount(accounts, currentIndex, modelId = null) { - if (accounts.length === 0) { - return { shouldWait: false, waitMs: 0, account: null }; - } - - // Clamp index to valid range - let index = currentIndex; - if (index >= accounts.length) { - index = 0; - } - - // Get current account directly (activeIndex = current account) - const account = accounts[index]; - - if (!account || account.isInvalid) { - return { shouldWait: false, waitMs: 0, account: null }; - } - - let waitMs = 0; - - // Check model-specific limit - if (modelId && account.modelRateLimits && account.modelRateLimits[modelId]) { - const limit = account.modelRateLimits[modelId]; - if (limit.isRateLimited && limit.resetTime) { - waitMs = limit.resetTime - Date.now(); - } - } - - // If wait time is within threshold, recommend waiting - if (waitMs > 0 && waitMs <= MAX_WAIT_BEFORE_ERROR_MS) { - return { shouldWait: true, waitMs, account }; - } - - return { shouldWait: false, waitMs: 0, account }; -} - -/** - * Pick an account with sticky selection preference. - * Prefers the current account for cache continuity. - * - * @param {Array} accounts - Array of account objects - * @param {number} currentIndex - Current account index - * @param {Function} onSave - Callback to save changes - * @param {string} [modelId] - Model ID to check rate limits for - * @returns {{account: Object|null, waitMs: number, newIndex: number}} - */ -export function pickStickyAccount(accounts, currentIndex, onSave, modelId = null) { - // First try to get the current sticky account - const { account: stickyAccount, newIndex: stickyIndex } = getCurrentStickyAccount(accounts, currentIndex, onSave, modelId); - if (stickyAccount) { - return { account: stickyAccount, waitMs: 0, newIndex: stickyIndex }; - } - - // Current account is rate-limited or invalid. - // CHECK IF OTHERS ARE AVAILABLE before deciding to wait. - const available = getAvailableAccounts(accounts, modelId); - if (available.length > 0) { - // Found a free account! Switch immediately. - const { account: nextAccount, newIndex } = pickNext(accounts, currentIndex, onSave, modelId); - if (nextAccount) { - logger.info(`[AccountManager] Switched to new account (failover): ${nextAccount.email}`); - return { account: nextAccount, waitMs: 0, newIndex }; - } - } - - // No other accounts available. Now checking if we should wait for current account. - const waitInfo = shouldWaitForCurrentAccount(accounts, currentIndex, modelId); - if (waitInfo.shouldWait) { - logger.info(`[AccountManager] Waiting ${formatDuration(waitInfo.waitMs)} for sticky account: ${waitInfo.account.email}`); - return { account: null, waitMs: waitInfo.waitMs, newIndex: currentIndex }; - } - - // Current account unavailable for too long/invalid, and no others available? - const { account: nextAccount, newIndex } = pickNext(accounts, currentIndex, onSave, modelId); - if (nextAccount) { - logger.info(`[AccountManager] Switched to new account for cache: ${nextAccount.email}`); - } - return { account: nextAccount, waitMs: 0, newIndex }; -} diff --git a/src/account-manager/strategies/base-strategy.js b/src/account-manager/strategies/base-strategy.js new file mode 100644 index 0000000..57007cc --- /dev/null +++ b/src/account-manager/strategies/base-strategy.js @@ -0,0 +1,104 @@ +/** + * Base Strategy + * + * Abstract base class defining the interface for account selection strategies. + * All strategies must implement the selectAccount method. + */ + +/** + * @typedef {Object} SelectionResult + * @property {Object|null} account - The selected account or null if none available + * @property {number} index - The index of the selected account + * @property {number} [waitMs] - Optional wait time before account becomes available + */ + +export class BaseStrategy { + /** + * Create a new BaseStrategy + * @param {Object} config - Strategy configuration + */ + constructor(config = {}) { + if (new.target === BaseStrategy) { + throw new Error('BaseStrategy is abstract and cannot be instantiated directly'); + } + this.config = config; + } + + /** + * Select an account for a request + * @param {Array} accounts - Array of account objects + * @param {string} modelId - The model ID for the request + * @param {Object} options - Additional options + * @param {number} options.currentIndex - Current account index + * @param {string} [options.sessionId] - Session ID for cache continuity + * @param {Function} [options.onSave] - Callback to save changes + * @returns {SelectionResult} The selected account and index + */ + selectAccount(accounts, modelId, options = {}) { + throw new Error('selectAccount must be implemented by subclass'); + } + + /** + * Called after a successful request + * @param {Object} account - The account that was used + * @param {string} modelId - The model ID that was used + */ + onSuccess(account, modelId) { + // Default: no-op, override in subclass if needed + } + + /** + * Called when a request is rate-limited + * @param {Object} account - The account that was rate-limited + * @param {string} modelId - The model ID that was rate-limited + */ + onRateLimit(account, modelId) { + // Default: no-op, override in subclass if needed + } + + /** + * Called when a request fails (non-rate-limit error) + * @param {Object} account - The account that failed + * @param {string} modelId - The model ID that failed + */ + onFailure(account, modelId) { + // Default: no-op, override in subclass if needed + } + + /** + * Check if an account is usable for a specific model + * @param {Object} account - Account object + * @param {string} modelId - Model ID to check + * @returns {boolean} True if account is usable + */ + isAccountUsable(account, modelId) { + if (!account || account.isInvalid) return false; + + // Skip disabled accounts + if (account.enabled === false) return false; + + // Check model-specific rate limit + if (modelId && account.modelRateLimits && account.modelRateLimits[modelId]) { + const limit = account.modelRateLimits[modelId]; + if (limit.isRateLimited && limit.resetTime > Date.now()) { + return false; + } + } + + return true; + } + + /** + * Get all usable accounts for a model + * @param {Array} accounts - Array of account objects + * @param {string} modelId - Model ID to check + * @returns {Array} Array of usable accounts with their original indices + */ + getUsableAccounts(accounts, modelId) { + return accounts + .map((account, index) => ({ account, index })) + .filter(({ account }) => this.isAccountUsable(account, modelId)); + } +} + +export default BaseStrategy; diff --git a/src/account-manager/strategies/hybrid-strategy.js b/src/account-manager/strategies/hybrid-strategy.js new file mode 100644 index 0000000..6814d7f --- /dev/null +++ b/src/account-manager/strategies/hybrid-strategy.js @@ -0,0 +1,195 @@ +/** + * Hybrid Strategy + * + * Smart selection based on health score, token bucket, and LRU freshness. + * Combines multiple signals for optimal account distribution. + * + * Scoring formula: + * score = (Health × 2) + ((Tokens / MaxTokens × 100) × 5) + (LRU × 0.1) + * + * Filters accounts that are: + * - Not rate-limited + * - Not invalid or disabled + * - Health score >= minUsable + * - Has tokens available + */ + +import { BaseStrategy } from './base-strategy.js'; +import { HealthTracker, TokenBucketTracker } from './trackers/index.js'; +import { logger } from '../../utils/logger.js'; + +// Default weights for scoring +const DEFAULT_WEIGHTS = { + health: 2, + tokens: 5, + lru: 0.1 +}; + +export class HybridStrategy extends BaseStrategy { + #healthTracker; + #tokenBucketTracker; + #weights; + + /** + * Create a new HybridStrategy + * @param {Object} config - Strategy configuration + * @param {Object} [config.healthScore] - Health tracker configuration + * @param {Object} [config.tokenBucket] - Token bucket configuration + * @param {Object} [config.weights] - Scoring weights + */ + constructor(config = {}) { + super(config); + this.#healthTracker = new HealthTracker(config.healthScore || {}); + this.#tokenBucketTracker = new TokenBucketTracker(config.tokenBucket || {}); + this.#weights = { ...DEFAULT_WEIGHTS, ...config.weights }; + } + + /** + * Select an account based on combined health, tokens, and LRU score + * + * @param {Array} accounts - Array of account objects + * @param {string} modelId - The model ID for the request + * @param {Object} options - Additional options + * @returns {SelectionResult} The selected account and index + */ + selectAccount(accounts, modelId, options = {}) { + const { onSave } = options; + + if (accounts.length === 0) { + return { account: null, index: 0, waitMs: 0 }; + } + + // Get candidates that pass all filters + const candidates = this.#getCandidates(accounts, modelId); + + if (candidates.length === 0) { + logger.debug('[HybridStrategy] No candidates available'); + return { account: null, index: 0, waitMs: 0 }; + } + + // Score and sort candidates + const scored = candidates.map(({ account, index }) => ({ + account, + index, + score: this.#calculateScore(account) + })); + + scored.sort((a, b) => b.score - a.score); + + // Select the best candidate + const best = scored[0]; + best.account.lastUsed = Date.now(); + + // Consume a token from the bucket + this.#tokenBucketTracker.consume(best.account.email); + + if (onSave) onSave(); + + const position = best.index + 1; + const total = accounts.length; + logger.info(`[HybridStrategy] Using account: ${best.account.email} (${position}/${total}, score: ${best.score.toFixed(1)})`); + + return { account: best.account, index: best.index, waitMs: 0 }; + } + + /** + * Called after a successful request + */ + onSuccess(account, modelId) { + if (account && account.email) { + this.#healthTracker.recordSuccess(account.email); + } + } + + /** + * Called when a request is rate-limited + */ + onRateLimit(account, modelId) { + if (account && account.email) { + this.#healthTracker.recordRateLimit(account.email); + } + } + + /** + * Called when a request fails + */ + onFailure(account, modelId) { + if (account && account.email) { + this.#healthTracker.recordFailure(account.email); + // Refund the token since the request didn't complete + this.#tokenBucketTracker.refund(account.email); + } + } + + /** + * Get candidates that pass all filters + * @private + */ + #getCandidates(accounts, modelId) { + return accounts + .map((account, index) => ({ account, index })) + .filter(({ account }) => { + // Basic usability check + if (!this.isAccountUsable(account, modelId)) { + return false; + } + + // Health score check + if (!this.#healthTracker.isUsable(account.email)) { + return false; + } + + // Token availability check + if (!this.#tokenBucketTracker.hasTokens(account.email)) { + return false; + } + + return true; + }); + } + + /** + * Calculate the combined score for an account + * @private + */ + #calculateScore(account) { + const email = account.email; + + // Health component (0-100 scaled by weight) + const health = this.#healthTracker.getScore(email); + const healthComponent = health * this.#weights.health; + + // Token component (0-100 scaled by weight) + const tokens = this.#tokenBucketTracker.getTokens(email); + const maxTokens = this.#tokenBucketTracker.getMaxTokens(); + const tokenRatio = tokens / maxTokens; + const tokenComponent = (tokenRatio * 100) * this.#weights.tokens; + + // LRU component (older = higher score) + // Use time since last use, capped at 1 hour for scoring + const lastUsed = account.lastUsed || 0; + const timeSinceLastUse = Math.min(Date.now() - lastUsed, 3600000); // Cap at 1 hour + const lruMinutes = timeSinceLastUse / 60000; + const lruComponent = lruMinutes * this.#weights.lru; + + return healthComponent + tokenComponent + lruComponent; + } + + /** + * Get the health tracker (for testing/debugging) + * @returns {HealthTracker} The health tracker instance + */ + getHealthTracker() { + return this.#healthTracker; + } + + /** + * Get the token bucket tracker (for testing/debugging) + * @returns {TokenBucketTracker} The token bucket tracker instance + */ + getTokenBucketTracker() { + return this.#tokenBucketTracker; + } +} + +export default HybridStrategy; diff --git a/src/account-manager/strategies/index.js b/src/account-manager/strategies/index.js new file mode 100644 index 0000000..239170a --- /dev/null +++ b/src/account-manager/strategies/index.js @@ -0,0 +1,85 @@ +/** + * Strategy Factory + * + * Creates and exports account selection strategy instances. + */ + +import { StickyStrategy } from './sticky-strategy.js'; +import { RoundRobinStrategy } from './round-robin-strategy.js'; +import { HybridStrategy } from './hybrid-strategy.js'; +import { logger } from '../../utils/logger.js'; +import { + SELECTION_STRATEGIES, + DEFAULT_SELECTION_STRATEGY +} from '../../constants.js'; + +// Re-export strategy constants for convenience +export const STRATEGY_NAMES = SELECTION_STRATEGIES; +export const DEFAULT_STRATEGY = DEFAULT_SELECTION_STRATEGY; + +// Strategy display labels +export const STRATEGY_LABELS = { + 'sticky': 'Sticky (Cache Optimized)', + 'round-robin': 'Round Robin (Load Balanced)', + 'hybrid': 'Hybrid (Smart Distribution)' +}; + +/** + * Create a strategy instance + * @param {string} strategyName - Name of the strategy ('sticky', 'round-robin', 'hybrid') + * @param {Object} config - Strategy configuration + * @returns {BaseStrategy} The strategy instance + */ +export function createStrategy(strategyName, config = {}) { + const name = (strategyName || DEFAULT_STRATEGY).toLowerCase(); + + switch (name) { + case 'sticky': + logger.debug('[Strategy] Creating StickyStrategy'); + return new StickyStrategy(config); + + case 'round-robin': + case 'roundrobin': + logger.debug('[Strategy] Creating RoundRobinStrategy'); + return new RoundRobinStrategy(config); + + case 'hybrid': + logger.debug('[Strategy] Creating HybridStrategy'); + return new HybridStrategy(config); + + default: + logger.warn(`[Strategy] Unknown strategy "${strategyName}", falling back to ${DEFAULT_STRATEGY}`); + return new HybridStrategy(config); + } +} + +/** + * Check if a strategy name is valid + * @param {string} name - Strategy name to check + * @returns {boolean} True if valid + */ +export function isValidStrategy(name) { + if (!name) return false; + const lower = name.toLowerCase(); + return STRATEGY_NAMES.includes(lower) || lower === 'roundrobin'; +} + +/** + * Get the display label for a strategy + * @param {string} name - Strategy name + * @returns {string} Display label + */ +export function getStrategyLabel(name) { + const lower = (name || DEFAULT_STRATEGY).toLowerCase(); + if (lower === 'roundrobin') return STRATEGY_LABELS['round-robin']; + return STRATEGY_LABELS[lower] || STRATEGY_LABELS[DEFAULT_STRATEGY]; +} + +// Re-export strategies for direct use +export { StickyStrategy } from './sticky-strategy.js'; +export { RoundRobinStrategy } from './round-robin-strategy.js'; +export { HybridStrategy } from './hybrid-strategy.js'; +export { BaseStrategy } from './base-strategy.js'; + +// Re-export trackers +export { HealthTracker, TokenBucketTracker } from './trackers/index.js'; diff --git a/src/account-manager/strategies/round-robin-strategy.js b/src/account-manager/strategies/round-robin-strategy.js new file mode 100644 index 0000000..1b79bd6 --- /dev/null +++ b/src/account-manager/strategies/round-robin-strategy.js @@ -0,0 +1,76 @@ +/** + * Round-Robin Strategy + * + * Rotates to the next account on every request for maximum throughput. + * Does not maintain cache continuity but maximizes concurrent requests. + */ + +import { BaseStrategy } from './base-strategy.js'; +import { logger } from '../../utils/logger.js'; + +export class RoundRobinStrategy extends BaseStrategy { + #cursor = 0; // Tracks current position in rotation + + /** + * Create a new RoundRobinStrategy + * @param {Object} config - Strategy configuration + */ + constructor(config = {}) { + super(config); + } + + /** + * Select the next available account in rotation + * + * @param {Array} accounts - Array of account objects + * @param {string} modelId - The model ID for the request + * @param {Object} options - Additional options + * @returns {SelectionResult} The selected account and index + */ + selectAccount(accounts, modelId, options = {}) { + const { onSave } = options; + + if (accounts.length === 0) { + return { account: null, index: 0, waitMs: 0 }; + } + + // Clamp cursor to valid range + if (this.#cursor >= accounts.length) { + this.#cursor = 0; + } + + // Start from the next position after the cursor + const startIndex = (this.#cursor + 1) % accounts.length; + + // Try each account starting from startIndex + for (let i = 0; i < accounts.length; i++) { + const idx = (startIndex + i) % accounts.length; + const account = accounts[idx]; + + if (this.isAccountUsable(account, modelId)) { + account.lastUsed = Date.now(); + this.#cursor = idx; + + if (onSave) onSave(); + + const position = idx + 1; + const total = accounts.length; + logger.info(`[RoundRobinStrategy] Using account: ${account.email} (${position}/${total})`); + + return { account, index: idx, waitMs: 0 }; + } + } + + // No usable accounts found + return { account: null, index: this.#cursor, waitMs: 0 }; + } + + /** + * Reset the cursor position + */ + resetCursor() { + this.#cursor = 0; + } +} + +export default RoundRobinStrategy; diff --git a/src/account-manager/strategies/sticky-strategy.js b/src/account-manager/strategies/sticky-strategy.js new file mode 100644 index 0000000..988f370 --- /dev/null +++ b/src/account-manager/strategies/sticky-strategy.js @@ -0,0 +1,138 @@ +/** + * Sticky Strategy + * + * Keeps using the same account until it becomes unavailable (rate-limited or invalid). + * Best for prompt caching as it maintains cache continuity across requests. + */ + +import { BaseStrategy } from './base-strategy.js'; +import { logger } from '../../utils/logger.js'; +import { formatDuration } from '../../utils/helpers.js'; +import { MAX_WAIT_BEFORE_ERROR_MS } from '../../constants.js'; + +export class StickyStrategy extends BaseStrategy { + /** + * Create a new StickyStrategy + * @param {Object} config - Strategy configuration + */ + constructor(config = {}) { + super(config); + } + + /** + * Select an account with sticky preference + * Prefers the current account for cache continuity, only switches when: + * - Current account is rate-limited for > 2 minutes + * - Current account is invalid + * - Current account is disabled + * + * @param {Array} accounts - Array of account objects + * @param {string} modelId - The model ID for the request + * @param {Object} options - Additional options + * @returns {SelectionResult} The selected account and index + */ + selectAccount(accounts, modelId, options = {}) { + const { currentIndex = 0, onSave } = options; + + if (accounts.length === 0) { + return { account: null, index: currentIndex, waitMs: 0 }; + } + + // Clamp index to valid range + let index = currentIndex >= accounts.length ? 0 : currentIndex; + const currentAccount = accounts[index]; + + // Check if current account is usable + if (this.isAccountUsable(currentAccount, modelId)) { + currentAccount.lastUsed = Date.now(); + if (onSave) onSave(); + return { account: currentAccount, index, waitMs: 0 }; + } + + // Current account is not usable - check if others are available + const usableAccounts = this.getUsableAccounts(accounts, modelId); + + if (usableAccounts.length > 0) { + // Found a free account - switch immediately + const { account: nextAccount, index: nextIndex } = this.#pickNext( + accounts, + index, + modelId, + onSave + ); + if (nextAccount) { + logger.info(`[StickyStrategy] Switched to new account (failover): ${nextAccount.email}`); + return { account: nextAccount, index: nextIndex, waitMs: 0 }; + } + } + + // No other accounts available - check if we should wait for current + const waitInfo = this.#shouldWaitForAccount(currentAccount, modelId); + if (waitInfo.shouldWait) { + logger.info(`[StickyStrategy] Waiting ${formatDuration(waitInfo.waitMs)} for sticky account: ${currentAccount.email}`); + return { account: null, index, waitMs: waitInfo.waitMs }; + } + + // Current account unavailable for too long, try to find any other + const { account: nextAccount, index: nextIndex } = this.#pickNext( + accounts, + index, + modelId, + onSave + ); + + return { account: nextAccount, index: nextIndex, waitMs: 0 }; + } + + /** + * Pick the next available account starting from after the current index + * @private + */ + #pickNext(accounts, currentIndex, modelId, onSave) { + for (let i = 1; i <= accounts.length; i++) { + const idx = (currentIndex + i) % accounts.length; + const account = accounts[idx]; + + if (this.isAccountUsable(account, modelId)) { + account.lastUsed = Date.now(); + if (onSave) onSave(); + + const position = idx + 1; + const total = accounts.length; + logger.info(`[StickyStrategy] Using account: ${account.email} (${position}/${total})`); + + return { account, index: idx }; + } + } + + return { account: null, index: currentIndex }; + } + + /** + * Check if we should wait for an account's rate limit to reset + * @private + */ + #shouldWaitForAccount(account, modelId) { + if (!account || account.isInvalid || account.enabled === false) { + return { shouldWait: false, waitMs: 0 }; + } + + let waitMs = 0; + + if (modelId && account.modelRateLimits && account.modelRateLimits[modelId]) { + const limit = account.modelRateLimits[modelId]; + if (limit.isRateLimited && limit.resetTime) { + waitMs = limit.resetTime - Date.now(); + } + } + + // Wait if within threshold + if (waitMs > 0 && waitMs <= MAX_WAIT_BEFORE_ERROR_MS) { + return { shouldWait: true, waitMs }; + } + + return { shouldWait: false, waitMs: 0 }; + } +} + +export default StickyStrategy; diff --git a/src/account-manager/strategies/trackers/health-tracker.js b/src/account-manager/strategies/trackers/health-tracker.js new file mode 100644 index 0000000..a53274f --- /dev/null +++ b/src/account-manager/strategies/trackers/health-tracker.js @@ -0,0 +1,162 @@ +/** + * Health Tracker + * + * Tracks per-account health scores to prioritize healthy accounts. + * Scores increase on success and decrease on failures/rate limits. + * Passive recovery over time helps accounts recover from temporary issues. + */ + +// Default configuration (matches opencode-antigravity-auth) +const DEFAULT_CONFIG = { + initial: 70, // Starting score for new accounts + successReward: 1, // Points on successful request + rateLimitPenalty: -10, // Points on rate limit + failurePenalty: -20, // Points on other failures + recoveryPerHour: 2, // Passive recovery rate + minUsable: 50, // Minimum score to be selected + maxScore: 100 // Maximum score cap +}; + +export class HealthTracker { + #scores = new Map(); // email -> { score, lastUpdated, consecutiveFailures } + #config; + + /** + * Create a new HealthTracker + * @param {Object} config - Health score configuration + */ + constructor(config = {}) { + this.#config = { ...DEFAULT_CONFIG, ...config }; + } + + /** + * Get the health score for an account + * @param {string} email - Account email + * @returns {number} Current health score (with passive recovery applied) + */ + getScore(email) { + const record = this.#scores.get(email); + if (!record) { + return this.#config.initial; + } + + // Apply passive recovery based on time elapsed + const now = Date.now(); + const hoursElapsed = (now - record.lastUpdated) / (1000 * 60 * 60); + const recovery = hoursElapsed * this.#config.recoveryPerHour; + const recoveredScore = Math.min( + this.#config.maxScore, + record.score + recovery + ); + + return recoveredScore; + } + + /** + * Record a successful request for an account + * @param {string} email - Account email + */ + recordSuccess(email) { + const currentScore = this.getScore(email); + const newScore = Math.min( + this.#config.maxScore, + currentScore + this.#config.successReward + ); + this.#scores.set(email, { + score: newScore, + lastUpdated: Date.now(), + consecutiveFailures: 0 // Reset on success + }); + } + + /** + * Record a rate limit for an account + * @param {string} email - Account email + */ + recordRateLimit(email) { + const record = this.#scores.get(email); + const currentScore = this.getScore(email); + const newScore = Math.max( + 0, + currentScore + this.#config.rateLimitPenalty + ); + this.#scores.set(email, { + score: newScore, + lastUpdated: Date.now(), + consecutiveFailures: (record?.consecutiveFailures ?? 0) + 1 + }); + } + + /** + * Record a failure for an account + * @param {string} email - Account email + */ + recordFailure(email) { + const record = this.#scores.get(email); + const currentScore = this.getScore(email); + const newScore = Math.max( + 0, + currentScore + this.#config.failurePenalty + ); + this.#scores.set(email, { + score: newScore, + lastUpdated: Date.now(), + consecutiveFailures: (record?.consecutiveFailures ?? 0) + 1 + }); + } + + /** + * Check if an account is usable based on health score + * @param {string} email - Account email + * @returns {boolean} True if account health score is above minimum threshold + */ + isUsable(email) { + return this.getScore(email) >= this.#config.minUsable; + } + + /** + * Get the minimum usable score threshold + * @returns {number} Minimum score for an account to be usable + */ + getMinUsable() { + return this.#config.minUsable; + } + + /** + * Get the maximum score cap + * @returns {number} Maximum health score + */ + getMaxScore() { + return this.#config.maxScore; + } + + /** + * Reset the score for an account (e.g., after re-authentication) + * @param {string} email - Account email + */ + reset(email) { + this.#scores.set(email, { + score: this.#config.initial, + lastUpdated: Date.now(), + consecutiveFailures: 0 + }); + } + + /** + * Get the consecutive failure count for an account + * @param {string} email - Account email + * @returns {number} Number of consecutive failures + */ + getConsecutiveFailures(email) { + return this.#scores.get(email)?.consecutiveFailures ?? 0; + } + + /** + * Clear all tracked scores + */ + clear() { + this.#scores.clear(); + } +} + +export default HealthTracker; diff --git a/src/account-manager/strategies/trackers/index.js b/src/account-manager/strategies/trackers/index.js new file mode 100644 index 0000000..974f527 --- /dev/null +++ b/src/account-manager/strategies/trackers/index.js @@ -0,0 +1,8 @@ +/** + * Trackers Index + * + * Exports all tracker classes for account selection strategies. + */ + +export { HealthTracker } from './health-tracker.js'; +export { TokenBucketTracker } from './token-bucket-tracker.js'; diff --git a/src/account-manager/strategies/trackers/token-bucket-tracker.js b/src/account-manager/strategies/trackers/token-bucket-tracker.js new file mode 100644 index 0000000..33d548c --- /dev/null +++ b/src/account-manager/strategies/trackers/token-bucket-tracker.js @@ -0,0 +1,121 @@ +/** + * Token Bucket Tracker + * + * Client-side rate limiting using the token bucket algorithm. + * Each account has a bucket of tokens that regenerate over time. + * Requests consume tokens; accounts without tokens are deprioritized. + */ + +// Default configuration (matches opencode-antigravity-auth) +const DEFAULT_CONFIG = { + maxTokens: 50, // Maximum token capacity + tokensPerMinute: 6, // Regeneration rate + initialTokens: 50 // Starting tokens +}; + +export class TokenBucketTracker { + #buckets = new Map(); // email -> { tokens, lastUpdated } + #config; + + /** + * Create a new TokenBucketTracker + * @param {Object} config - Token bucket configuration + */ + constructor(config = {}) { + this.#config = { ...DEFAULT_CONFIG, ...config }; + } + + /** + * Get the current token count for an account + * @param {string} email - Account email + * @returns {number} Current token count (with regeneration applied) + */ + getTokens(email) { + const bucket = this.#buckets.get(email); + if (!bucket) { + return this.#config.initialTokens; + } + + // Apply token regeneration based on time elapsed + const now = Date.now(); + const minutesElapsed = (now - bucket.lastUpdated) / (1000 * 60); + const regenerated = minutesElapsed * this.#config.tokensPerMinute; + const currentTokens = Math.min( + this.#config.maxTokens, + bucket.tokens + regenerated + ); + + return currentTokens; + } + + /** + * Check if an account has tokens available + * @param {string} email - Account email + * @returns {boolean} True if account has at least 1 token + */ + hasTokens(email) { + return this.getTokens(email) >= 1; + } + + /** + * Consume a token from an account's bucket + * @param {string} email - Account email + * @returns {boolean} True if token was consumed, false if no tokens available + */ + consume(email) { + const currentTokens = this.getTokens(email); + if (currentTokens < 1) { + return false; + } + + this.#buckets.set(email, { + tokens: currentTokens - 1, + lastUpdated: Date.now() + }); + return true; + } + + /** + * Refund a token to an account's bucket (e.g., on request failure before processing) + * @param {string} email - Account email + */ + refund(email) { + const currentTokens = this.getTokens(email); + const newTokens = Math.min( + this.#config.maxTokens, + currentTokens + 1 + ); + this.#buckets.set(email, { + tokens: newTokens, + lastUpdated: Date.now() + }); + } + + /** + * Get the maximum token capacity + * @returns {number} Maximum tokens per bucket + */ + getMaxTokens() { + return this.#config.maxTokens; + } + + /** + * Reset the bucket for an account + * @param {string} email - Account email + */ + reset(email) { + this.#buckets.set(email, { + tokens: this.#config.initialTokens, + lastUpdated: Date.now() + }); + } + + /** + * Clear all tracked buckets + */ + clear() { + this.#buckets.clear(); + } +} + +export default TokenBucketTracker; diff --git a/src/cloudcode/message-handler.js b/src/cloudcode/message-handler.js index 970e834..586ee16 100644 --- a/src/cloudcode/message-handler.js +++ b/src/cloudcode/message-handler.js @@ -10,6 +10,11 @@ import { MAX_RETRIES, MAX_WAIT_BEFORE_ERROR_MS, DEFAULT_COOLDOWN_MS, + RATE_LIMIT_DEDUP_WINDOW_MS, + MAX_CONSECUTIVE_FAILURES, + EXTENDED_COOLDOWN_MS, + CAPACITY_RETRY_DELAY_MS, + MAX_CAPACITY_RETRIES, isThinkingModel } from '../constants.js'; import { convertGoogleToAnthropic } from '../format/index.js'; @@ -21,6 +26,85 @@ import { buildCloudCodeRequest, buildHeaders } from './request-builder.js'; import { parseThinkingSSEResponse } from './sse-parser.js'; import { getFallbackModel } from '../fallback-config.js'; +/** + * Gap 1: Rate limit deduplication - prevents thundering herd on concurrent rate limits + * Tracks last rate limit timestamp per model to skip duplicate retries + */ +const lastRateLimitTimestamps = new Map(); // modelId -> timestamp + +/** + * Check if we should skip retry due to recent rate limit on this model + * @param {string} model - Model ID + * @returns {boolean} True if retry should be skipped (within dedup window) + */ +function shouldSkipRetryDueToDedup(model) { + const lastTimestamp = lastRateLimitTimestamps.get(model); + if (!lastTimestamp) return false; + + const elapsed = Date.now() - lastTimestamp; + if (elapsed < RATE_LIMIT_DEDUP_WINDOW_MS) { + logger.debug(`[CloudCode] Rate limit on ${model} within dedup window (${elapsed}ms ago), skipping retry`); + return true; + } + return false; +} + +/** + * Record rate limit timestamp for deduplication + * @param {string} model - Model ID + */ +function recordRateLimitTimestamp(model) { + lastRateLimitTimestamps.set(model, Date.now()); +} + +/** + * Clear rate limit timestamp after successful retry + * @param {string} model - Model ID + */ +function clearRateLimitTimestamp(model) { + lastRateLimitTimestamps.delete(model); +} + +/** + * Gap 3: Detect permanent authentication failures that require re-authentication + * These should mark the account as invalid rather than just clearing cache + * @param {string} errorText - Error message from API + * @returns {boolean} True if permanent auth failure + */ +function isPermanentAuthFailure(errorText) { + const lower = (errorText || '').toLowerCase(); + return lower.includes('invalid_grant') || + lower.includes('token revoked') || + lower.includes('token has been expired or revoked') || + lower.includes('token_revoked') || + lower.includes('invalid_client') || + lower.includes('credentials are invalid'); +} + +/** + * Gap 4: Detect if 429 error is due to model capacity (not user quota) + * Capacity issues should retry on same account with shorter delay + * @param {string} errorText - Error message from API + * @returns {boolean} True if capacity exhausted (not quota) + */ +function isModelCapacityExhausted(errorText) { + const lower = (errorText || '').toLowerCase(); + return lower.includes('model_capacity_exhausted') || + lower.includes('capacity_exhausted') || + lower.includes('model is currently overloaded') || + lower.includes('service temporarily unavailable'); +} + +// Periodically clean up stale dedup timestamps (every 60 seconds) +setInterval(() => { + const cutoff = Date.now() - 60000; // 1 minute + for (const [model, timestamp] of lastRateLimitTimestamps.entries()) { + if (timestamp < cutoff) { + lastRateLimitTimestamps.delete(model); + } + } +}, 60000); + /** * Send a non-streaming request to Cloud Code with multi-account support * Uses SSE endpoint for thinking models (non-streaming doesn't return thinking blocks) @@ -83,10 +167,14 @@ export async function sendMessage(anthropicRequest, accountManager, fallbackEnab throw new Error('No accounts available'); } - // Pick sticky account (prefers current for cache continuity) - let account = accountManager.getCurrentStickyAccount(model); - if (!account) { - account = accountManager.pickNext(model); + // Select account using configured strategy + const { account, waitMs } = accountManager.selectAccount(model); + + // If strategy returns a wait time, sleep and retry + if (!account && waitMs > 0) { + logger.info(`[CloudCode] Waiting ${formatDuration(waitMs)} for account...`); + await sleep(waitMs + 500); + continue; } if (!account) { @@ -101,11 +189,14 @@ export async function sendMessage(anthropicRequest, accountManager, fallbackEnab logger.debug(`[CloudCode] Sending request for model: ${model}`); - // Try each endpoint + // Try each endpoint with index-based loop for capacity retry support let lastError = null; let retriedOnce = false; // Track if we've already retried for short rate limit + let capacityRetryCount = 0; // Gap 4: Track capacity exhaustion retries + let endpointIndex = 0; - for (const endpoint of ANTIGRAVITY_ENDPOINT_FALLBACKS) { + while (endpointIndex < ANTIGRAVITY_ENDPOINT_FALLBACKS.length) { + const endpoint = ANTIGRAVITY_ENDPOINT_FALLBACKS[endpointIndex]; try { const url = isThinking ? `${endpoint}/v1internal:streamGenerateContent?alt=sse` @@ -122,16 +213,45 @@ export async function sendMessage(anthropicRequest, accountManager, fallbackEnab logger.warn(`[CloudCode] Error at ${endpoint}: ${response.status} - ${errorText}`); if (response.status === 401) { - // Auth error - clear caches and retry with fresh token - logger.warn('[CloudCode] Auth error, refreshing token...'); + // Gap 3: Check for permanent auth failures + if (isPermanentAuthFailure(errorText)) { + logger.error(`[CloudCode] Permanent auth failure for ${account.email}: ${errorText.substring(0, 100)}`); + accountManager.markInvalid(account.email, 'Token revoked - re-authentication required'); + throw new Error(`AUTH_INVALID_PERMANENT: ${errorText}`); + } + + // Transient auth error - clear caches and retry with fresh token + logger.warn('[CloudCode] Transient auth error, refreshing token...'); accountManager.clearTokenCache(account.email); accountManager.clearProjectCache(account.email); + endpointIndex++; continue; } if (response.status === 429) { const resetMs = parseResetTime(response, errorText); + // Gap 4: Check if capacity issue (NOT quota) - retry SAME endpoint + if (isModelCapacityExhausted(errorText)) { + if (capacityRetryCount < MAX_CAPACITY_RETRIES) { + capacityRetryCount++; + const waitMs = resetMs || CAPACITY_RETRY_DELAY_MS; + logger.info(`[CloudCode] Model capacity exhausted, retry ${capacityRetryCount}/${MAX_CAPACITY_RETRIES} after ${formatDuration(waitMs)}...`); + await sleep(waitMs); + // Don't increment endpointIndex - retry same endpoint + continue; + } + // Max capacity retries exceeded - treat as quota exhaustion + logger.warn(`[CloudCode] Max capacity retries (${MAX_CAPACITY_RETRIES}) exceeded, switching account`); + } + + // Gap 1: Check deduplication window to prevent thundering herd + if (shouldSkipRetryDueToDedup(model)) { + logger.info(`[CloudCode] Skipping retry due to recent rate limit, switching account...`); + accountManager.markRateLimited(account.email, resetMs || DEFAULT_COOLDOWN_MS, model); + throw new Error(`RATE_LIMITED_DEDUP: ${errorText}`); + } + // Decision: wait and retry OR switch account if (resetMs && resetMs > DEFAULT_COOLDOWN_MS) { // Long-term quota exhaustion (> 10s) - switch to next account @@ -144,31 +264,11 @@ export async function sendMessage(anthropicRequest, accountManager, fallbackEnab if (!retriedOnce) { retriedOnce = true; + recordRateLimitTimestamp(model); // Gap 1: Record before retry logger.info(`[CloudCode] Short rate limit (${formatDuration(waitMs)}), waiting and retrying...`); await sleep(waitMs); - // Retry same endpoint - const retryResponse = await fetch(url, { - method: 'POST', - headers: buildHeaders(token, model, isThinking ? 'text/event-stream' : 'application/json'), - body: JSON.stringify(payload) - }); - - if (retryResponse.ok) { - // Process retry response - if (isThinking) { - return await parseThinkingSSEResponse(retryResponse, anthropicRequest.model); - } - const data = await retryResponse.json(); - logger.debug('[CloudCode] Response received after retry'); - return convertGoogleToAnthropic(data, anthropicRequest.model); - } - - // Retry also failed - parse new reset time - const retryErrorText = await retryResponse.text(); - const retryResetMs = parseResetTime(retryResponse, retryErrorText); - logger.warn(`[CloudCode] Retry also failed, marking and switching...`); - accountManager.markRateLimited(account.email, retryResetMs || waitMs, model); - throw new Error(`RATE_LIMITED_AFTER_RETRY: ${retryErrorText}`); + // Don't increment endpointIndex - retry same endpoint + continue; } else { // Already retried once, mark and switch accountManager.markRateLimited(account.email, waitMs, model); @@ -184,18 +284,26 @@ export async function sendMessage(anthropicRequest, accountManager, fallbackEnab logger.warn(`[CloudCode] ${response.status} error, waiting 1s before retry...`); await sleep(1000); } + endpointIndex++; continue; } } // For thinking models, parse SSE and accumulate all parts if (isThinking) { - return await parseThinkingSSEResponse(response, anthropicRequest.model); + const result = await parseThinkingSSEResponse(response, anthropicRequest.model); + // Gap 1: Clear timestamp on success + clearRateLimitTimestamp(model); + accountManager.notifySuccess(account, model); + return result; } // Non-thinking models use regular JSON const data = await response.json(); logger.debug('[CloudCode] Response received'); + // Gap 1: Clear timestamp on success + clearRateLimitTimestamp(model); + accountManager.notifySuccess(account, model); return convertGoogleToAnthropic(data, anthropicRequest.model); } catch (endpointError) { @@ -204,6 +312,7 @@ export async function sendMessage(anthropicRequest, accountManager, fallbackEnab } logger.warn(`[CloudCode] Error at ${endpoint}:`, endpointError.message); lastError = endpointError; + endpointIndex++; } } @@ -219,7 +328,8 @@ export async function sendMessage(anthropicRequest, accountManager, fallbackEnab } catch (error) { if (isRateLimitError(error)) { - // Rate limited - already marked, continue to next account + // Rate limited - already marked, notify strategy and continue to next account + accountManager.notifyRateLimit(account, model); logger.info(`[CloudCode] Account ${account.email} rate-limited, trying next...`); continue; } @@ -230,15 +340,31 @@ export async function sendMessage(anthropicRequest, accountManager, fallbackEnab } // Handle 5xx errors if (error.message.includes('API error 5') || error.message.includes('500') || error.message.includes('503')) { - logger.warn(`[CloudCode] Account ${account.email} failed with 5xx error, trying next...`); - accountManager.pickNext(model); + accountManager.notifyFailure(account, model); + + // Gap 2: Check consecutive failures for extended cooldown + const consecutiveFailures = accountManager.getHealthTracker()?.getConsecutiveFailures(account.email) || 0; + if (consecutiveFailures >= MAX_CONSECUTIVE_FAILURES) { + logger.warn(`[CloudCode] Account ${account.email} has ${consecutiveFailures} consecutive failures, applying extended cooldown (${formatDuration(EXTENDED_COOLDOWN_MS)})`); + accountManager.markRateLimited(account.email, EXTENDED_COOLDOWN_MS, model); + } else { + logger.warn(`[CloudCode] Account ${account.email} failed with 5xx error, trying next...`); + } continue; } if (isNetworkError(error)) { - logger.warn(`[CloudCode] Network error for ${account.email}, trying next account... (${error.message})`); + accountManager.notifyFailure(account, model); + + // Gap 2: Check consecutive failures for extended cooldown + const consecutiveFailures = accountManager.getHealthTracker()?.getConsecutiveFailures(account.email) || 0; + if (consecutiveFailures >= MAX_CONSECUTIVE_FAILURES) { + logger.warn(`[CloudCode] Account ${account.email} has ${consecutiveFailures} consecutive network failures, applying extended cooldown (${formatDuration(EXTENDED_COOLDOWN_MS)})`); + accountManager.markRateLimited(account.email, EXTENDED_COOLDOWN_MS, model); + } else { + logger.warn(`[CloudCode] Network error for ${account.email}, trying next account... (${error.message})`); + } await sleep(1000); - accountManager.pickNext(model); continue; } diff --git a/src/cloudcode/streaming-handler.js b/src/cloudcode/streaming-handler.js index f5d76f0..df74a68 100644 --- a/src/cloudcode/streaming-handler.js +++ b/src/cloudcode/streaming-handler.js @@ -10,7 +10,12 @@ import { MAX_RETRIES, MAX_EMPTY_RESPONSE_RETRIES, MAX_WAIT_BEFORE_ERROR_MS, - DEFAULT_COOLDOWN_MS + DEFAULT_COOLDOWN_MS, + RATE_LIMIT_DEDUP_WINDOW_MS, + MAX_CONSECUTIVE_FAILURES, + EXTENDED_COOLDOWN_MS, + CAPACITY_RETRY_DELAY_MS, + MAX_CAPACITY_RETRIES } from '../constants.js'; import { isRateLimitError, isAuthError, isEmptyResponseError } from '../errors.js'; import { formatDuration, sleep, isNetworkError } from '../utils/helpers.js'; @@ -21,6 +26,83 @@ import { streamSSEResponse } from './sse-streamer.js'; import { getFallbackModel } from '../fallback-config.js'; import crypto from 'crypto'; +/** + * Gap 1: Rate limit deduplication - prevents thundering herd on concurrent rate limits + * Tracks last rate limit timestamp per model to skip duplicate retries + */ +const lastRateLimitTimestamps = new Map(); // modelId -> timestamp + +/** + * Check if we should skip retry due to recent rate limit on this model + * @param {string} model - Model ID + * @returns {boolean} True if retry should be skipped (within dedup window) + */ +function shouldSkipRetryDueToDedup(model) { + const lastTimestamp = lastRateLimitTimestamps.get(model); + if (!lastTimestamp) return false; + + const elapsed = Date.now() - lastTimestamp; + if (elapsed < RATE_LIMIT_DEDUP_WINDOW_MS) { + logger.debug(`[CloudCode] Rate limit on ${model} within dedup window (${elapsed}ms ago), skipping retry`); + return true; + } + return false; +} + +/** + * Record rate limit timestamp for deduplication + * @param {string} model - Model ID + */ +function recordRateLimitTimestamp(model) { + lastRateLimitTimestamps.set(model, Date.now()); +} + +/** + * Clear rate limit timestamp after successful retry + * @param {string} model - Model ID + */ +function clearRateLimitTimestamp(model) { + lastRateLimitTimestamps.delete(model); +} + +/** + * Gap 3: Detect permanent authentication failures that require re-authentication + * @param {string} errorText - Error message from API + * @returns {boolean} True if permanent auth failure + */ +function isPermanentAuthFailure(errorText) { + const lower = (errorText || '').toLowerCase(); + return lower.includes('invalid_grant') || + lower.includes('token revoked') || + lower.includes('token has been expired or revoked') || + lower.includes('token_revoked') || + lower.includes('invalid_client') || + lower.includes('credentials are invalid'); +} + +/** + * Gap 4: Detect if 429 error is due to model capacity (not user quota) + * @param {string} errorText - Error message from API + * @returns {boolean} True if capacity exhausted (not quota) + */ +function isModelCapacityExhausted(errorText) { + const lower = (errorText || '').toLowerCase(); + return lower.includes('model_capacity_exhausted') || + lower.includes('capacity_exhausted') || + lower.includes('model is currently overloaded') || + lower.includes('service temporarily unavailable'); +} + +// Periodically clean up stale dedup timestamps (every 60 seconds) +setInterval(() => { + const cutoff = Date.now() - 60000; // 1 minute + for (const [model, timestamp] of lastRateLimitTimestamps.entries()) { + if (timestamp < cutoff) { + lastRateLimitTimestamps.delete(model); + } + } +}, 60000); + /** * Send a streaming request to Cloud Code with multi-account support * Streams events in real-time as they arrive from the server @@ -83,10 +165,14 @@ export async function* sendMessageStream(anthropicRequest, accountManager, fallb throw new Error('No accounts available'); } - // Pick sticky account (prefers current for cache continuity) - let account = accountManager.getCurrentStickyAccount(model); - if (!account) { - account = accountManager.pickNext(model); + // Select account using configured strategy + const { account, waitMs } = accountManager.selectAccount(model); + + // If strategy returns a wait time, sleep and retry + if (!account && waitMs > 0) { + logger.info(`[CloudCode] Waiting ${formatDuration(waitMs)} for account...`); + await sleep(waitMs + 500); + continue; } if (!account) { @@ -101,11 +187,14 @@ export async function* sendMessageStream(anthropicRequest, accountManager, fallb logger.debug(`[CloudCode] Starting stream for model: ${model}`); - // Try each endpoint for streaming + // Try each endpoint with index-based loop for capacity retry support let lastError = null; let retriedOnce = false; // Track if we've already retried for short rate limit + let capacityRetryCount = 0; // Gap 4: Track capacity exhaustion retries + let endpointIndex = 0; - for (const endpoint of ANTIGRAVITY_ENDPOINT_FALLBACKS) { + while (endpointIndex < ANTIGRAVITY_ENDPOINT_FALLBACKS.length) { + const endpoint = ANTIGRAVITY_ENDPOINT_FALLBACKS[endpointIndex]; try { const url = `${endpoint}/v1internal:streamGenerateContent?alt=sse`; @@ -120,15 +209,44 @@ export async function* sendMessageStream(anthropicRequest, accountManager, fallb logger.warn(`[CloudCode] Stream error at ${endpoint}: ${response.status} - ${errorText}`); if (response.status === 401) { - // Auth error - clear caches and retry + // Gap 3: Check for permanent auth failures + if (isPermanentAuthFailure(errorText)) { + logger.error(`[CloudCode] Permanent auth failure for ${account.email}: ${errorText.substring(0, 100)}`); + accountManager.markInvalid(account.email, 'Token revoked - re-authentication required'); + throw new Error(`AUTH_INVALID_PERMANENT: ${errorText}`); + } + + // Transient auth error - clear caches and retry accountManager.clearTokenCache(account.email); accountManager.clearProjectCache(account.email); + endpointIndex++; continue; } if (response.status === 429) { const resetMs = parseResetTime(response, errorText); + // Gap 4: Check if capacity issue (NOT quota) - retry SAME endpoint + if (isModelCapacityExhausted(errorText)) { + if (capacityRetryCount < MAX_CAPACITY_RETRIES) { + capacityRetryCount++; + const waitMs = resetMs || CAPACITY_RETRY_DELAY_MS; + logger.info(`[CloudCode] Model capacity exhausted, retry ${capacityRetryCount}/${MAX_CAPACITY_RETRIES} after ${formatDuration(waitMs)}...`); + await sleep(waitMs); + // Don't increment endpointIndex - retry same endpoint + continue; + } + // Max capacity retries exceeded - treat as quota exhaustion + logger.warn(`[CloudCode] Max capacity retries (${MAX_CAPACITY_RETRIES}) exceeded, switching account`); + } + + // Gap 1: Check deduplication window to prevent thundering herd + if (shouldSkipRetryDueToDedup(model)) { + logger.info(`[CloudCode] Skipping retry due to recent rate limit, switching account...`); + accountManager.markRateLimited(account.email, resetMs || DEFAULT_COOLDOWN_MS, model); + throw new Error(`RATE_LIMITED_DEDUP: ${errorText}`); + } + // Decision: wait and retry OR switch account if (resetMs && resetMs > DEFAULT_COOLDOWN_MS) { // Long-term quota exhaustion (> 10s) - switch to next account @@ -141,28 +259,11 @@ export async function* sendMessageStream(anthropicRequest, accountManager, fallb if (!retriedOnce) { retriedOnce = true; + recordRateLimitTimestamp(model); // Gap 1: Record before retry logger.info(`[CloudCode] Short rate limit (${formatDuration(waitMs)}), waiting and retrying...`); await sleep(waitMs); - // Retry same endpoint - const retryResponse = await fetch(url, { - method: 'POST', - headers: buildHeaders(token, model, 'text/event-stream'), - body: JSON.stringify(payload) - }); - - if (retryResponse.ok) { - // Stream the retry response - yield* streamSSEResponse(retryResponse, anthropicRequest.model); - logger.debug('[CloudCode] Stream completed after retry'); - return; - } - - // Retry also failed - parse new reset time - const retryErrorText = await retryResponse.text(); - const retryResetMs = parseResetTime(retryResponse, retryErrorText); - logger.warn(`[CloudCode] Retry also failed, marking and switching...`); - accountManager.markRateLimited(account.email, retryResetMs || waitMs, model); - throw new Error(`RATE_LIMITED_AFTER_RETRY: ${retryErrorText}`); + // Don't increment endpointIndex - retry same endpoint + continue; } else { // Already retried once, mark and switch accountManager.markRateLimited(account.email, waitMs, model); @@ -179,6 +280,7 @@ export async function* sendMessageStream(anthropicRequest, accountManager, fallb await sleep(1000); } + endpointIndex++; continue; } @@ -189,6 +291,9 @@ export async function* sendMessageStream(anthropicRequest, accountManager, fallb try { yield* streamSSEResponse(currentResponse, anthropicRequest.model); logger.debug('[CloudCode] Stream completed'); + // Gap 1: Clear timestamp on success + clearRateLimitTimestamp(model); + accountManager.notifySuccess(account, model); return; } catch (streamError) { // Only retry on EmptyResponseError @@ -226,8 +331,13 @@ export async function* sendMessageStream(anthropicRequest, accountManager, fallb throw new Error(`429 RESOURCE_EXHAUSTED during retry: ${retryErrorText}`); } - // Auth error - clear caches and throw with recognizable message + // Auth error - check for permanent failure if (currentResponse.status === 401) { + if (isPermanentAuthFailure(retryErrorText)) { + logger.error(`[CloudCode] Permanent auth failure during retry for ${account.email}`); + accountManager.markInvalid(account.email, 'Token revoked - re-authentication required'); + throw new Error(`AUTH_INVALID_PERMANENT: ${retryErrorText}`); + } accountManager.clearTokenCache(account.email); accountManager.clearProjectCache(account.email); throw new Error(`401 AUTH_INVALID during retry: ${retryErrorText}`); @@ -261,6 +371,7 @@ export async function* sendMessageStream(anthropicRequest, accountManager, fallb } logger.warn(`[CloudCode] Stream error at ${endpoint}:`, endpointError.message); lastError = endpointError; + endpointIndex++; } } @@ -276,7 +387,8 @@ export async function* sendMessageStream(anthropicRequest, accountManager, fallb } catch (error) { if (isRateLimitError(error)) { - // Rate limited - already marked, continue to next account + // Rate limited - already marked, notify strategy and continue to next account + accountManager.notifyRateLimit(account, model); logger.info(`[CloudCode] Account ${account.email} rate-limited, trying next...`); continue; } @@ -287,15 +399,31 @@ export async function* sendMessageStream(anthropicRequest, accountManager, fallb } // Handle 5xx errors if (error.message.includes('API error 5') || error.message.includes('500') || error.message.includes('503')) { - logger.warn(`[CloudCode] Account ${account.email} failed with 5xx stream error, trying next...`); - accountManager.pickNext(model); + accountManager.notifyFailure(account, model); + + // Gap 2: Check consecutive failures for extended cooldown + const consecutiveFailures = accountManager.getHealthTracker()?.getConsecutiveFailures(account.email) || 0; + if (consecutiveFailures >= MAX_CONSECUTIVE_FAILURES) { + logger.warn(`[CloudCode] Account ${account.email} has ${consecutiveFailures} consecutive failures, applying extended cooldown (${formatDuration(EXTENDED_COOLDOWN_MS)})`); + accountManager.markRateLimited(account.email, EXTENDED_COOLDOWN_MS, model); + } else { + logger.warn(`[CloudCode] Account ${account.email} failed with 5xx stream error, trying next...`); + } continue; } if (isNetworkError(error)) { - logger.warn(`[CloudCode] Network error for ${account.email} (stream), trying next account... (${error.message})`); + accountManager.notifyFailure(account, model); + + // Gap 2: Check consecutive failures for extended cooldown + const consecutiveFailures = accountManager.getHealthTracker()?.getConsecutiveFailures(account.email) || 0; + if (consecutiveFailures >= MAX_CONSECUTIVE_FAILURES) { + logger.warn(`[CloudCode] Account ${account.email} has ${consecutiveFailures} consecutive network failures, applying extended cooldown (${formatDuration(EXTENDED_COOLDOWN_MS)})`); + accountManager.markRateLimited(account.email, EXTENDED_COOLDOWN_MS, model); + } else { + logger.warn(`[CloudCode] Network error for ${account.email} (stream), trying next account... (${error.message})`); + } await sleep(1000); - accountManager.pickNext(model); continue; } diff --git a/src/config.js b/src/config.js index 3aa3109..a05b3c2 100644 --- a/src/config.js +++ b/src/config.js @@ -15,7 +15,26 @@ const DEFAULT_CONFIG = { persistTokenCache: false, defaultCooldownMs: 10000, // 10 seconds maxWaitBeforeErrorMs: 120000, // 2 minutes - modelMapping: {} + modelMapping: {}, + // Account selection strategy configuration + accountSelection: { + strategy: 'hybrid', // 'sticky' | 'round-robin' | 'hybrid' + // Hybrid strategy tuning (optional - sensible defaults) + healthScore: { + initial: 70, // Starting score for new accounts + successReward: 1, // Points on successful request + rateLimitPenalty: -10, // Points on rate limit + failurePenalty: -20, // Points on other failures + recoveryPerHour: 2, // Passive recovery rate + minUsable: 50, // Minimum score to be selected + maxScore: 100 // Maximum score cap + }, + tokenBucket: { + maxTokens: 50, // Maximum token capacity + tokensPerMinute: 6, // Regeneration rate + initialTokens: 50 // Starting tokens + } + } }; // Config locations diff --git a/src/constants.js b/src/constants.js index 4eb7a30..e8136e3 100644 --- a/src/constants.js +++ b/src/constants.js @@ -103,9 +103,24 @@ export const MAX_ACCOUNTS = config?.maxAccounts || 10; // From config or 10 // Rate limit wait thresholds export const MAX_WAIT_BEFORE_ERROR_MS = config?.maxWaitBeforeErrorMs || 120000; // From config or 2 minutes +// Gap 1: Retry deduplication - prevents thundering herd on concurrent rate limits +export const RATE_LIMIT_DEDUP_WINDOW_MS = config?.rateLimitDedupWindowMs || 5000; // 5 seconds + +// Gap 2: Consecutive failure tracking - extended cooldown after repeated failures +export const MAX_CONSECUTIVE_FAILURES = config?.maxConsecutiveFailures || 3; +export const EXTENDED_COOLDOWN_MS = config?.extendedCooldownMs || 60000; // 1 minute + +// Gap 4: Capacity exhaustion - shorter retry for model capacity issues (not quota) +export const CAPACITY_RETRY_DELAY_MS = config?.capacityRetryDelayMs || 2000; // 2 seconds +export const MAX_CAPACITY_RETRIES = config?.maxCapacityRetries || 3; + // Thinking model constants export const MIN_SIGNATURE_LENGTH = 50; // Minimum valid thinking signature length +// Account selection strategies +export const SELECTION_STRATEGIES = ['sticky', 'round-robin', 'hybrid']; +export const DEFAULT_SELECTION_STRATEGY = 'hybrid'; + // Gemini-specific limits export const GEMINI_MAX_OUTPUT_TOKENS = 16384; @@ -235,6 +250,11 @@ export default { MAX_EMPTY_RESPONSE_RETRIES, MAX_ACCOUNTS, MAX_WAIT_BEFORE_ERROR_MS, + RATE_LIMIT_DEDUP_WINDOW_MS, + MAX_CONSECUTIVE_FAILURES, + EXTENDED_COOLDOWN_MS, + CAPACITY_RETRY_DELAY_MS, + MAX_CAPACITY_RETRIES, MIN_SIGNATURE_LENGTH, GEMINI_MAX_OUTPUT_TOKENS, GEMINI_SKIP_SIGNATURE, diff --git a/src/errors.js b/src/errors.js index fac9410..c5a14f2 100644 --- a/src/errors.js +++ b/src/errors.js @@ -149,6 +149,23 @@ export class EmptyResponseError extends AntigravityError { } } +/** + * Capacity exhausted error - Google's model is at capacity (not user quota) + * Should retry on same account with shorter delay, not switch accounts immediately + * Different from QUOTA_EXHAUSTED which indicates user's daily/hourly limit + */ +export class CapacityExhaustedError extends AntigravityError { + /** + * @param {string} message - Error message + * @param {number|null} retryAfterMs - Suggested retry delay in ms + */ + constructor(message = 'Model capacity exhausted', retryAfterMs = null) { + super(message, 'CAPACITY_EXHAUSTED', true, { retryAfterMs }); + this.name = 'CapacityExhaustedError'; + this.retryAfterMs = retryAfterMs; + } +} + /** * Check if an error is a rate limit error * Works with both custom error classes and legacy string-based errors @@ -188,6 +205,22 @@ export function isEmptyResponseError(error) { error?.name === 'EmptyResponseError'; } +/** + * Check if an error is a capacity exhausted error (model overload, not user quota) + * This is different from quota exhaustion - capacity issues are temporary infrastructure + * limits that should be retried on the SAME account with shorter delays + * @param {Error} error - Error to check + * @returns {boolean} + */ +export function isCapacityExhaustedError(error) { + if (error instanceof CapacityExhaustedError) return true; + const msg = (error.message || '').toLowerCase(); + return msg.includes('model_capacity_exhausted') || + msg.includes('capacity_exhausted') || + msg.includes('model is currently overloaded') || + msg.includes('service temporarily unavailable'); +} + export default { AntigravityError, RateLimitError, @@ -197,7 +230,9 @@ export default { ApiError, NativeModuleError, EmptyResponseError, + CapacityExhaustedError, isRateLimitError, isAuthError, - isEmptyResponseError + isEmptyResponseError, + isCapacityExhaustedError }; diff --git a/src/index.js b/src/index.js index f17e9f6..8864766 100644 --- a/src/index.js +++ b/src/index.js @@ -3,9 +3,10 @@ * Entry point - starts the proxy server */ -import app from './server.js'; +import app, { accountManager } from './server.js'; import { DEFAULT_PORT } from './constants.js'; import { logger } from './utils/logger.js'; +import { getStrategyLabel, STRATEGY_NAMES, DEFAULT_STRATEGY } from './account-manager/strategies/index.js'; import path from 'path'; import os from 'os'; @@ -14,6 +15,21 @@ const args = process.argv.slice(2); const isDebug = args.includes('--debug') || process.env.DEBUG === 'true'; const isFallbackEnabled = args.includes('--fallback') || process.env.FALLBACK === 'true'; +// Parse --strategy flag (format: --strategy=sticky or --strategy sticky) +let strategyOverride = null; +for (let i = 0; i < args.length; i++) { + if (args[i].startsWith('--strategy=')) { + strategyOverride = args[i].split('=')[1]; + } else if (args[i] === '--strategy' && args[i + 1]) { + strategyOverride = args[i + 1]; + } +} +// Validate strategy +if (strategyOverride && !STRATEGY_NAMES.includes(strategyOverride.toLowerCase())) { + logger.warn(`[Startup] Invalid strategy "${strategyOverride}". Valid options: ${STRATEGY_NAMES.join(', ')}. Using default.`); + strategyOverride = null; +} + // Initialize logger logger.setDebug(isDebug); @@ -45,6 +61,7 @@ const server = app.listen(PORT, () => { // Build Control section dynamically let controlSection = '║ Control: ║\n'; + controlSection += '║ --strategy= Set selection strategy (sticky/hybrid) ║\n'; if (!isDebug) { controlSection += '║ --debug Enable debug logging ║\n'; } @@ -53,17 +70,18 @@ const server = app.listen(PORT, () => { } controlSection += '║ Ctrl+C Stop server ║'; - // Build status section if any modes are active - let statusSection = ''; - if (isDebug || isFallbackEnabled) { - statusSection = '║ ║\n'; - statusSection += '║ Active Modes: ║\n'; - if (isDebug) { - statusSection += '║ ✓ Debug mode enabled ║\n'; - } - if (isFallbackEnabled) { - statusSection += '║ ✓ Model fallback enabled ║\n'; - } + // Get the strategy label (accountManager will be initialized by now) + const strategyLabel = accountManager.getStrategyLabel(); + + // Build status section - always show strategy, plus any active modes + let statusSection = '║ ║\n'; + statusSection += '║ Active Modes: ║\n'; + statusSection += `${border} ${align4(`✓ Strategy: ${strategyLabel}`)}${border}\n`; + if (isDebug) { + statusSection += '║ ✓ Debug mode enabled ║\n'; + } + if (isFallbackEnabled) { + statusSection += '║ ✓ Model fallback enabled ║\n'; } logger.log(` diff --git a/src/server.js b/src/server.js index 827eb58..ef5eb6c 100644 --- a/src/server.js +++ b/src/server.js @@ -26,13 +26,23 @@ import usageStats from './modules/usage-stats.js'; const args = process.argv.slice(2); const FALLBACK_ENABLED = args.includes('--fallback') || process.env.FALLBACK === 'true'; +// Parse --strategy flag (format: --strategy=sticky or --strategy sticky) +let STRATEGY_OVERRIDE = null; +for (let i = 0; i < args.length; i++) { + if (args[i].startsWith('--strategy=')) { + STRATEGY_OVERRIDE = args[i].split('=')[1]; + } else if (args[i] === '--strategy' && args[i + 1]) { + STRATEGY_OVERRIDE = args[i + 1]; + } +} + const app = express(); // Disable x-powered-by header for security app.disable('x-powered-by'); // Initialize account manager (will be fully initialized on first request or startup) -const accountManager = new AccountManager(); +export const accountManager = new AccountManager(); // Track initialization status let isInitialized = false; @@ -50,7 +60,7 @@ async function ensureInitialized() { initPromise = (async () => { try { - await accountManager.initialize(); + await accountManager.initialize(STRATEGY_OVERRIDE); isInitialized = true; const status = accountManager.getStatus(); logger.success(`[Server] Account pool initialized: ${status.summary}`); diff --git a/src/webui/index.js b/src/webui/index.js index 0386550..d449297 100644 --- a/src/webui/index.js +++ b/src/webui/index.js @@ -282,7 +282,7 @@ export function mountWebUI(app, dirname, accountManager) { */ app.post('/api/config', (req, res) => { try { - const { debug, logLevel, maxRetries, retryBaseMs, retryMaxMs, persistTokenCache, defaultCooldownMs, maxWaitBeforeErrorMs } = req.body; + const { debug, logLevel, maxRetries, retryBaseMs, retryMaxMs, persistTokenCache, defaultCooldownMs, maxWaitBeforeErrorMs, accountSelection } = req.body; // Only allow updating specific fields (security) const updates = {}; @@ -308,6 +308,16 @@ export function mountWebUI(app, dirname, accountManager) { if (typeof maxWaitBeforeErrorMs === 'number' && maxWaitBeforeErrorMs >= 0 && maxWaitBeforeErrorMs <= 600000) { updates.maxWaitBeforeErrorMs = maxWaitBeforeErrorMs; } + // Account selection strategy validation + if (accountSelection && typeof accountSelection === 'object') { + const validStrategies = ['sticky', 'round-robin', 'hybrid']; + if (accountSelection.strategy && validStrategies.includes(accountSelection.strategy)) { + updates.accountSelection = { + ...(config.accountSelection || {}), + strategy: accountSelection.strategy + }; + } + } if (Object.keys(updates).length === 0) { return res.status(400).json({ diff --git a/tests/run-all.cjs b/tests/run-all.cjs index 01256cd..b8bec8f 100644 --- a/tests/run-all.cjs +++ b/tests/run-all.cjs @@ -9,6 +9,7 @@ const { spawn } = require('child_process'); const path = require('path'); const tests = [ + { name: 'Account Selection Strategies', file: 'test-strategies.cjs' }, { name: 'Thinking Signatures', file: 'test-thinking-signatures.cjs' }, { name: 'Multi-turn Tools (Non-Streaming)', file: 'test-multiturn-thinking-tools.cjs' }, { name: 'Multi-turn Tools (Streaming)', file: 'test-multiturn-thinking-tools-streaming.cjs' }, diff --git a/tests/test-strategies.cjs b/tests/test-strategies.cjs new file mode 100644 index 0000000..0be8162 --- /dev/null +++ b/tests/test-strategies.cjs @@ -0,0 +1,795 @@ +/** + * Test Account Selection Strategies - Unit Tests + * + * Tests the strategy pattern implementation for account selection: + * - HealthTracker: health score tracking with passive recovery + * - TokenBucketTracker: token bucket rate limiting + * - StickyStrategy: cache-optimized sticky selection + * - RoundRobinStrategy: load-balanced rotation + * - HybridStrategy: smart multi-signal distribution + * - Strategy Factory: createStrategy, isValidStrategy, getStrategyLabel + */ + +// Since we're in CommonJS and the module is ESM, we need to use dynamic import +async function runTests() { + console.log('╔══════════════════════════════════════════════════════════════╗'); + console.log('║ ACCOUNT SELECTION STRATEGY TEST SUITE ║'); + console.log('╚══════════════════════════════════════════════════════════════╝\n'); + + // Dynamic imports for ESM modules + const { HealthTracker } = await import('../src/account-manager/strategies/trackers/health-tracker.js'); + const { TokenBucketTracker } = await import('../src/account-manager/strategies/trackers/token-bucket-tracker.js'); + const { StickyStrategy } = await import('../src/account-manager/strategies/sticky-strategy.js'); + const { RoundRobinStrategy } = await import('../src/account-manager/strategies/round-robin-strategy.js'); + const { HybridStrategy } = await import('../src/account-manager/strategies/hybrid-strategy.js'); + const { BaseStrategy } = await import('../src/account-manager/strategies/base-strategy.js'); + const { + createStrategy, + isValidStrategy, + getStrategyLabel, + STRATEGY_NAMES, + DEFAULT_STRATEGY + } = await import('../src/account-manager/strategies/index.js'); + + let passed = 0; + let failed = 0; + + function test(name, fn) { + try { + fn(); + console.log(`✓ ${name}`); + passed++; + } catch (e) { + console.log(`✗ ${name}`); + console.log(` Error: ${e.message}`); + failed++; + } + } + + function assertEqual(actual, expected, message = '') { + if (actual !== expected) { + throw new Error(`${message}\nExpected: ${expected}\nActual: ${actual}`); + } + } + + function assertDeepEqual(actual, expected, message = '') { + if (JSON.stringify(actual) !== JSON.stringify(expected)) { + throw new Error(`${message}\nExpected: ${JSON.stringify(expected, null, 2)}\nActual: ${JSON.stringify(actual, null, 2)}`); + } + } + + function assertTrue(value, message = '') { + if (!value) { + throw new Error(message || 'Expected true but got false'); + } + } + + function assertFalse(value, message = '') { + if (value) { + throw new Error(message || 'Expected false but got true'); + } + } + + function assertNull(value, message = '') { + if (value !== null) { + throw new Error(`${message}\nExpected null but got: ${value}`); + } + } + + function assertWithin(actual, min, max, message = '') { + if (actual < min || actual > max) { + throw new Error(`${message}\nExpected value between ${min} and ${max}, got: ${actual}`); + } + } + + // Helper to create mock accounts + function createMockAccounts(count = 3, options = {}) { + return Array.from({ length: count }, (_, i) => ({ + email: `account${i + 1}@example.com`, + enabled: true, + isInvalid: false, + lastUsed: Date.now() - (i * 60000), // Stagger by 1 minute + modelRateLimits: {}, + ...options + })); + } + + // ========================================================================== + // HEALTH TRACKER TESTS + // ========================================================================== + console.log('\n─── HealthTracker Tests ───'); + + test('HealthTracker: initial score is 70 by default', () => { + const tracker = new HealthTracker(); + const score = tracker.getScore('new@example.com'); + assertEqual(score, 70, 'Default initial score should be 70'); + }); + + test('HealthTracker: custom initial score', () => { + const tracker = new HealthTracker({ initial: 80 }); + const score = tracker.getScore('new@example.com'); + assertEqual(score, 80, 'Custom initial score should be 80'); + }); + + test('HealthTracker: recordSuccess increases score', () => { + const tracker = new HealthTracker({ initial: 70, successReward: 1 }); + tracker.recordSuccess('test@example.com'); + const score = tracker.getScore('test@example.com'); + assertEqual(score, 71, 'Score should increase by 1 on success'); + }); + + test('HealthTracker: recordRateLimit decreases score', () => { + const tracker = new HealthTracker({ initial: 70, rateLimitPenalty: -10 }); + tracker.recordRateLimit('test@example.com'); + const score = tracker.getScore('test@example.com'); + assertEqual(score, 60, 'Score should decrease by 10 on rate limit'); + }); + + test('HealthTracker: recordFailure decreases score', () => { + const tracker = new HealthTracker({ initial: 70, failurePenalty: -20 }); + tracker.recordFailure('test@example.com'); + const score = tracker.getScore('test@example.com'); + assertEqual(score, 50, 'Score should decrease by 20 on failure'); + }); + + test('HealthTracker: score cannot exceed maxScore', () => { + const tracker = new HealthTracker({ initial: 99, maxScore: 100, successReward: 5 }); + tracker.recordSuccess('test@example.com'); + const score = tracker.getScore('test@example.com'); + assertEqual(score, 100, 'Score should be capped at maxScore'); + }); + + test('HealthTracker: score cannot go below 0', () => { + const tracker = new HealthTracker({ initial: 10, failurePenalty: -50 }); + tracker.recordFailure('test@example.com'); + const score = tracker.getScore('test@example.com'); + assertEqual(score, 0, 'Score should not go below 0'); + }); + + test('HealthTracker: isUsable returns true when score >= minUsable', () => { + const tracker = new HealthTracker({ initial: 50, minUsable: 50 }); + assertTrue(tracker.isUsable('test@example.com'), 'Should be usable at minUsable'); + }); + + test('HealthTracker: isUsable returns false when score < minUsable', () => { + const tracker = new HealthTracker({ initial: 49, minUsable: 50 }); + assertFalse(tracker.isUsable('test@example.com'), 'Should not be usable below minUsable'); + }); + + test('HealthTracker: reset restores initial score', () => { + const tracker = new HealthTracker({ initial: 70 }); + tracker.recordFailure('test@example.com'); // Score drops + tracker.reset('test@example.com'); + const score = tracker.getScore('test@example.com'); + assertEqual(score, 70, 'Reset should restore initial score'); + }); + + test('HealthTracker: clear removes all scores', () => { + const tracker = new HealthTracker({ initial: 70 }); + tracker.recordSuccess('a@example.com'); + tracker.recordSuccess('b@example.com'); + tracker.clear(); + // After clear, new accounts should get initial score + assertEqual(tracker.getScore('a@example.com'), 70); + assertEqual(tracker.getScore('b@example.com'), 70); + }); + + test('HealthTracker: getConsecutiveFailures returns 0 for new account', () => { + const tracker = new HealthTracker(); + assertEqual(tracker.getConsecutiveFailures('new@example.com'), 0); + }); + + test('HealthTracker: recordRateLimit increments consecutiveFailures', () => { + const tracker = new HealthTracker(); + tracker.recordRateLimit('test@example.com'); + assertEqual(tracker.getConsecutiveFailures('test@example.com'), 1); + tracker.recordRateLimit('test@example.com'); + assertEqual(tracker.getConsecutiveFailures('test@example.com'), 2); + }); + + test('HealthTracker: recordFailure increments consecutiveFailures', () => { + const tracker = new HealthTracker(); + tracker.recordFailure('test@example.com'); + assertEqual(tracker.getConsecutiveFailures('test@example.com'), 1); + }); + + test('HealthTracker: recordSuccess resets consecutiveFailures', () => { + const tracker = new HealthTracker(); + tracker.recordRateLimit('test@example.com'); + tracker.recordRateLimit('test@example.com'); + assertEqual(tracker.getConsecutiveFailures('test@example.com'), 2); + tracker.recordSuccess('test@example.com'); + assertEqual(tracker.getConsecutiveFailures('test@example.com'), 0); + }); + + test('HealthTracker: reset clears consecutiveFailures', () => { + const tracker = new HealthTracker(); + tracker.recordFailure('test@example.com'); + tracker.recordFailure('test@example.com'); + assertEqual(tracker.getConsecutiveFailures('test@example.com'), 2); + tracker.reset('test@example.com'); + assertEqual(tracker.getConsecutiveFailures('test@example.com'), 0); + }); + + // ========================================================================== + // TOKEN BUCKET TRACKER TESTS + // ========================================================================== + console.log('\n─── TokenBucketTracker Tests ───'); + + test('TokenBucketTracker: initial tokens is 50 by default', () => { + const tracker = new TokenBucketTracker(); + const tokens = tracker.getTokens('new@example.com'); + assertEqual(tokens, 50, 'Default initial tokens should be 50'); + }); + + test('TokenBucketTracker: custom initial tokens', () => { + const tracker = new TokenBucketTracker({ initialTokens: 30 }); + const tokens = tracker.getTokens('new@example.com'); + assertEqual(tokens, 30, 'Custom initial tokens should be 30'); + }); + + test('TokenBucketTracker: consume decreases tokens', () => { + const tracker = new TokenBucketTracker({ initialTokens: 10, maxTokens: 10 }); + const consumed = tracker.consume('test@example.com'); + assertTrue(consumed, 'Consume should return true'); + assertEqual(tracker.getTokens('test@example.com'), 9, 'Tokens should decrease by 1'); + }); + + test('TokenBucketTracker: consume fails when no tokens', () => { + const tracker = new TokenBucketTracker({ initialTokens: 0, maxTokens: 10 }); + const consumed = tracker.consume('test@example.com'); + assertFalse(consumed, 'Consume should return false when no tokens'); + }); + + test('TokenBucketTracker: hasTokens returns true when tokens > 0', () => { + const tracker = new TokenBucketTracker({ initialTokens: 1 }); + assertTrue(tracker.hasTokens('test@example.com'), 'Should have tokens'); + }); + + test('TokenBucketTracker: hasTokens returns false when tokens < 1', () => { + const tracker = new TokenBucketTracker({ initialTokens: 0 }); + assertFalse(tracker.hasTokens('test@example.com'), 'Should not have tokens'); + }); + + test('TokenBucketTracker: refund increases tokens', () => { + const tracker = new TokenBucketTracker({ initialTokens: 5, maxTokens: 10 }); + tracker.consume('test@example.com'); // 5 -> 4 + tracker.refund('test@example.com'); // 4 -> 5 + assertEqual(tracker.getTokens('test@example.com'), 5, 'Refund should restore token'); + }); + + test('TokenBucketTracker: refund cannot exceed maxTokens', () => { + const tracker = new TokenBucketTracker({ initialTokens: 10, maxTokens: 10 }); + tracker.refund('test@example.com'); + assertEqual(tracker.getTokens('test@example.com'), 10, 'Refund should not exceed max'); + }); + + test('TokenBucketTracker: getMaxTokens returns configured max', () => { + const tracker = new TokenBucketTracker({ maxTokens: 100 }); + assertEqual(tracker.getMaxTokens(), 100, 'getMaxTokens should return 100'); + }); + + test('TokenBucketTracker: reset restores initial tokens', () => { + const tracker = new TokenBucketTracker({ initialTokens: 50, maxTokens: 50 }); + tracker.consume('test@example.com'); + tracker.consume('test@example.com'); + tracker.reset('test@example.com'); + assertEqual(tracker.getTokens('test@example.com'), 50, 'Reset should restore initial'); + }); + + // ========================================================================== + // BASE STRATEGY TESTS + // ========================================================================== + console.log('\n─── BaseStrategy Tests ───'); + + test('BaseStrategy: cannot be instantiated directly', () => { + try { + new BaseStrategy(); + throw new Error('Should have thrown'); + } catch (e) { + assertTrue(e.message.includes('abstract'), 'Should throw abstract error'); + } + }); + + test('BaseStrategy: isAccountUsable returns false for null account', () => { + // Create a minimal subclass to test + class TestStrategy extends BaseStrategy { + selectAccount() { return { account: null, index: 0 }; } + } + const strategy = new TestStrategy(); + assertFalse(strategy.isAccountUsable(null, 'model'), 'Null account should not be usable'); + }); + + test('BaseStrategy: isAccountUsable returns false for invalid account', () => { + class TestStrategy extends BaseStrategy { + selectAccount() { return { account: null, index: 0 }; } + } + const strategy = new TestStrategy(); + const account = { email: 'test@example.com', isInvalid: true }; + assertFalse(strategy.isAccountUsable(account, 'model'), 'Invalid account should not be usable'); + }); + + test('BaseStrategy: isAccountUsable returns false for disabled account', () => { + class TestStrategy extends BaseStrategy { + selectAccount() { return { account: null, index: 0 }; } + } + const strategy = new TestStrategy(); + const account = { email: 'test@example.com', enabled: false }; + assertFalse(strategy.isAccountUsable(account, 'model'), 'Disabled account should not be usable'); + }); + + test('BaseStrategy: isAccountUsable returns false for rate-limited model', () => { + class TestStrategy extends BaseStrategy { + selectAccount() { return { account: null, index: 0 }; } + } + const strategy = new TestStrategy(); + const account = { + email: 'test@example.com', + modelRateLimits: { + 'claude-sonnet': { + isRateLimited: true, + resetTime: Date.now() + 60000 // 1 minute in future + } + } + }; + assertFalse(strategy.isAccountUsable(account, 'claude-sonnet'), 'Rate-limited model should not be usable'); + }); + + test('BaseStrategy: isAccountUsable returns true for expired rate limit', () => { + class TestStrategy extends BaseStrategy { + selectAccount() { return { account: null, index: 0 }; } + } + const strategy = new TestStrategy(); + const account = { + email: 'test@example.com', + modelRateLimits: { + 'claude-sonnet': { + isRateLimited: true, + resetTime: Date.now() - 1000 // 1 second in past + } + } + }; + assertTrue(strategy.isAccountUsable(account, 'claude-sonnet'), 'Expired rate limit should be usable'); + }); + + test('BaseStrategy: getUsableAccounts filters correctly', () => { + class TestStrategy extends BaseStrategy { + selectAccount() { return { account: null, index: 0 }; } + } + const strategy = new TestStrategy(); + const accounts = [ + { email: 'a@example.com', enabled: true }, + { email: 'b@example.com', enabled: false }, + { email: 'c@example.com', enabled: true, isInvalid: true }, + { email: 'd@example.com', enabled: true } + ]; + const usable = strategy.getUsableAccounts(accounts, 'model'); + assertEqual(usable.length, 2, 'Should have 2 usable accounts'); + assertEqual(usable[0].account.email, 'a@example.com'); + assertEqual(usable[1].account.email, 'd@example.com'); + }); + + // ========================================================================== + // STICKY STRATEGY TESTS + // ========================================================================== + console.log('\n─── StickyStrategy Tests ───'); + + test('StickyStrategy: returns null for empty accounts', () => { + const strategy = new StickyStrategy(); + const result = strategy.selectAccount([], 'model', { currentIndex: 0 }); + assertNull(result.account, 'Should return null for empty accounts'); + }); + + test('StickyStrategy: keeps using current account when available', () => { + const strategy = new StickyStrategy(); + const accounts = createMockAccounts(3); + + const result1 = strategy.selectAccount(accounts, 'model', { currentIndex: 0 }); + assertEqual(result1.account.email, 'account1@example.com'); + assertEqual(result1.index, 0); + + const result2 = strategy.selectAccount(accounts, 'model', { currentIndex: 0 }); + assertEqual(result2.account.email, 'account1@example.com', 'Should stick to same account'); + assertEqual(result2.index, 0); + }); + + test('StickyStrategy: switches when current account is rate-limited', () => { + const strategy = new StickyStrategy(); + const accounts = createMockAccounts(3); + // Rate-limit account1 for 5 minutes (longer than MAX_WAIT) + accounts[0].modelRateLimits = { + 'model': { isRateLimited: true, resetTime: Date.now() + 300000 } + }; + + const result = strategy.selectAccount(accounts, 'model', { currentIndex: 0 }); + assertEqual(result.account.email, 'account2@example.com', 'Should switch to next available'); + assertEqual(result.index, 1); + }); + + test('StickyStrategy: returns waitMs when current account has short rate limit', () => { + const strategy = new StickyStrategy(); + const accounts = createMockAccounts(1); // Only one account + // Rate-limit for 30 seconds (less than MAX_WAIT of 2 minutes) + accounts[0].modelRateLimits = { + 'model': { isRateLimited: true, resetTime: Date.now() + 30000 } + }; + + const result = strategy.selectAccount(accounts, 'model', { currentIndex: 0 }); + assertNull(result.account, 'Should return null when waiting'); + assertWithin(result.waitMs, 29000, 31000, 'Should return ~30s wait time'); + }); + + test('StickyStrategy: switches when current account is disabled', () => { + const strategy = new StickyStrategy(); + const accounts = createMockAccounts(3); + accounts[0].enabled = false; + + const result = strategy.selectAccount(accounts, 'model', { currentIndex: 0 }); + assertEqual(result.account.email, 'account2@example.com', 'Should switch to next'); + }); + + test('StickyStrategy: switches when current account is invalid', () => { + const strategy = new StickyStrategy(); + const accounts = createMockAccounts(3); + accounts[0].isInvalid = true; + + const result = strategy.selectAccount(accounts, 'model', { currentIndex: 0 }); + assertEqual(result.account.email, 'account2@example.com', 'Should switch to next'); + }); + + test('StickyStrategy: wraps around when at end of list', () => { + const strategy = new StickyStrategy(); + const accounts = createMockAccounts(3); + accounts[2].isInvalid = true; // Last account invalid + + const result = strategy.selectAccount(accounts, 'model', { currentIndex: 2 }); + assertEqual(result.account.email, 'account1@example.com', 'Should wrap to first'); + assertEqual(result.index, 0); + }); + + test('StickyStrategy: clamps invalid currentIndex', () => { + const strategy = new StickyStrategy(); + const accounts = createMockAccounts(3); + + const result = strategy.selectAccount(accounts, 'model', { currentIndex: 10 }); + assertEqual(result.account.email, 'account1@example.com', 'Should clamp to valid index'); + assertEqual(result.index, 0); + }); + + // ========================================================================== + // ROUND-ROBIN STRATEGY TESTS + // ========================================================================== + console.log('\n─── RoundRobinStrategy Tests ───'); + + test('RoundRobinStrategy: returns null for empty accounts', () => { + const strategy = new RoundRobinStrategy(); + const result = strategy.selectAccount([], 'model'); + assertNull(result.account, 'Should return null for empty accounts'); + }); + + test('RoundRobinStrategy: rotates through accounts', () => { + const strategy = new RoundRobinStrategy(); + const accounts = createMockAccounts(3); + + const r1 = strategy.selectAccount(accounts, 'model'); + const r2 = strategy.selectAccount(accounts, 'model'); + const r3 = strategy.selectAccount(accounts, 'model'); + const r4 = strategy.selectAccount(accounts, 'model'); + + // First call starts at cursor 0, looks at (0+1)%3 = 1 + // Then cursor becomes 1, next looks at (1+1)%3 = 2 + // Then cursor becomes 2, next looks at (2+1)%3 = 0 + // Then cursor becomes 0, next looks at (0+1)%3 = 1 + assertEqual(r1.account.email, 'account2@example.com', 'First should be account2'); + assertEqual(r2.account.email, 'account3@example.com', 'Second should be account3'); + assertEqual(r3.account.email, 'account1@example.com', 'Third should wrap to account1'); + assertEqual(r4.account.email, 'account2@example.com', 'Fourth should continue rotation'); + }); + + test('RoundRobinStrategy: skips unavailable accounts', () => { + const strategy = new RoundRobinStrategy(); + const accounts = createMockAccounts(3); + accounts[1].enabled = false; // Disable account2 + + const r1 = strategy.selectAccount(accounts, 'model'); + const r2 = strategy.selectAccount(accounts, 'model'); + const r3 = strategy.selectAccount(accounts, 'model'); + + // account2 is skipped + assertEqual(r1.account.email, 'account3@example.com'); + assertEqual(r2.account.email, 'account1@example.com'); + assertEqual(r3.account.email, 'account3@example.com'); + }); + + test('RoundRobinStrategy: returns null when all accounts unavailable', () => { + const strategy = new RoundRobinStrategy(); + const accounts = createMockAccounts(3); + accounts.forEach(a => a.enabled = false); + + const result = strategy.selectAccount(accounts, 'model'); + assertNull(result.account, 'Should return null when all unavailable'); + }); + + test('RoundRobinStrategy: resetCursor resets position', () => { + const strategy = new RoundRobinStrategy(); + const accounts = createMockAccounts(3); + + strategy.selectAccount(accounts, 'model'); // Moves cursor + strategy.selectAccount(accounts, 'model'); // Moves cursor + strategy.resetCursor(); + + const result = strategy.selectAccount(accounts, 'model'); + assertEqual(result.account.email, 'account2@example.com', 'Should start from beginning after reset'); + }); + + // ========================================================================== + // HYBRID STRATEGY TESTS + // ========================================================================== + console.log('\n─── HybridStrategy Tests ───'); + + test('HybridStrategy: returns null for empty accounts', () => { + const strategy = new HybridStrategy(); + const result = strategy.selectAccount([], 'model'); + assertNull(result.account, 'Should return null for empty accounts'); + }); + + test('HybridStrategy: selects best scored account', () => { + const strategy = new HybridStrategy({ + healthScore: { initial: 70 }, + tokenBucket: { initialTokens: 50, maxTokens: 50 } + }); + const accounts = createMockAccounts(3); + // Make account3 older (higher LRU score) + accounts[2].lastUsed = Date.now() - 3600000; // 1 hour ago + + const result = strategy.selectAccount(accounts, 'model'); + // account3 should win due to higher LRU score + assertEqual(result.account.email, 'account3@example.com', 'Oldest account should be selected'); + }); + + test('HybridStrategy: filters out unhealthy accounts', () => { + const strategy = new HybridStrategy({ + healthScore: { initial: 40, minUsable: 50 }, + tokenBucket: { initialTokens: 50, maxTokens: 50 } + }); + const accounts = createMockAccounts(3); + + // All accounts start with health 40, which is below minUsable 50 + const result = strategy.selectAccount(accounts, 'model'); + assertNull(result.account, 'Should filter all accounts with low health'); + }); + + test('HybridStrategy: filters out accounts without tokens', () => { + const strategy = new HybridStrategy({ + healthScore: { initial: 70 }, + tokenBucket: { initialTokens: 0, maxTokens: 50 } + }); + const accounts = createMockAccounts(3); + + const result = strategy.selectAccount(accounts, 'model'); + assertNull(result.account, 'Should filter all accounts without tokens'); + }); + + test('HybridStrategy: consumes token on selection', () => { + const strategy = new HybridStrategy({ + healthScore: { initial: 70 }, + tokenBucket: { initialTokens: 10, maxTokens: 50 } + }); + const accounts = createMockAccounts(1); + + strategy.selectAccount(accounts, 'model'); + const tracker = strategy.getTokenBucketTracker(); + assertEqual(tracker.getTokens(accounts[0].email), 9, 'Token should be consumed'); + }); + + test('HybridStrategy: onSuccess increases health', () => { + const strategy = new HybridStrategy({ + healthScore: { initial: 70, successReward: 5 } + }); + const account = { email: 'test@example.com' }; + + strategy.onSuccess(account, 'model'); + const tracker = strategy.getHealthTracker(); + assertEqual(tracker.getScore('test@example.com'), 75, 'Health should increase'); + }); + + test('HybridStrategy: onRateLimit decreases health', () => { + const strategy = new HybridStrategy({ + healthScore: { initial: 70, rateLimitPenalty: -10 } + }); + const account = { email: 'test@example.com' }; + + strategy.onRateLimit(account, 'model'); + const tracker = strategy.getHealthTracker(); + assertEqual(tracker.getScore('test@example.com'), 60, 'Health should decrease'); + }); + + test('HybridStrategy: onFailure decreases health and refunds token', () => { + const strategy = new HybridStrategy({ + healthScore: { initial: 70, failurePenalty: -20 }, + tokenBucket: { initialTokens: 10, maxTokens: 50 } + }); + const accounts = createMockAccounts(1); + + // First consume a token + strategy.selectAccount(accounts, 'model'); + const tokensBefore = strategy.getTokenBucketTracker().getTokens(accounts[0].email); + + // Then fail + strategy.onFailure(accounts[0], 'model'); + + const healthTracker = strategy.getHealthTracker(); + const tokenTracker = strategy.getTokenBucketTracker(); + + assertEqual(healthTracker.getScore(accounts[0].email), 50, 'Health should decrease by 20'); + assertEqual(tokenTracker.getTokens(accounts[0].email), tokensBefore + 1, 'Token should be refunded'); + }); + + test('HybridStrategy: scoring formula weights work correctly', () => { + // Test that health, tokens, and LRU all contribute to score + const strategy = new HybridStrategy({ + healthScore: { initial: 100 }, + tokenBucket: { initialTokens: 50, maxTokens: 50 }, + weights: { health: 2, tokens: 5, lru: 0.1 } + }); + + const accounts = [ + { email: 'high-health@example.com', enabled: true, lastUsed: Date.now() }, + { email: 'old-account@example.com', enabled: true, lastUsed: Date.now() - 3600000 } + ]; + + // Both have same health and tokens, but old-account has higher LRU + const result = strategy.selectAccount(accounts, 'model'); + assertEqual(result.account.email, 'old-account@example.com', 'Older account should win with LRU weight'); + }); + + // ========================================================================== + // STRATEGY FACTORY TESTS + // ========================================================================== + console.log('\n─── Strategy Factory Tests ───'); + + test('createStrategy: creates StickyStrategy for "sticky"', () => { + const strategy = createStrategy('sticky'); + assertTrue(strategy instanceof StickyStrategy, 'Should create StickyStrategy'); + }); + + test('createStrategy: creates RoundRobinStrategy for "round-robin"', () => { + const strategy = createStrategy('round-robin'); + assertTrue(strategy instanceof RoundRobinStrategy, 'Should create RoundRobinStrategy'); + }); + + test('createStrategy: creates RoundRobinStrategy for "roundrobin"', () => { + const strategy = createStrategy('roundrobin'); + assertTrue(strategy instanceof RoundRobinStrategy, 'Should accept roundrobin alias'); + }); + + test('createStrategy: creates HybridStrategy for "hybrid"', () => { + const strategy = createStrategy('hybrid'); + assertTrue(strategy instanceof HybridStrategy, 'Should create HybridStrategy'); + }); + + test('createStrategy: falls back to HybridStrategy for unknown strategy', () => { + const strategy = createStrategy('unknown'); + assertTrue(strategy instanceof HybridStrategy, 'Should fall back to HybridStrategy'); + }); + + test('createStrategy: uses default when null', () => { + const strategy = createStrategy(null); + assertTrue(strategy instanceof HybridStrategy, 'Null should use default HybridStrategy'); + }); + + test('createStrategy: is case-insensitive', () => { + const s1 = createStrategy('STICKY'); + const s2 = createStrategy('Hybrid'); + const s3 = createStrategy('ROUND-ROBIN'); + assertTrue(s1 instanceof StickyStrategy); + assertTrue(s2 instanceof HybridStrategy); + assertTrue(s3 instanceof RoundRobinStrategy); + }); + + test('isValidStrategy: returns true for valid strategies', () => { + assertTrue(isValidStrategy('sticky')); + assertTrue(isValidStrategy('round-robin')); + assertTrue(isValidStrategy('hybrid')); + assertTrue(isValidStrategy('roundrobin')); + }); + + test('isValidStrategy: returns false for invalid strategies', () => { + assertFalse(isValidStrategy('invalid')); + assertFalse(isValidStrategy('')); + assertFalse(isValidStrategy(null)); + assertFalse(isValidStrategy(undefined)); + }); + + test('getStrategyLabel: returns correct labels', () => { + assertEqual(getStrategyLabel('sticky'), 'Sticky (Cache Optimized)'); + assertEqual(getStrategyLabel('round-robin'), 'Round Robin (Load Balanced)'); + assertEqual(getStrategyLabel('roundrobin'), 'Round Robin (Load Balanced)'); + assertEqual(getStrategyLabel('hybrid'), 'Hybrid (Smart Distribution)'); + }); + + test('getStrategyLabel: returns default label for unknown', () => { + assertEqual(getStrategyLabel('unknown'), 'Hybrid (Smart Distribution)'); + assertEqual(getStrategyLabel(null), 'Hybrid (Smart Distribution)'); + }); + + test('STRATEGY_NAMES contains all valid strategies', () => { + assertDeepEqual(STRATEGY_NAMES, ['sticky', 'round-robin', 'hybrid']); + }); + + test('DEFAULT_STRATEGY is hybrid', () => { + assertEqual(DEFAULT_STRATEGY, 'hybrid'); + }); + + // ========================================================================== + // INTEGRATION TESTS + // ========================================================================== + console.log('\n─── Integration Tests ───'); + + test('Integration: Hybrid strategy recovers from rate limits', () => { + const strategy = new HybridStrategy({ + healthScore: { initial: 70, rateLimitPenalty: -10, minUsable: 50 } + }); + const accounts = createMockAccounts(2); + + // Rate limit first account multiple times + for (let i = 0; i < 3; i++) { + strategy.onRateLimit(accounts[0], 'model'); + } + + // Health of first account should be 40 (below minUsable 50) + const healthTracker = strategy.getHealthTracker(); + assertEqual(healthTracker.getScore(accounts[0].email), 40); + assertFalse(healthTracker.isUsable(accounts[0].email)); + + // Selection should prefer second account + const result = strategy.selectAccount(accounts, 'model'); + assertEqual(result.account.email, 'account2@example.com'); + }); + + test('Integration: Token consumption limits requests', () => { + const strategy = new HybridStrategy({ + tokenBucket: { initialTokens: 2, maxTokens: 10 } + }); + const accounts = createMockAccounts(1); + + // Consume all tokens + strategy.selectAccount(accounts, 'model'); // 2 -> 1 + strategy.selectAccount(accounts, 'model'); // 1 -> 0 + + // Third request should fail (no tokens) + const result = strategy.selectAccount(accounts, 'model'); + assertNull(result.account, 'Should return null when tokens exhausted'); + }); + + test('Integration: Multi-model rate limiting is independent', () => { + const strategy = new StickyStrategy(); + const accounts = createMockAccounts(2); + + // Rate limit account1 for model-a only + accounts[0].modelRateLimits = { + 'model-a': { isRateLimited: true, resetTime: Date.now() + 300000 } + }; + + // model-a should switch to account2 + const resultA = strategy.selectAccount(accounts, 'model-a', { currentIndex: 0 }); + assertEqual(resultA.account.email, 'account2@example.com'); + + // model-b should still use account1 + const resultB = strategy.selectAccount(accounts, 'model-b', { currentIndex: 0 }); + assertEqual(resultB.account.email, 'account1@example.com'); + }); + + // Summary + console.log('\n' + '═'.repeat(60)); + console.log(`Tests completed: ${passed} passed, ${failed} failed`); + + if (failed > 0) { + process.exit(1); + } +} + +runTests().catch(err => { + console.error('Test suite failed:', err); + process.exit(1); +});