/**
* Rate Limiting Middleware
* Implements token bucket algorithm for rate limiting MCP requests
*/
import { EventEmitter } from 'events';
/**
* Token Bucket Rate Limiter
* Allows burst traffic while maintaining average rate limits
*/
export class RateLimiter extends EventEmitter {
constructor(options = {}) {
super();
// Configuration
this.config = {
tokensPerInterval: options.tokensPerInterval || 10, // Requests allowed per interval
interval: options.interval || 60000, // Time window in ms (default: 1 minute)
maxBurst: options.maxBurst || 20, // Maximum burst capacity
enableLogging: options.enableLogging !== false,
blockDuration: options.blockDuration || 60000, // How long to block after limit exceeded
...options
};
// Token bucket state
this.tokens = this.config.maxBurst;
this.lastRefill = Date.now();
// Track blocked clients
this.blockedClients = new Map();
// Statistics
this.stats = {
totalRequests: 0,
allowedRequests: 0,
blockedRequests: 0,
throttledRequests: 0,
currentTokens: this.tokens
};
// Start token refill timer
this.startRefillTimer();
}
/**
* Check if a request should be allowed
* @param {string} clientId - Unique identifier for the client
* @param {number} cost - Token cost for this request (default: 1)
* @returns {Object} { allowed: boolean, reason?: string, retryAfter?: number }
*/
async checkLimit(clientId, cost = 1) {
this.stats.totalRequests++;
// Check if client is blocked
if (this.isBlocked(clientId)) {
this.stats.blockedRequests++;
const blockExpiry = this.blockedClients.get(clientId);
const retryAfter = Math.max(0, blockExpiry - Date.now());
if (this.config.enableLogging) {
console.warn(`Rate limit: Client ${clientId} is blocked for ${retryAfter}ms`);
}
return {
allowed: false,
reason: 'Rate limit exceeded - client blocked',
retryAfter
};
}
// Refill tokens based on elapsed time
this.refillTokens();
// Check if we have enough tokens
if (this.tokens >= cost) {
this.tokens -= cost;
this.stats.allowedRequests++;
this.stats.currentTokens = this.tokens;
this.emit('request:allowed', { clientId, cost, tokensRemaining: this.tokens });
return {
allowed: true,
tokensRemaining: this.tokens
};
} else {
this.stats.throttledRequests++;
// Calculate when tokens will be available
const tokensNeeded = cost - this.tokens;
const refillRate = this.config.tokensPerInterval / this.config.interval;
const retryAfter = Math.ceil(tokensNeeded / refillRate);
// Block client if they're repeatedly hitting limits
this.handleRateLimitExceeded(clientId);
if (this.config.enableLogging) {
console.warn(`Rate limit: Not enough tokens for ${clientId}. Need ${cost}, have ${this.tokens}`);
}
this.emit('request:throttled', { clientId, cost, tokensAvailable: this.tokens, retryAfter });
return {
allowed: false,
reason: 'Rate limit exceeded - insufficient tokens',
retryAfter,
tokensAvailable: this.tokens
};
}
}
/**
* Refill tokens based on elapsed time
*/
refillTokens() {
const now = Date.now();
const elapsed = now - this.lastRefill;
if (elapsed > 0) {
const refillRate = this.config.tokensPerInterval / this.config.interval;
const tokensToAdd = elapsed * refillRate;
this.tokens = Math.min(this.config.maxBurst, this.tokens + tokensToAdd);
this.lastRefill = now;
this.stats.currentTokens = this.tokens;
}
}
/**
* Start automatic token refill timer
*/
startRefillTimer() {
// Refill tokens periodically
this.refillInterval = setInterval(() => {
this.refillTokens();
this.cleanupBlockedClients();
}, Math.min(this.config.interval / 10, 1000)); // Refill every 100ms or 1s, whichever is smaller
}
/**
* Stop the refill timer
*/
stopRefillTimer() {
if (this.refillInterval) {
clearInterval(this.refillInterval);
this.refillInterval = null;
}
}
/**
* Check if a client is blocked
*/
isBlocked(clientId) {
const blockExpiry = this.blockedClients.get(clientId);
if (!blockExpiry) return false;
if (Date.now() >= blockExpiry) {
this.blockedClients.delete(clientId);
return false;
}
return true;
}
/**
* Handle rate limit exceeded - potentially block the client
*/
handleRateLimitExceeded(clientId) {
// Simple strategy: block after multiple violations
// You can implement more sophisticated strategies here
const violations = this.getViolations(clientId);
if (violations >= 3) {
const blockUntil = Date.now() + this.config.blockDuration;
this.blockedClients.set(clientId, blockUntil);
if (this.config.enableLogging) {
console.error(`Rate limit: Blocking client ${clientId} for ${this.config.blockDuration}ms due to repeated violations`);
}
this.emit('client:blocked', { clientId, blockUntil, violations });
}
}
/**
* Track violations per client (simple in-memory tracking)
*/
violations = new Map();
getViolations(clientId) {
const current = this.violations.get(clientId) || 0;
this.violations.set(clientId, current + 1);
// Reset violations after some time
setTimeout(() => {
this.violations.delete(clientId);
}, this.config.interval * 5);
return current + 1;
}
/**
* Clean up expired blocked clients
*/
cleanupBlockedClients() {
const now = Date.now();
for (const [clientId, expiry] of this.blockedClients.entries()) {
if (now >= expiry) {
this.blockedClients.delete(clientId);
this.emit('client:unblocked', { clientId });
}
}
}
/**
* Reset rate limiter state
*/
reset() {
this.tokens = this.config.maxBurst;
this.lastRefill = Date.now();
this.blockedClients.clear();
this.violations.clear();
this.stats = {
totalRequests: 0,
allowedRequests: 0,
blockedRequests: 0,
throttledRequests: 0,
currentTokens: this.tokens
};
}
/**
* Get current statistics
*/
getStats() {
return {
...this.stats,
blockedClients: this.blockedClients.size,
config: this.config
};
}
/**
* Destroy the rate limiter
*/
destroy() {
this.stopRefillTimer();
this.removeAllListeners();
this.blockedClients.clear();
this.violations.clear();
}
}
/**
* Create a rate limiter middleware for MCP server
*/
export function createRateLimitMiddleware(options = {}) {
const limiter = new RateLimiter(options);
return {
limiter,
/**
* Middleware function to check rate limits
*/
async checkRequest(clientId, requestType = 'default', cost = 1) {
// Different costs for different operations
const operationCosts = {
generate_image: 5, // Image generation is expensive
upscale_image: 3, // Upscaling is moderately expensive
remove_background: 3,
connect_comfyui: 1,
disconnect_comfyui: 1,
check_models: 1,
default: 1
};
const actualCost = operationCosts[requestType] || cost;
return limiter.checkLimit(clientId, actualCost);
},
/**
* Express-style middleware (if needed)
*/
middleware: (req, res, next) => {
const clientId = req.ip || req.connection.remoteAddress || 'unknown';
const requestType = req.body?.method || 'default';
limiter.checkLimit(clientId, 1).then(result => {
if (result.allowed) {
next();
} else {
res.status(429).json({
error: 'Rate limit exceeded',
reason: result.reason,
retryAfter: result.retryAfter
});
}
});
}
};
}
/**
* Sliding Window Rate Limiter (alternative implementation)
* More accurate but uses more memory
*/
export class SlidingWindowRateLimiter {
constructor(options = {}) {
this.windowSize = options.windowSize || 60000; // 1 minute window
this.maxRequests = options.maxRequests || 60; // Max requests per window
this.requests = new Map(); // clientId -> timestamp array
}
async checkLimit(clientId) {
const now = Date.now();
const windowStart = now - this.windowSize;
// Get or create request history for client
let clientRequests = this.requests.get(clientId) || [];
// Remove old requests outside the window
clientRequests = clientRequests.filter(timestamp => timestamp > windowStart);
// Check if limit exceeded
if (clientRequests.length >= this.maxRequests) {
const oldestRequest = clientRequests[0];
const retryAfter = (oldestRequest + this.windowSize) - now;
return {
allowed: false,
reason: 'Rate limit exceeded',
retryAfter: Math.max(0, retryAfter)
};
}
// Add current request
clientRequests.push(now);
this.requests.set(clientId, clientRequests);
return {
allowed: true,
remaining: this.maxRequests - clientRequests.length
};
}
cleanup() {
const now = Date.now();
const windowStart = now - this.windowSize;
for (const [clientId, requests] of this.requests.entries()) {
const validRequests = requests.filter(timestamp => timestamp > windowStart);
if (validRequests.length === 0) {
this.requests.delete(clientId);
} else {
this.requests.set(clientId, validRequests);
}
}
}
}
