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 <noreply@anthropic.com>
This commit is contained in:
Badri Narayanan S
50
CLAUDE.md
50
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`
|
||||
|
||||
37
README.md
37
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
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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] || '';
|
||||
}
|
||||
});
|
||||
|
||||
@@ -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
|
||||
};
|
||||
|
||||
/**
|
||||
|
||||
@@ -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}",
|
||||
|
||||
@@ -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",
|
||||
};
|
||||
|
||||
@@ -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",
|
||||
};
|
||||
|
||||
@@ -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",
|
||||
};
|
||||
|
||||
@@ -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: "选择了无效的策略",
|
||||
};
|
||||
|
||||
@@ -936,6 +936,35 @@
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- 🔀 Account Selection Strategy -->
|
||||
<div class="space-y-3">
|
||||
<div class="flex items-center gap-2 mb-1 px-1">
|
||||
<span class="text-[10px] uppercase text-gray-500 font-bold tracking-widest"
|
||||
x-text="$store.global.t('accountSelectionStrategy')">Account Selection Strategy</span>
|
||||
<div class="h-px flex-1 bg-space-border/30"></div>
|
||||
</div>
|
||||
|
||||
<div class="form-control view-card border-space-border/50 hover:border-neon-cyan/50">
|
||||
<div class="flex items-center justify-between">
|
||||
<div class="flex flex-col gap-1 flex-1">
|
||||
<span class="text-sm font-medium text-gray-200"
|
||||
x-text="$store.global.t('selectionStrategy')">Selection Strategy</span>
|
||||
<span class="text-[11px] text-gray-500"
|
||||
x-text="currentStrategyDescription()">How accounts are selected for requests</span>
|
||||
</div>
|
||||
<select
|
||||
class="select bg-space-800 border-space-border text-gray-200 focus:border-neon-cyan focus:ring-neon-cyan/20 w-64"
|
||||
:value="serverConfig.accountSelection?.strategy || 'hybrid'"
|
||||
@change="toggleStrategy($el.value)"
|
||||
aria-label="Account selection strategy">
|
||||
<option value="hybrid" x-text="$store.global.t('strategyHybridLabel')">Hybrid (Smart Distribution)</option>
|
||||
<option value="sticky" x-text="$store.global.t('strategyStickyLabel')">Sticky (Cache Optimized)</option>
|
||||
<option value="round-robin" x-text="$store.global.t('strategyRoundRobinLabel')">Round Robin (Load Balanced)</option>
|
||||
</select>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- ▼ Advanced Tuning (Fixed Logic) -->
|
||||
<div class="view-card !p-0 border-space-border/50">
|
||||
<div class="flex items-center justify-between p-4 cursor-pointer hover:bg-white/5 transition-colors"
|
||||
@@ -1095,6 +1124,134 @@
|
||||
x-text="$store.global.t('maxWaitDesc')">If all accounts are rate-limited longer than this, error immediately.</p>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- Error Handling Tuning -->
|
||||
<div class="space-y-4 pt-2 border-t border-space-border/10">
|
||||
<div class="flex items-center gap-2 mb-2">
|
||||
<span class="text-[10px] text-gray-500 font-bold uppercase tracking-widest"
|
||||
x-text="$store.global.t('errorHandlingTuning')">Error Handling Tuning</span>
|
||||
</div>
|
||||
|
||||
<div class="form-control">
|
||||
<label class="label pt-0">
|
||||
<span class="label-text text-gray-400 text-xs"
|
||||
x-text="$store.global.t('rateLimitDedupWindow')">Rate Limit Dedup Window</span>
|
||||
<span class="label-text-alt font-mono text-neon-cyan text-xs font-semibold"
|
||||
x-text="Math.round((serverConfig.rateLimitDedupWindowMs || 5000) / 1000) + 's'"></span>
|
||||
</label>
|
||||
<div class="flex gap-3 items-center">
|
||||
<input type="range" min="1000" max="30000" step="1000"
|
||||
class="custom-range custom-range-cyan flex-1"
|
||||
:value="serverConfig.rateLimitDedupWindowMs || 5000"
|
||||
:style="`background-size: ${((serverConfig.rateLimitDedupWindowMs || 5000) - 1000) / 290}% 100%`"
|
||||
@input="toggleRateLimitDedupWindowMs($event.target.value)"
|
||||
aria-label="Rate limit dedup window slider">
|
||||
<input type="number" min="1000" max="30000" step="1000"
|
||||
class="input input-xs input-bordered w-20 bg-space-800 border-space-border text-white font-mono text-center"
|
||||
:value="serverConfig.rateLimitDedupWindowMs || 5000"
|
||||
@change="toggleRateLimitDedupWindowMs($event.target.value)"
|
||||
aria-label="Rate limit dedup window value">
|
||||
</div>
|
||||
<p class="text-[9px] text-gray-600 mt-1 leading-tight"
|
||||
x-text="$store.global.t('rateLimitDedupWindowDesc')">Prevents concurrent retry storms.</p>
|
||||
</div>
|
||||
|
||||
<div class="form-control">
|
||||
<label class="label pt-0">
|
||||
<span class="label-text text-gray-400 text-xs"
|
||||
x-text="$store.global.t('maxConsecutiveFailures')">Max Consecutive Failures</span>
|
||||
<span class="label-text-alt font-mono text-neon-cyan text-xs font-semibold"
|
||||
x-text="serverConfig.maxConsecutiveFailures || 3"></span>
|
||||
</label>
|
||||
<div class="flex gap-3 items-center">
|
||||
<input type="range" min="1" max="10" step="1"
|
||||
class="custom-range custom-range-cyan flex-1"
|
||||
:value="serverConfig.maxConsecutiveFailures || 3"
|
||||
:style="`background-size: ${((serverConfig.maxConsecutiveFailures || 3) - 1) / 0.09}% 100%`"
|
||||
@input="toggleMaxConsecutiveFailures($event.target.value)"
|
||||
aria-label="Max consecutive failures slider">
|
||||
<input type="number" min="1" max="10" step="1"
|
||||
class="input input-xs input-bordered w-16 bg-space-800 border-space-border text-white font-mono text-center"
|
||||
:value="serverConfig.maxConsecutiveFailures || 3"
|
||||
@change="toggleMaxConsecutiveFailures($event.target.value)"
|
||||
aria-label="Max consecutive failures value">
|
||||
</div>
|
||||
<p class="text-[9px] text-gray-600 mt-1 leading-tight"
|
||||
x-text="$store.global.t('maxConsecutiveFailuresDesc')">Failures before extended cooldown.</p>
|
||||
</div>
|
||||
|
||||
<div class="form-control">
|
||||
<label class="label pt-0">
|
||||
<span class="label-text text-gray-400 text-xs"
|
||||
x-text="$store.global.t('extendedCooldown')">Extended Cooldown</span>
|
||||
<span class="label-text-alt font-mono text-neon-cyan text-xs font-semibold"
|
||||
x-text="((serverConfig.extendedCooldownMs || 60000) >= 60000 ? Math.round((serverConfig.extendedCooldownMs || 60000) / 60000) + 'm' : Math.round((serverConfig.extendedCooldownMs || 60000) / 1000) + 's')"></span>
|
||||
</label>
|
||||
<div class="flex gap-3 items-center">
|
||||
<input type="range" min="10000" max="300000" step="10000"
|
||||
class="custom-range custom-range-cyan flex-1"
|
||||
:value="serverConfig.extendedCooldownMs || 60000"
|
||||
:style="`background-size: ${((serverConfig.extendedCooldownMs || 60000) - 10000) / 2900}% 100%`"
|
||||
@input="toggleExtendedCooldownMs($event.target.value)"
|
||||
aria-label="Extended cooldown slider">
|
||||
<input type="number" min="10000" max="300000" step="10000"
|
||||
class="input input-xs input-bordered w-24 bg-space-800 border-space-border text-white font-mono text-center"
|
||||
:value="serverConfig.extendedCooldownMs || 60000"
|
||||
@change="toggleExtendedCooldownMs($event.target.value)"
|
||||
aria-label="Extended cooldown value">
|
||||
</div>
|
||||
<p class="text-[9px] text-gray-600 mt-1 leading-tight"
|
||||
x-text="$store.global.t('extendedCooldownDesc')">Applied after max consecutive failures.</p>
|
||||
</div>
|
||||
|
||||
<div class="form-control">
|
||||
<label class="label pt-0">
|
||||
<span class="label-text text-gray-400 text-xs"
|
||||
x-text="$store.global.t('capacityRetryDelay')">Capacity Retry Delay</span>
|
||||
<span class="label-text-alt font-mono text-neon-cyan text-xs font-semibold"
|
||||
x-text="Math.round((serverConfig.capacityRetryDelayMs || 2000) / 1000) + 's'"></span>
|
||||
</label>
|
||||
<div class="flex gap-3 items-center">
|
||||
<input type="range" min="500" max="10000" step="500"
|
||||
class="custom-range custom-range-cyan flex-1"
|
||||
:value="serverConfig.capacityRetryDelayMs || 2000"
|
||||
:style="`background-size: ${((serverConfig.capacityRetryDelayMs || 2000) - 500) / 95}% 100%`"
|
||||
@input="toggleCapacityRetryDelayMs($event.target.value)"
|
||||
aria-label="Capacity retry delay slider">
|
||||
<input type="number" min="500" max="10000" step="500"
|
||||
class="input input-xs input-bordered w-20 bg-space-800 border-space-border text-white font-mono text-center"
|
||||
:value="serverConfig.capacityRetryDelayMs || 2000"
|
||||
@change="toggleCapacityRetryDelayMs($event.target.value)"
|
||||
aria-label="Capacity retry delay value">
|
||||
</div>
|
||||
<p class="text-[9px] text-gray-600 mt-1 leading-tight"
|
||||
x-text="$store.global.t('capacityRetryDelayDesc')">Delay for capacity (not quota) issues.</p>
|
||||
</div>
|
||||
|
||||
<div class="form-control">
|
||||
<label class="label pt-0">
|
||||
<span class="label-text text-gray-400 text-xs"
|
||||
x-text="$store.global.t('maxCapacityRetries')">Max Capacity Retries</span>
|
||||
<span class="label-text-alt font-mono text-neon-cyan text-xs font-semibold"
|
||||
x-text="serverConfig.maxCapacityRetries || 3"></span>
|
||||
</label>
|
||||
<div class="flex gap-3 items-center">
|
||||
<input type="range" min="1" max="10" step="1"
|
||||
class="custom-range custom-range-cyan flex-1"
|
||||
:value="serverConfig.maxCapacityRetries || 3"
|
||||
:style="`background-size: ${((serverConfig.maxCapacityRetries || 3) - 1) / 0.09}% 100%`"
|
||||
@input="toggleMaxCapacityRetries($event.target.value)"
|
||||
aria-label="Max capacity retries slider">
|
||||
<input type="number" min="1" max="10" step="1"
|
||||
class="input input-xs input-bordered w-16 bg-space-800 border-space-border text-white font-mono text-center"
|
||||
:value="serverConfig.maxCapacityRetries || 3"
|
||||
@change="toggleMaxCapacityRetries($event.target.value)"
|
||||
aria-label="Max capacity retries value">
|
||||
</div>
|
||||
<p class="text-[9px] text-gray-600 mt-1 leading-tight"
|
||||
x-text="$store.global.t('maxCapacityRetriesDesc')">Retries before switching accounts.</p>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
/**
|
||||
|
||||
@@ -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 };
|
||||
}
|
||||
104
src/account-manager/strategies/base-strategy.js
Normal file
104
src/account-manager/strategies/base-strategy.js
Normal file
@@ -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;
|
||||
195
src/account-manager/strategies/hybrid-strategy.js
Normal file
195
src/account-manager/strategies/hybrid-strategy.js
Normal file
@@ -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;
|
||||
85
src/account-manager/strategies/index.js
Normal file
85
src/account-manager/strategies/index.js
Normal file
@@ -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';
|
||||
76
src/account-manager/strategies/round-robin-strategy.js
Normal file
76
src/account-manager/strategies/round-robin-strategy.js
Normal file
@@ -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;
|
||||
138
src/account-manager/strategies/sticky-strategy.js
Normal file
138
src/account-manager/strategies/sticky-strategy.js
Normal file
@@ -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;
|
||||
162
src/account-manager/strategies/trackers/health-tracker.js
Normal file
162
src/account-manager/strategies/trackers/health-tracker.js
Normal file
@@ -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;
|
||||
8
src/account-manager/strategies/trackers/index.js
Normal file
8
src/account-manager/strategies/trackers/index.js
Normal file
@@ -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';
|
||||
121
src/account-manager/strategies/trackers/token-bucket-tracker.js
Normal file
121
src/account-manager/strategies/trackers/token-bucket-tracker.js
Normal file
@@ -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;
|
||||
@@ -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')) {
|
||||
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...`);
|
||||
accountManager.pickNext(model);
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
if (isNetworkError(error)) {
|
||||
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;
|
||||
}
|
||||
|
||||
|
||||
@@ -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')) {
|
||||
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...`);
|
||||
accountManager.pickNext(model);
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
if (isNetworkError(error)) {
|
||||
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;
|
||||
}
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
};
|
||||
|
||||
30
src/index.js
30
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=<s> Set selection strategy (sticky/hybrid) ║\n';
|
||||
if (!isDebug) {
|
||||
controlSection += '║ --debug Enable debug logging ║\n';
|
||||
}
|
||||
@@ -53,18 +70,19 @@ 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';
|
||||
// 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(`
|
||||
╔══════════════════════════════════════════════════════════════╗
|
||||
|
||||
@@ -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}`);
|
||||
|
||||
@@ -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({
|
||||
|
||||
@@ -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' },
|
||||
|
||||
795
tests/test-strategies.cjs
Normal file
795
tests/test-strategies.cjs
Normal file
@@ -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);
|
||||
});
|
||||
Reference in New Issue
Block a user