Files
openmaic/OpenMAIC/lib/qa/rate-limit.ts
2026-08-16 14:58:47 +08:00

101 lines
3.4 KiB
TypeScript

// In-memory sliding-window rate limiter for the learner Q&A and TTS endpoints.
//
// The L2 hot path stays LLM-free; `qa` and `tts` are the two controlled LLM /
// TTS endpoints and must be throttled (design doc §4.2). Per-process state is
// fine for these (a multi-instance deployment should move this to a shared
// store), and windows are reaped lazily so the table cannot grow unbounded.
export interface RateLimitDecision {
allowed: boolean;
/** Requests remaining in the window (0 when blocked). */
remaining: number;
/** Milliseconds until the oldest request leaves the window. */
retryAfterMs: number;
}
export interface SlidingWindowLimiterOptions {
/** Window length, e.g. 60_000 (1 minute). */
windowMs: number;
/** Max requests per key per window. */
max: number;
/** Hard cap for process-local buckets, including one shared overflow bucket. */
maxKeys?: number;
}
interface Bucket {
/** Timestamps of requests inside the current window (ms). */
hits: number[];
}
export function createSlidingWindowLimiter(options: SlidingWindowLimiterOptions) {
const { windowMs, max, maxKeys = 10_000 } = options;
const buckets = new Map<string, Bucket>();
const overflowKey = '__openmaic_rate_limit_overflow__';
let lastReapAt = 0;
const reap = (now: number) => {
const cutoff = now - windowMs;
for (const [key, bucket] of buckets) {
bucket.hits = bucket.hits.filter((t) => t > cutoff);
if (bucket.hits.length === 0) buckets.delete(key);
}
lastReapAt = now;
};
const boundedKey = (key: string, now: number): string => {
if (buckets.has(key)) return key;
if (now - lastReapAt >= windowMs) reap(now);
// Reserve one slot for a shared overflow bucket. Unknown clients cannot
// grow the Map without bound or force a full-table reap on every request.
if (buckets.size < Math.max(0, maxKeys - 1)) return key;
return overflowKey;
};
return {
check(key: string, now: number = Date.now()): RateLimitDecision {
const resolvedKey = boundedKey(key, now);
let bucket = buckets.get(resolvedKey);
if (!bucket) {
bucket = { hits: [] };
buckets.set(resolvedKey, bucket);
}
const cutoff = now - windowMs;
bucket.hits = bucket.hits.filter((t) => t > cutoff);
if (bucket.hits.length >= max) {
const oldest = Math.min(...bucket.hits);
return {
allowed: false,
remaining: 0,
retryAfterMs: Math.max(0, windowMs - (now - oldest)),
};
}
bucket.hits.push(now);
return { allowed: true, remaining: max - bucket.hits.length, retryAfterMs: 0 };
},
/** Drop all state (tests). */
clear() {
buckets.clear();
lastReapAt = 0;
},
/** Process-local cardinality, exposed for deterministic capacity tests. */
size() {
return buckets.size;
},
};
}
export function shouldTrustLearnerRateLimitProxyHeaders(): boolean {
return process.env.LEARNER_RATE_LIMIT_TRUST_PROXY_HEADERS === 'true';
}
/** Client IP for rate limiting. Proxy headers are opt-in and must be overwritten upstream. */
export function clientIp(request: { headers: Headers; ip?: string }): string {
if (shouldTrustLearnerRateLimitProxyHeaders()) {
const forwarded = request.headers.get('x-forwarded-for');
if (forwarded) return forwarded.split(',')[0].trim();
const realIp = request.headers.get('x-real-ip');
if (realIp) return realIp;
}
return request.ip || 'unknown';
}