Merge pull request #173 from jgor20/fix/quota-aware-selection

feat(strategy): add quota-awareness to hybrid account selection

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;

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';

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;

src/config.js

@@ -34,6 +34,11 @@ 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
         }
     }
 };

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
     // ==========================================================================