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

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
This commit is contained in:
jgor20
2026-01-21 11:15:38 +00:00
parent e51e3ff56a
commit 2f5babba99
5 changed files with 422 additions and 7 deletions

View File

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

View File

@@ -6,3 +6,4 @@
export { HealthTracker } from './health-tracker.js';
export { TokenBucketTracker } from './token-bucket-tracker.js';
export { QuotaTracker } from './quota-tracker.js';

View File

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

View File

@@ -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)
}
}
};