From 2f5babba990717a03416a3cdffa5b5d440c159c8 Mon Sep 17 00:00:00 2001 From: jgor20 <102353650+jgor20@users.noreply.github.com> Date: Wed, 21 Jan 2026 11:15:38 +0000 Subject: [PATCH 1/2] feat(strategy): add quota-awareness to hybrid account selection MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The hybrid strategy now considers account quota levels when selecting accounts, preventing any single account from being drained to 0%. - Add QuotaTracker class to track per-account quota levels - Exclude accounts with critical quota (<5%) from selection - Add quota component to scoring formula (weight: 3) - Fall back to critical accounts when no alternatives exist - Add 18 new tests for quota-aware selection Scoring formula: Health×2 + Tokens×5 + Quota×3 + LRU×0.1 An attempt at resolving badrisnarayanan/antigravity-claude-proxy#171 --- .../strategies/hybrid-strategy.js | 56 +++- .../strategies/trackers/index.js | 1 + .../strategies/trackers/quota-tracker.js | 120 +++++++++ src/config.js | 6 + tests/test-strategies.cjs | 246 ++++++++++++++++++ 5 files changed, 422 insertions(+), 7 deletions(-) create mode 100644 src/account-manager/strategies/trackers/quota-tracker.js diff --git a/src/account-manager/strategies/hybrid-strategy.js b/src/account-manager/strategies/hybrid-strategy.js index 6814d7f..acda2a9 100644 --- a/src/account-manager/strategies/hybrid-strategy.js +++ b/src/account-manager/strategies/hybrid-strategy.js @@ -1,33 +1,36 @@ /** * Hybrid Strategy * - * Smart selection based on health score, token bucket, and LRU freshness. + * Smart selection based on health score, token bucket, quota, and LRU freshness. * Combines multiple signals for optimal account distribution. * * Scoring formula: - * score = (Health × 2) + ((Tokens / MaxTokens × 100) × 5) + (LRU × 0.1) + * score = (Health × 2) + ((Tokens / MaxTokens × 100) × 5) + (Quota × 3) + (LRU × 0.1) * * Filters accounts that are: * - Not rate-limited * - Not invalid or disabled * - Health score >= minUsable * - Has tokens available + * - Quota not critically low (< 5%) */ import { BaseStrategy } from './base-strategy.js'; -import { HealthTracker, TokenBucketTracker } from './trackers/index.js'; +import { HealthTracker, TokenBucketTracker, QuotaTracker } from './trackers/index.js'; import { logger } from '../../utils/logger.js'; // Default weights for scoring const DEFAULT_WEIGHTS = { health: 2, tokens: 5, + quota: 3, lru: 0.1 }; export class HybridStrategy extends BaseStrategy { #healthTracker; #tokenBucketTracker; + #quotaTracker; #weights; /** @@ -35,12 +38,14 @@ export class HybridStrategy extends BaseStrategy { * @param {Object} config - Strategy configuration * @param {Object} [config.healthScore] - Health tracker configuration * @param {Object} [config.tokenBucket] - Token bucket configuration + * @param {Object} [config.quota] - Quota tracker configuration * @param {Object} [config.weights] - Scoring weights */ constructor(config = {}) { super(config); this.#healthTracker = new HealthTracker(config.healthScore || {}); this.#tokenBucketTracker = new TokenBucketTracker(config.tokenBucket || {}); + this.#quotaTracker = new QuotaTracker(config.quota || {}); this.#weights = { ...DEFAULT_WEIGHTS, ...config.weights }; } @@ -71,7 +76,7 @@ export class HybridStrategy extends BaseStrategy { const scored = candidates.map(({ account, index }) => ({ account, index, - score: this.#calculateScore(account) + score: this.#calculateScore(account, modelId) })); scored.sort((a, b) => b.score - a.score); @@ -126,7 +131,7 @@ export class HybridStrategy extends BaseStrategy { * @private */ #getCandidates(accounts, modelId) { - return accounts + const candidates = accounts .map((account, index) => ({ account, index })) .filter(({ account }) => { // Basic usability check @@ -144,15 +149,40 @@ export class HybridStrategy extends BaseStrategy { return false; } + // Quota availability check (exclude critically low quota) + if (this.#quotaTracker.isQuotaCritical(account, modelId)) { + logger.debug(`[HybridStrategy] Excluding ${account.email}: quota critically low for ${modelId}`); + return false; + } + return true; }); + + // If no candidates after quota filter, fall back to all usable accounts + // (better to use critical quota than fail entirely) + if (candidates.length === 0) { + const fallback = accounts + .map((account, index) => ({ account, index })) + .filter(({ account }) => { + if (!this.isAccountUsable(account, modelId)) return false; + if (!this.#healthTracker.isUsable(account.email)) return false; + if (!this.#tokenBucketTracker.hasTokens(account.email)) return false; + return true; + }); + if (fallback.length > 0) { + logger.warn('[HybridStrategy] All accounts have critical quota, using fallback'); + return fallback; + } + } + + return candidates; } /** * Calculate the combined score for an account * @private */ - #calculateScore(account) { + #calculateScore(account, modelId) { const email = account.email; // Health component (0-100 scaled by weight) @@ -165,6 +195,10 @@ export class HybridStrategy extends BaseStrategy { const tokenRatio = tokens / maxTokens; const tokenComponent = (tokenRatio * 100) * this.#weights.tokens; + // Quota component (0-100 scaled by weight) + const quotaScore = this.#quotaTracker.getScore(account, modelId); + const quotaComponent = quotaScore * this.#weights.quota; + // LRU component (older = higher score) // Use time since last use, capped at 1 hour for scoring const lastUsed = account.lastUsed || 0; @@ -172,7 +206,7 @@ export class HybridStrategy extends BaseStrategy { const lruMinutes = timeSinceLastUse / 60000; const lruComponent = lruMinutes * this.#weights.lru; - return healthComponent + tokenComponent + lruComponent; + return healthComponent + tokenComponent + quotaComponent + lruComponent; } /** @@ -190,6 +224,14 @@ export class HybridStrategy extends BaseStrategy { getTokenBucketTracker() { return this.#tokenBucketTracker; } + + /** + * Get the quota tracker (for testing/debugging) + * @returns {QuotaTracker} The quota tracker instance + */ + getQuotaTracker() { + return this.#quotaTracker; + } } export default HybridStrategy; diff --git a/src/account-manager/strategies/trackers/index.js b/src/account-manager/strategies/trackers/index.js index 974f527..f2f9839 100644 --- a/src/account-manager/strategies/trackers/index.js +++ b/src/account-manager/strategies/trackers/index.js @@ -6,3 +6,4 @@ export { HealthTracker } from './health-tracker.js'; export { TokenBucketTracker } from './token-bucket-tracker.js'; +export { QuotaTracker } from './quota-tracker.js'; diff --git a/src/account-manager/strategies/trackers/quota-tracker.js b/src/account-manager/strategies/trackers/quota-tracker.js new file mode 100644 index 0000000..f3e014b --- /dev/null +++ b/src/account-manager/strategies/trackers/quota-tracker.js @@ -0,0 +1,120 @@ +/** + * Quota Tracker + * + * Tracks per-account quota levels to prioritize accounts with available quota. + * Uses quota data from account.quota.models[modelId].remainingFraction. + * Accounts below critical threshold are excluded from selection. + */ + +// Default configuration +const DEFAULT_CONFIG = { + lowThreshold: 0.10, // 10% - reduce score + criticalThreshold: 0.05, // 5% - exclude from candidates + staleMs: 300000, // 5 min - max age of quota data to trust + unknownScore: 50 // Score for accounts with unknown quota +}; + +export class QuotaTracker { + #config; + + /** + * Create a new QuotaTracker + * @param {Object} config - Quota tracker configuration + */ + constructor(config = {}) { + this.#config = { ...DEFAULT_CONFIG, ...config }; + } + + /** + * Get the quota fraction for an account and model + * @param {Object} account - Account object + * @param {string} modelId - Model ID to check + * @returns {number|null} Remaining fraction (0-1) or null if unknown + */ + getQuotaFraction(account, modelId) { + if (!account?.quota?.models?.[modelId]) return null; + const fraction = account.quota.models[modelId].remainingFraction; + return typeof fraction === 'number' ? fraction : null; + } + + /** + * Check if quota data is fresh enough to be trusted + * @param {Object} account - Account object + * @returns {boolean} True if quota data is fresh + */ + isQuotaFresh(account) { + if (!account?.quota?.lastChecked) return false; + return (Date.now() - account.quota.lastChecked) < this.#config.staleMs; + } + + /** + * Check if an account has critically low quota for a model + * @param {Object} account - Account object + * @param {string} modelId - Model ID to check + * @returns {boolean} True if quota is at or below critical threshold + */ + isQuotaCritical(account, modelId) { + const fraction = this.getQuotaFraction(account, modelId); + // Unknown quota = not critical (assume OK) + if (fraction === null) return false; + // Only apply critical check if data is fresh + if (!this.isQuotaFresh(account)) return false; + return fraction <= this.#config.criticalThreshold; + } + + /** + * Check if an account has low (but not critical) quota for a model + * @param {Object} account - Account object + * @param {string} modelId - Model ID to check + * @returns {boolean} True if quota is below low threshold but above critical + */ + isQuotaLow(account, modelId) { + const fraction = this.getQuotaFraction(account, modelId); + if (fraction === null) return false; + return fraction <= this.#config.lowThreshold && fraction > this.#config.criticalThreshold; + } + + /** + * Get a score (0-100) for an account based on quota + * Higher score = more quota available + * @param {Object} account - Account object + * @param {string} modelId - Model ID to check + * @returns {number} Score from 0-100 + */ + getScore(account, modelId) { + const fraction = this.getQuotaFraction(account, modelId); + + // Unknown quota = middle score + if (fraction === null) { + return this.#config.unknownScore; + } + + // Convert fraction (0-1) to score (0-100) + let score = fraction * 100; + + // Apply small penalty for stale data (reduce confidence) + if (!this.isQuotaFresh(account)) { + score *= 0.9; // 10% penalty for stale data + } + + return score; + } + + /** + * Get the critical threshold + * @returns {number} Critical threshold (0-1) + */ + getCriticalThreshold() { + return this.#config.criticalThreshold; + } + + /** + * Get the low threshold + * @returns {number} Low threshold (0-1) + */ + getLowThreshold() { + return this.#config.lowThreshold; + } +} + +export default QuotaTracker; diff --git a/src/config.js b/src/config.js index 953bf48..397dc2e 100644 --- a/src/config.js +++ b/src/config.js @@ -34,6 +34,12 @@ const DEFAULT_CONFIG = { maxTokens: 50, // Maximum token capacity tokensPerMinute: 6, // Regeneration rate initialTokens: 50 // Starting tokens + }, + quota: { + lowThreshold: 0.10, // 10% - reduce score + criticalThreshold: 0.05, // 5% - exclude from candidates + staleMs: 300000, // 5 min - max age of quota data to trust + weight: 3 // Scoring weight (same scale as health/tokens) } } }; diff --git a/tests/test-strategies.cjs b/tests/test-strategies.cjs index 0be8162..9cb7107 100644 --- a/tests/test-strategies.cjs +++ b/tests/test-strategies.cjs @@ -19,6 +19,7 @@ async function runTests() { // 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 { QuotaTracker } = await import('../src/account-manager/strategies/trackers/quota-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'); @@ -277,6 +278,149 @@ async function runTests() { assertEqual(tracker.getTokens('test@example.com'), 50, 'Reset should restore initial'); }); + // ========================================================================== + // QUOTA TRACKER TESTS + // ========================================================================== + console.log('\n─── QuotaTracker Tests ───'); + + test('QuotaTracker: getQuotaFraction returns null for missing data', () => { + const tracker = new QuotaTracker(); + const account = { email: 'test@example.com' }; + assertNull(tracker.getQuotaFraction(account, 'model'), 'Missing quota should return null'); + }); + + test('QuotaTracker: getQuotaFraction returns correct value', () => { + const tracker = new QuotaTracker(); + const account = { + email: 'test@example.com', + quota: { + models: { 'model': { remainingFraction: 0.75 } }, + lastChecked: Date.now() + } + }; + assertEqual(tracker.getQuotaFraction(account, 'model'), 0.75); + }); + + test('QuotaTracker: isQuotaFresh returns false when no lastChecked', () => { + const tracker = new QuotaTracker(); + const account = { email: 'test@example.com' }; + assertFalse(tracker.isQuotaFresh(account), 'Missing lastChecked should not be fresh'); + }); + + test('QuotaTracker: isQuotaFresh returns true for recent data', () => { + const tracker = new QuotaTracker({ staleMs: 300000 }); // 5 min + const account = { + email: 'test@example.com', + quota: { lastChecked: Date.now() - 60000 } // 1 min ago + }; + assertTrue(tracker.isQuotaFresh(account), 'Recent data should be fresh'); + }); + + test('QuotaTracker: isQuotaFresh returns false for stale data', () => { + const tracker = new QuotaTracker({ staleMs: 300000 }); // 5 min + const account = { + email: 'test@example.com', + quota: { lastChecked: Date.now() - 600000 } // 10 min ago + }; + assertFalse(tracker.isQuotaFresh(account), 'Old data should be stale'); + }); + + test('QuotaTracker: isQuotaCritical returns false for unknown quota', () => { + const tracker = new QuotaTracker({ criticalThreshold: 0.05 }); + const account = { email: 'test@example.com' }; + assertFalse(tracker.isQuotaCritical(account, 'model'), 'Unknown quota should not be critical'); + }); + + test('QuotaTracker: isQuotaCritical returns true when quota <= threshold', () => { + const tracker = new QuotaTracker({ criticalThreshold: 0.05 }); + const account = { + email: 'test@example.com', + quota: { + models: { 'model': { remainingFraction: 0.04 } }, + lastChecked: Date.now() + } + }; + assertTrue(tracker.isQuotaCritical(account, 'model'), 'Low quota should be critical'); + }); + + test('QuotaTracker: isQuotaCritical returns false when quota > threshold', () => { + const tracker = new QuotaTracker({ criticalThreshold: 0.05 }); + const account = { + email: 'test@example.com', + quota: { + models: { 'model': { remainingFraction: 0.10 } }, + lastChecked: Date.now() + } + }; + assertFalse(tracker.isQuotaCritical(account, 'model'), 'Higher quota should not be critical'); + }); + + test('QuotaTracker: isQuotaCritical returns false for stale data', () => { + const tracker = new QuotaTracker({ criticalThreshold: 0.05, staleMs: 300000 }); + const account = { + email: 'test@example.com', + quota: { + models: { 'model': { remainingFraction: 0.01 } }, + lastChecked: Date.now() - 600000 // 10 min ago (stale) + } + }; + assertFalse(tracker.isQuotaCritical(account, 'model'), 'Stale critical data should be ignored'); + }); + + test('QuotaTracker: isQuotaLow returns true for low but not critical quota', () => { + const tracker = new QuotaTracker({ lowThreshold: 0.10, criticalThreshold: 0.05 }); + const account = { + email: 'test@example.com', + quota: { + models: { 'model': { remainingFraction: 0.08 } }, + lastChecked: Date.now() + } + }; + assertTrue(tracker.isQuotaLow(account, 'model'), 'Quota at 8% should be low'); + }); + + test('QuotaTracker: isQuotaLow returns false for critical quota', () => { + const tracker = new QuotaTracker({ lowThreshold: 0.10, criticalThreshold: 0.05 }); + const account = { + email: 'test@example.com', + quota: { + models: { 'model': { remainingFraction: 0.03 } }, + lastChecked: Date.now() + } + }; + assertFalse(tracker.isQuotaLow(account, 'model'), 'Critical quota should not be just low'); + }); + + test('QuotaTracker: getScore returns unknownScore for missing quota', () => { + const tracker = new QuotaTracker({ unknownScore: 50 }); + const account = { email: 'test@example.com' }; + assertEqual(tracker.getScore(account, 'model'), 50, 'Unknown quota should return default score'); + }); + + test('QuotaTracker: getScore returns 0-100 based on fraction', () => { + const tracker = new QuotaTracker(); + const account = { + email: 'test@example.com', + quota: { + models: { 'model': { remainingFraction: 0.75 } }, + lastChecked: Date.now() + } + }; + assertEqual(tracker.getScore(account, 'model'), 75, 'Score should be fraction * 100'); + }); + + test('QuotaTracker: getScore applies penalty for stale data', () => { + const tracker = new QuotaTracker({ staleMs: 300000 }); + const account = { + email: 'test@example.com', + quota: { + models: { 'model': { remainingFraction: 1.0 } }, + lastChecked: Date.now() - 600000 // 10 min ago + } + }; + assertEqual(tracker.getScore(account, 'model'), 90, 'Stale data should have 10% penalty'); + }); + // ========================================================================== // BASE STRATEGY TESTS // ========================================================================== @@ -643,6 +787,108 @@ async function runTests() { assertEqual(result.account.email, 'old-account@example.com', 'Older account should win with LRU weight'); }); + test('HybridStrategy: filters out accounts with critical quota', () => { + const strategy = new HybridStrategy({ + healthScore: { initial: 70 }, + tokenBucket: { initialTokens: 50, maxTokens: 50 }, + quota: { criticalThreshold: 0.05, staleMs: 300000 } + }); + + const accounts = [ + { + email: 'critical@example.com', + enabled: true, + lastUsed: Date.now() - 3600000, // Older (would normally win LRU) + quota: { + models: { 'model': { remainingFraction: 0.02 } }, + lastChecked: Date.now() + } + }, + { + email: 'healthy@example.com', + enabled: true, + lastUsed: Date.now() + } + ]; + + const result = strategy.selectAccount(accounts, 'model'); + assertEqual(result.account.email, 'healthy@example.com', 'Critical quota account should be excluded'); + }); + + test('HybridStrategy: prefers higher quota accounts', () => { + const strategy = new HybridStrategy({ + healthScore: { initial: 70 }, + tokenBucket: { initialTokens: 50, maxTokens: 50 }, + quota: { weight: 3 }, + weights: { health: 2, tokens: 5, quota: 3, lru: 0.1 } + }); + + // Create accounts with same lastUsed (equal LRU) + const now = Date.now(); + const accounts = [ + { + email: 'low-quota@example.com', + enabled: true, + lastUsed: now, + quota: { + models: { 'model': { remainingFraction: 0.20 } }, + lastChecked: now + } + }, + { + email: 'high-quota@example.com', + enabled: true, + lastUsed: now, + quota: { + models: { 'model': { remainingFraction: 0.80 } }, + lastChecked: now + } + } + ]; + + const result = strategy.selectAccount(accounts, 'model'); + assertEqual(result.account.email, 'high-quota@example.com', 'Higher quota account should be preferred'); + }); + + test('HybridStrategy: falls back when all accounts have critical quota', () => { + const strategy = new HybridStrategy({ + healthScore: { initial: 70 }, + tokenBucket: { initialTokens: 50, maxTokens: 50 }, + quota: { criticalThreshold: 0.05, staleMs: 300000 } + }); + + const accounts = [ + { + email: 'critical1@example.com', + enabled: true, + lastUsed: Date.now() - 60000, + quota: { + models: { 'model': { remainingFraction: 0.02 } }, + lastChecked: Date.now() + } + }, + { + email: 'critical2@example.com', + enabled: true, + lastUsed: Date.now(), + quota: { + models: { 'model': { remainingFraction: 0.01 } }, + lastChecked: Date.now() + } + } + ]; + + // Should fall back and select an account even though all are critical + const result = strategy.selectAccount(accounts, 'model'); + assertTrue(result.account !== null, 'Should fall back to critical quota accounts when no alternatives'); + }); + + test('HybridStrategy: getQuotaTracker returns tracker', () => { + const strategy = new HybridStrategy(); + const tracker = strategy.getQuotaTracker(); + assertTrue(tracker instanceof QuotaTracker, 'Should return QuotaTracker instance'); + }); + // ========================================================================== // STRATEGY FACTORY TESTS // ========================================================================== From 7ed9305b5bdf69924243faea150ee5c2944bffef Mon Sep 17 00:00:00 2001 From: jgor20 <102353650+jgor20@users.noreply.github.com> Date: Wed, 21 Jan 2026 20:52:14 +0000 Subject: [PATCH 2/2] fix: remove unused weight from quota config The quota scoring weight is managed in HybridStrategy's DEFAULT_WEIGHTS, not in the config.quota block. Removed to avoid confusion. --- src/config.js | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/config.js b/src/config.js index 397dc2e..c41df9c 100644 --- a/src/config.js +++ b/src/config.js @@ -38,8 +38,7 @@ const DEFAULT_CONFIG = { quota: { lowThreshold: 0.10, // 10% - reduce score criticalThreshold: 0.05, // 5% - exclude from candidates - staleMs: 300000, // 5 min - max age of quota data to trust - weight: 3 // Scoring weight (same scale as health/tokens) + staleMs: 300000 // 5 min - max age of quota data to trust } } };