74 lines
2.8 KiB
JavaScript
74 lines
2.8 KiB
JavaScript
const SECOND = 1_000;
|
|
const FIVE_HOURS = 5 * 60 * 60 * 1_000;
|
|
|
|
export const RATE_LIMITS = {
|
|
requestsPerSecond: 2,
|
|
requestsPerFiveHours: 100,
|
|
spendPerFiveHours: 10,
|
|
};
|
|
|
|
function prune(timestamps, now, window) {
|
|
return timestamps.filter((timestamp) => timestamp > now - window);
|
|
}
|
|
|
|
export function clientIp(request) {
|
|
const railwayIp = request.get('x-real-ip');
|
|
return railwayIp?.split(',')[0].trim() || request.ip || request.socket.remoteAddress || 'unknown';
|
|
}
|
|
|
|
export function createRateLimiter({ now = () => Date.now() } = {}) {
|
|
const clients = new Map();
|
|
|
|
function stateFor(ip, timestamp) {
|
|
let state = clients.get(ip);
|
|
if (!state) {
|
|
state = { requests: [], spend: [] };
|
|
clients.set(ip, state);
|
|
}
|
|
state.requests = prune(state.requests, timestamp, FIVE_HOURS);
|
|
state.spend = state.spend.filter(({ timestamp: spentAt }) => spentAt > timestamp - FIVE_HOURS);
|
|
return state;
|
|
}
|
|
|
|
return {
|
|
check(ip) {
|
|
const timestamp = now();
|
|
const state = stateFor(ip, timestamp);
|
|
const recentRequests = state.requests.filter((requestAt) => requestAt > timestamp - SECOND);
|
|
const spend = state.spend.reduce((total, entry) => total + entry.amount, 0);
|
|
if (recentRequests.length >= RATE_LIMITS.requestsPerSecond) {
|
|
return { allowed: false, retryAfter: Math.max(1, Math.ceil((recentRequests[0] + SECOND - timestamp) / 1_000)), reason: 'rate' };
|
|
}
|
|
if (state.requests.length >= RATE_LIMITS.requestsPerFiveHours) {
|
|
return { allowed: false, retryAfter: Math.max(1, Math.ceil((state.requests[0] + FIVE_HOURS - timestamp) / 1_000)), reason: 'rate' };
|
|
}
|
|
if (spend >= RATE_LIMITS.spendPerFiveHours) {
|
|
const firstSpend = state.spend[0];
|
|
return { allowed: false, retryAfter: Math.max(1, Math.ceil((firstSpend.timestamp + FIVE_HOURS - timestamp) / 1_000)), reason: 'spend' };
|
|
}
|
|
state.requests.push(timestamp);
|
|
return { allowed: true };
|
|
},
|
|
|
|
recordCost(ip, amount) {
|
|
if (!Number.isFinite(amount) || amount < 0) return;
|
|
const timestamp = now();
|
|
const state = stateFor(ip, timestamp);
|
|
state.spend.push({ timestamp, amount });
|
|
},
|
|
};
|
|
}
|
|
|
|
export function numericCost(payload) {
|
|
const cost = payload?.cost ?? payload?.usage?.cost;
|
|
return typeof cost === 'number' && Number.isFinite(cost) ? cost : null;
|
|
}
|
|
|
|
export function rateLimitError(response, anthropic = false, result) {
|
|
response.set('retry-after', String(result.retryAfter));
|
|
if (anthropic) {
|
|
return response.status(429).json({ type: 'error', error: { type: 'rate_limit_error', message: `Rate limit exceeded (${result.reason === 'spend' ? 'spend' : 'request'} limit)` } });
|
|
}
|
|
return response.status(429).json({ error: { message: `Rate limit exceeded (${result.reason === 'spend' ? 'spend' : 'request'} limit)`, type: 'rate_limit_error' } });
|
|
}
|