tryaii-mcp-server
Version:
TryAII MCP Server - 15+ AI models with comparison, cost tracking, and collective intelligence
294 lines • 12.1 kB
JavaScript
import { logger } from '../utils/logger.js';
export class RateLimitService {
static instance;
usageStore = new Map();
blockStore = new Map(); // userId -> unblock timestamp
tokenBuckets = new Map();
// Rate limit tiers based on user type
rateLimits = {
free: {
windowMs: 60 * 1000, // 1 minute
maxRequests: 10,
maxCost: 0.10, // $0.10 per minute
blockDurationMs: 60 * 1000, // 1 minute block
algorithm: 'sliding'
},
pro: {
windowMs: 60 * 1000, // 1 minute
maxRequests: 100,
maxCost: 1.00, // $1.00 per minute
blockDurationMs: 30 * 1000, // 30 second block
algorithm: 'token-bucket'
},
enterprise: {
windowMs: 60 * 1000, // 1 minute
maxRequests: 1000,
maxCost: 10.00, // $10.00 per minute
blockDurationMs: 10 * 1000, // 10 second block
algorithm: 'token-bucket'
},
admin: {
windowMs: 60 * 1000,
maxRequests: 10000,
maxCost: 100.00,
blockDurationMs: 0, // No blocking for admins
algorithm: 'fixed'
}
};
constructor() { }
static getInstance() {
if (!RateLimitService.instance) {
RateLimitService.instance = new RateLimitService();
}
return RateLimitService.instance;
}
async checkRateLimit(userId, userTier = 'free', cost = 0) {
const now = Date.now();
const tier = this.rateLimits[userTier] || this.rateLimits.free;
if (!tier) {
throw new Error(`Rate limit configuration not found for tier: ${userTier}`);
}
// Check if user is currently blocked
const blockUntil = this.blockStore.get(userId);
if (blockUntil && now < blockUntil) {
logger.warn('Rate limit: User blocked', {
userId,
userTier,
retryAfter: Math.ceil((blockUntil - now) / 1000)
});
return {
allowed: false,
remaining: 0,
resetTime: blockUntil,
retryAfter: Math.ceil((blockUntil - now) / 1000),
reason: 'Rate limit exceeded - temporarily blocked'
};
}
// Apply rate limiting based on algorithm
let result;
switch (tier.algorithm) {
case 'sliding':
result = this.applySlidingWindowRateLimit(userId, tier, cost, now);
break;
case 'token-bucket':
result = this.applyTokenBucketRateLimit(userId, tier, cost, now);
break;
case 'fixed':
default:
result = this.applyFixedWindowRateLimit(userId, tier, cost, now);
break;
}
// If rate limit exceeded and blocking is enabled, add to block store
if (!result.allowed && tier.blockDurationMs > 0) {
const blockUntil = now + tier.blockDurationMs;
this.blockStore.set(userId, blockUntil);
result.retryAfter = Math.ceil(tier.blockDurationMs / 1000);
logger.warn('Rate limit exceeded - user blocked', {
userId,
userTier,
blockDurationMs: tier.blockDurationMs,
cost,
algorithm: tier.algorithm
});
}
// Log rate limit check
logger.debug('Rate limit check', {
userId,
userTier,
cost,
allowed: result.allowed,
remaining: result.remaining,
algorithm: tier.algorithm
});
return result;
}
applySlidingWindowRateLimit(userId, tier, cost, now) {
const windowStart = now - tier.windowMs;
const userUsage = this.usageStore.get(userId) || [];
// Remove old entries outside the sliding window
const validUsage = userUsage.filter(record => record.timestamp > windowStart);
// Calculate current usage
const totalRequests = validUsage.reduce((sum, record) => sum + record.requests, 0);
const totalCost = validUsage.reduce((sum, record) => sum + record.cost, 0);
// Check limits
const requestsAllowed = totalRequests + 1 <= tier.maxRequests;
const costAllowed = totalCost + cost <= tier.maxCost;
const allowed = requestsAllowed && costAllowed;
if (allowed) {
// Add new usage record
validUsage.push({ timestamp: now, requests: 1, cost });
this.usageStore.set(userId, validUsage);
}
return {
allowed,
remaining: Math.max(0, tier.maxRequests - totalRequests - (allowed ? 1 : 0)),
resetTime: now + tier.windowMs,
reason: !allowed ? (!requestsAllowed ? 'Request limit exceeded' : 'Cost limit exceeded') : undefined
};
}
applyTokenBucketRateLimit(userId, tier, cost, now) {
let bucket = this.tokenBuckets.get(userId);
if (!bucket) {
bucket = { tokens: tier.maxRequests, lastRefill: now };
this.tokenBuckets.set(userId, bucket);
}
// Refill tokens based on time elapsed
const timeSinceRefill = now - bucket.lastRefill;
const tokensToAdd = Math.floor(timeSinceRefill / tier.windowMs * tier.maxRequests);
if (tokensToAdd > 0) {
bucket.tokens = Math.min(tier.maxRequests, bucket.tokens + tokensToAdd);
bucket.lastRefill = now;
}
// Check if request can be processed
const tokensNeeded = Math.max(1, Math.ceil(cost * tier.maxRequests / tier.maxCost));
const allowed = bucket.tokens >= tokensNeeded;
if (allowed) {
bucket.tokens -= tokensNeeded;
}
// Calculate when bucket will have enough tokens again
const nextRefillTime = bucket.lastRefill + tier.windowMs;
return {
allowed,
remaining: bucket.tokens,
resetTime: nextRefillTime,
reason: !allowed ? 'Token bucket depleted' : undefined
};
}
applyFixedWindowRateLimit(userId, tier, cost, now) {
const windowStart = Math.floor(now / tier.windowMs) * tier.windowMs;
const userUsage = this.usageStore.get(userId) || [];
// Find current window usage
const currentWindow = userUsage.find(record => record.timestamp >= windowStart && record.timestamp < windowStart + tier.windowMs);
const currentRequests = currentWindow?.requests || 0;
const currentCost = currentWindow?.cost || 0;
// Check limits
const requestsAllowed = currentRequests + 1 <= tier.maxRequests;
const costAllowed = currentCost + cost <= tier.maxCost;
const allowed = requestsAllowed && costAllowed;
if (allowed) {
if (currentWindow) {
currentWindow.requests += 1;
currentWindow.cost += cost;
}
else {
userUsage.push({
timestamp: windowStart,
requests: 1,
cost
});
}
this.usageStore.set(userId, userUsage);
}
return {
allowed,
remaining: Math.max(0, tier.maxRequests - currentRequests - (allowed ? 1 : 0)),
resetTime: windowStart + tier.windowMs,
reason: !allowed ? (!requestsAllowed ? 'Request limit exceeded' : 'Cost limit exceeded') : undefined
};
}
// Get current usage statistics for a user
async getUserUsage(userId, userTier = 'free') {
const now = Date.now();
const tier = this.rateLimits[userTier] || this.rateLimits.free;
if (!tier) {
throw new Error(`Rate limit configuration not found for tier: ${userTier}`);
}
const blockUntil = this.blockStore.get(userId);
// Calculate current usage based on algorithm
let currentRequests = 0;
let currentCost = 0;
let resetTime = now + tier.windowMs;
if (tier.algorithm === 'sliding') {
const windowStart = now - tier.windowMs;
const userUsage = this.usageStore.get(userId) || [];
const validUsage = userUsage.filter(record => record.timestamp > windowStart);
currentRequests = validUsage.reduce((sum, record) => sum + record.requests, 0);
currentCost = validUsage.reduce((sum, record) => sum + record.cost, 0);
}
else if (tier.algorithm === 'token-bucket') {
const bucket = this.tokenBuckets.get(userId);
if (bucket) {
currentRequests = tier.maxRequests - bucket.tokens;
resetTime = bucket.lastRefill + tier.windowMs;
}
}
else { // fixed window
const windowStart = Math.floor(now / tier.windowMs) * tier.windowMs;
const userUsage = this.usageStore.get(userId) || [];
const currentWindow = userUsage.find(record => record.timestamp >= windowStart && record.timestamp < windowStart + tier.windowMs);
currentRequests = currentWindow?.requests || 0;
currentCost = currentWindow?.cost || 0;
resetTime = windowStart + tier.windowMs;
}
return {
currentPeriod: {
requests: currentRequests,
cost: currentCost
},
remaining: {
requests: Math.max(0, tier.maxRequests - currentRequests),
cost: Math.max(0, tier.maxCost - currentCost)
},
resetTime,
tier: userTier,
blocked: blockUntil ? now < blockUntil : false,
blockExpiresAt: blockUntil && now < blockUntil ? blockUntil : undefined
};
}
// Admin function to reset rate limits for a user
async resetUserLimits(userId) {
this.usageStore.delete(userId);
this.blockStore.delete(userId);
this.tokenBuckets.delete(userId);
logger.info('Rate limits reset for user', { userId });
}
// Admin function to get all users' rate limit status
async getAllUsersStatus() {
const status = new Map();
// Combine all user IDs from different stores
const userIds = new Set([
...this.usageStore.keys(),
...this.blockStore.keys(),
...this.tokenBuckets.keys()
]);
for (const userId of userIds) {
const usage = await this.getUserUsage(userId);
status.set(userId, usage);
}
return status;
}
// Cleanup old data periodically
async cleanup() {
const now = Date.now();
const maxAge = 24 * 60 * 60 * 1000; // 24 hours
// Clean up usage store
for (const [userId, records] of this.usageStore.entries()) {
const validRecords = records.filter(record => now - record.timestamp < maxAge);
if (validRecords.length === 0) {
this.usageStore.delete(userId);
}
else {
this.usageStore.set(userId, validRecords);
}
}
// Clean up block store
for (const [userId, blockUntil] of this.blockStore.entries()) {
if (now >= blockUntil) {
this.blockStore.delete(userId);
}
}
// Clean up old token buckets
for (const [userId, bucket] of this.tokenBuckets.entries()) {
if (now - bucket.lastRefill > maxAge) {
this.tokenBuckets.delete(userId);
}
}
logger.debug('Rate limit cleanup completed', {
usageStoreSize: this.usageStore.size,
blockStoreSize: this.blockStore.size,
tokenBucketsSize: this.tokenBuckets.size
});
}
}
export const rateLimitService = RateLimitService.getInstance();
//# sourceMappingURL=RateLimitService.js.map