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:
@@ -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;
|
||||
|
||||
@@ -6,3 +6,4 @@
|
||||
|
||||
export { HealthTracker } from './health-tracker.js';
|
||||
export { TokenBucketTracker } from './token-bucket-tracker.js';
|
||||
export { QuotaTracker } from './quota-tracker.js';
|
||||
|
||||
120
src/account-manager/strategies/trackers/quota-tracker.js
Normal file
120
src/account-manager/strategies/trackers/quota-tracker.js
Normal 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;
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
Reference in New Issue
Block a user