项目文件夹

文件
2026-07-13 13:39:12 +08:00

946 行
34 KiB
TypeScript

此文件含有模棱两可的 Unicode 字符
此文件含有可能会与其他字符混淆的 Unicode 字符。 如果您是想特意这样的,可以安全地忽略该警告。 使用 Escape 按钮显示他们。
/**
* Rate Limit Manager — Adaptive rate limiting using Bottleneck
*
* Creates per-provider+connection limiters that auto-learn rate limits
* from API response headers (x-ratelimit-*, retry-after, anthropic-ratelimit-*).
*
* Default: ENABLED for API key providers (safety net), DISABLED for OAuth.
* Can be toggled per provider connection via dashboard.
*/
import Bottleneck from "bottleneck";
import { parseRetryAfterFromBody } from "./accountFallback.ts";
import { getProviderCategory } from "../config/providerRegistry.ts";
import { getCodexRateLimitKey } from "../executors/codex.ts";
import { awaitProviderDefaultSlot } from "./providerDefaultRateLimit.ts";
import {
DEFAULT_RESILIENCE_SETTINGS,
resolveResilienceSettings,
type RequestQueueSettings,
} from "../../src/lib/resilience/settings";
import {
STANDARD_HEADERS,
ANTHROPIC_HEADERS,
parseResetTime,
toPlainHeaders,
} from "./rateLimitManager/headers";
interface LearnedLimitEntry {
provider: string;
connectionId: string;
lastUpdated: number;
limit?: number;
remaining?: number;
minTime?: number;
}
interface LimiterUpdateSettings {
maxConcurrent?: number | null;
minTime: number;
reservoir?: number | null;
reservoirRefreshAmount?: number | null;
reservoirRefreshInterval?: number | null;
}
type JsonRecord = Record<string, unknown>;
function toRecord(value: unknown): JsonRecord {
return value && typeof value === "object" && !Array.isArray(value) ? (value as JsonRecord) : {};
}
function toNumber(value: unknown, fallback = 0): number {
const parsed =
typeof value === "number"
? value
: typeof value === "string" && value.trim().length > 0
? Number(value)
: Number.NaN;
return Number.isFinite(parsed) ? parsed : fallback;
}
function isNodeTestRunnerChild(): boolean {
return typeof process.env.NODE_TEST_CONTEXT === "string";
}
function logRateLimit(...args: unknown[]): void {
if (!isNodeTestRunnerChild()) console.log(...args);
}
function warnRateLimit(...args: unknown[]): void {
if (!isNodeTestRunnerChild()) console.warn(...args);
}
function errorRateLimit(...args: unknown[]): void {
if (!isNodeTestRunnerChild()) console.error(...args);
}
// Store limiters keyed by "provider:connectionId" (and optionally ":model")
const limiters = new Map<string, Bottleneck>();
// Store connections that have rate limit protection enabled
const enabledConnections = new Set<string>();
// Store per-connection rate limit overrides (RPM, TPM, TPD, minTime, maxConcurrent)
// Populated from provider_connections.rateLimitOverrides on startup and refresh.
const connectionRateLimitOverrides = new Map<string, Record<string, number>>();
// Store learned limits for persistence (debounced)
const learnedLimits: Record<string, LearnedLimitEntry> = {};
const MAX_LEARNED_LIMITS = 200;
const INACTIVE_LIMITER_MS = 10 * 60 * 1000;
const limiterLastUsed = new Map<string, number>();
let persistTimer: ReturnType<typeof setTimeout> | null = null;
const pendingAsyncOperations = new Set<Promise<unknown>>();
const PERSIST_DEBOUNCE_MS = 60_000; // Debounce persistence to every 60s max
// Track initialization
let initialized = false;
let currentRequestQueueSettings: RequestQueueSettings = DEFAULT_RESILIENCE_SETTINGS.requestQueue;
// Watchdog: detect Bottleneck limiters that are wedged (queue has work, but no
// jobs are dispatched). When the reservoir/refresh state desyncs from reality,
// this catches it and force-resets so traffic isn't stuck forever.
const lastDispatchAt = new Map<string, number>();
let watchdogInterval: ReturnType<typeof setInterval> | null = null;
const WATCHDOG_INTERVAL_MS = 30_000;
// Threshold has to exceed any *legitimate* gap between dispatches:
// - default reservoirRefreshInterval is 60s
// - adaptive minTime can climb to ~60s for 1-RPM providers (see updateFromHeaders)
// 120s gives a 2× margin against both, while still catching the actual wedge
// case we observed (queue stalled for 3+ minutes with no progress).
const WEDGE_THRESHOLD_MS = 120_000;
/**
* Env-var override for the auto-enable safety net. Highest priority — wins
* over the persisted dashboard setting. Use to disable in an incident without
* needing dashboard access.
* RATE_LIMIT_AUTO_ENABLE=false → never auto-enable
* RATE_LIMIT_AUTO_ENABLE=true → force on regardless of dashboard
* (unset) → use dashboard setting
*/
function isAutoEnableActive(settings: RequestQueueSettings): boolean {
const env = process.env.RATE_LIMIT_AUTO_ENABLE?.trim().toLowerCase();
if (env === "false" || env === "0" || env === "off") return false;
if (env === "true" || env === "1" || env === "on") return true;
return settings.autoEnableApiKeyProviders;
}
// Sentinels for "no rate limit" / effectively infinite capacity. The reservoir
// value uses Number.MAX_SAFE_INTEGER so the bucket can never realistically be
// exhausted; maxConcurrent uses a smaller-but-still-vast ceiling since
// Bottleneck tracks concurrent jobs in memory and an unbounded number would
// risk internal counter overflow under sustained pressure.
const EFFECTIVELY_INFINITE = Number.MAX_SAFE_INTEGER;
const EFFECTIVELY_INFINITE_CONCURRENCY = 1000;
// Resolve an RPM override. 0 or missing means "infinite" (no rate cap).
function resolveRpm(override: number | undefined | null): number {
return typeof override === "number" && override > 0 ? override : EFFECTIVELY_INFINITE;
}
// Resolve a minTime override. 0 or missing means "no minimum gap".
function resolveMinTime(override: number | undefined | null): number {
return typeof override === "number" && override > 0 ? override : 0;
}
// Resolve a maxConcurrent override. 0 or missing means "effectively infinite".
function resolveMaxConcurrent(override: number | undefined | null): number {
return typeof override === "number" && override > 0 ? override : EFFECTIVELY_INFINITE_CONCURRENCY;
}
function buildLimiterDefaults() {
// 0 or missing values mean "infinite" / no rate limit applies. This treats
// the global request-queue settings the same way per-connection overrides
// are interpreted (see resolveRpm / resolveMinTime / resolveMaxConcurrent).
return {
maxConcurrent: resolveMaxConcurrent(currentRequestQueueSettings.concurrentRequests),
minTime: resolveMinTime(currentRequestQueueSettings.minTimeBetweenRequestsMs),
reservoir: resolveRpm(currentRequestQueueSettings.requestsPerMinute),
reservoirRefreshAmount: resolveRpm(currentRequestQueueSettings.requestsPerMinute),
reservoirRefreshInterval: 60 * 1000,
};
}
function updateAllLimiterSettings() {
const defaults = buildLimiterDefaults();
for (const limiter of limiters.values()) {
limiter.updateSettings(defaults);
}
}
function reconcileEnabledConnections(
connectionsRaw: unknown[],
requestQueueSettings: RequestQueueSettings
) {
const nextEnabledConnections = new Set<string>();
let explicitCount = 0;
let autoCount = 0;
for (const connRaw of connectionsRaw) {
const conn = toRecord(connRaw);
const connectionId = typeof conn.id === "string" ? conn.id : "";
const provider = typeof conn.provider === "string" ? conn.provider : "";
const isActive = conn.isActive === true;
const rateLimitProtection = conn.rateLimitProtection === true;
if (!connectionId || !provider) continue;
if (rateLimitProtection) {
nextEnabledConnections.add(connectionId);
explicitCount++;
continue;
}
if (
isAutoEnableActive(requestQueueSettings) &&
getProviderCategory(provider) === "apikey" &&
isActive
) {
nextEnabledConnections.add(connectionId);
autoCount++;
// Route through getLimiter so the `queued`/`executing` listeners and
// lastDispatchAt heartbeat are wired up — otherwise the watchdog sees
// `stalledMs = now - 0` and falsely flags healthy idle limiters as wedged.
getLimiter(provider, connectionId);
}
}
for (const connectionId of Array.from(enabledConnections)) {
if (!nextEnabledConnections.has(connectionId)) {
disableRateLimitProtection(connectionId);
}
}
for (const connectionId of nextEnabledConnections) {
enabledConnections.add(connectionId);
}
return {
explicitCount,
autoCount,
};
}
function watchdogTick() {
const now = Date.now();
// Clean up idle limiters that haven't been used recently
for (const [key, limiter] of Array.from(limiters)) {
const lastUsed = limiterLastUsed.get(key) ?? 0;
if (now - lastUsed > INACTIVE_LIMITER_MS) {
const counts = limiter.counts();
if (counts.QUEUED === 0 && counts.RUNNING === 0 && counts.EXECUTING === 0) {
limiters.delete(key);
lastDispatchAt.delete(key);
limiterLastUsed.delete(key);
logRateLimit(
`🧹 [RATE-LIMIT] Evicting idle limiter: ${key} (inactive for ${Math.round((now - lastUsed) / 1000)}s)`
);
trackAsyncOperation(limiter.disconnect());
}
}
}
for (const [key, limiter] of Array.from(limiters)) {
const counts = limiter.counts();
if (counts.QUEUED === 0) continue;
if (counts.RUNNING > 0 || counts.EXECUTING > 0) continue;
const lastDispatch = lastDispatchAt.get(key);
// No heartbeat yet → seed it and skip this tick. Prevents false wedge
// detection on a brand-new limiter or one created outside getLimiter.
if (lastDispatch === undefined) {
lastDispatchAt.set(key, now);
continue;
}
const stalledMs = now - lastDispatch;
if (stalledMs < WEDGE_THRESHOLD_MS) continue;
warnRateLimit(
`🚨 [RATE-LIMIT] WEDGED: ${key} queued=${counts.QUEUED} running=0 executing=0 stalled=${stalledMs}ms — force-resetting`
);
limiters.delete(key);
lastDispatchAt.delete(key);
limiterLastUsed.delete(key);
// Do NOT call limiter.stop() — it permanently rejects future .schedule() calls with
// "This limiter has been stopped". In-flight requests still holding a reference to
// the old instance cannot be redirected to a new one, causing spurious 502 bursts.
// Call disconnect() (not stop()) to release Bottleneck's internal heartbeat timer
// without poisoning the queue for any remaining in-flight jobs. This prevents the
// heartbeat-timer memory leak observed when many limiters are evicted at runtime.
// getLimiter() lazily allocates a fresh Bottleneck on the next call.
trackAsyncOperation(limiter.disconnect());
}
}
let shutdownHandlersRegistered = false;
export function startRateLimitWatchdog(): void {
if (watchdogInterval) return;
watchdogInterval = setInterval(watchdogTick, WATCHDOG_INTERVAL_MS);
watchdogInterval.unref?.();
// Register SIGTERM/SIGINT shutdown handlers once, lazily, on first watchdog start.
// Registering here (rather than at module load) avoids interfering with test runner
// subprocess IPC teardown — the test suite does not call startRateLimitWatchdog().
if (!shutdownHandlersRegistered) {
shutdownHandlersRegistered = true;
process.once("SIGTERM", shutdownLimiters);
process.once("SIGINT", shutdownLimiters);
}
}
export function stopRateLimitWatchdog(): void {
if (!watchdogInterval) return;
clearInterval(watchdogInterval);
watchdogInterval = null;
}
/**
* Gracefully stop all limiters for process shutdown.
* ONLY call this from SIGTERM/SIGINT handlers — not during runtime resets.
* Calling .stop() during runtime (e.g. on 429 or connection disable) permanently
* rejects future .schedule() calls, causing 502 bursts. This function is the
* sole legitimate use of limiter.stop() in this module.
*/
function shutdownLimiters(): void {
for (const limiter of limiters.values()) {
limiter.stop({ dropWaitingJobs: false });
}
limiters.clear();
lastDispatchAt.clear();
limiterLastUsed.clear();
}
// Only register shutdown handlers when there are active limiters to shut down.
// Guard with once() so repeated registrations (e.g. test resets) don't stack.
// Note: these are registered lazily in startRateLimitWatchdog() to avoid
// interfering with test runner subprocess IPC teardown.
function trackAsyncOperation<T>(promise: Promise<T>): Promise<T> {
pendingAsyncOperations.add(promise);
// Do not use a fire-and-forget `.finally()` here: it creates a derived
// Promise that mirrors rejections from `promise`. When the caller intentionally
// tracks a background cleanup without awaiting it, that derived Promise can be
// reported as an unhandled rejection during Node's test-runner IPC teardown.
void promise.then(
() => {
pendingAsyncOperations.delete(promise);
},
() => {
pendingAsyncOperations.delete(promise);
}
);
return promise;
}
/**
* Initialize rate limit protection from persisted connection settings.
* Called once on app startup.
*/
export async function initializeRateLimits() {
if (initialized) return;
initialized = true;
try {
const { getProviderConnections, getSettings } = await import("@/lib/localDb");
const [connections, settings] = await Promise.all([getProviderConnections(), getSettings()]);
const resilience = resolveResilienceSettings(settings);
currentRequestQueueSettings = { ...resilience.requestQueue };
const { explicitCount, autoCount } = reconcileEnabledConnections(
connections as unknown[],
currentRequestQueueSettings
);
updateAllLimiterSettings();
// Load per-connection rate limit overrides
connectionRateLimitOverrides.clear();
for (const conn of connections as Array<Record<string, unknown>>) {
const overrides = conn.rateLimitOverrides;
if (overrides && typeof overrides === "object" && !Array.isArray(overrides)) {
connectionRateLimitOverrides.set(String(conn.id), overrides as Record<string, number>);
}
}
if (explicitCount > 0 || autoCount > 0) {
logRateLimit(
`🛡️ [RATE-LIMIT] Loaded ${explicitCount} explicit + ${autoCount} auto-enabled protection(s)`
);
}
// Load persisted learned limits
await loadPersistedLimits();
// Watchdog runs unconditionally — cheap, only fires when something is
// actually wedged.
startRateLimitWatchdog();
} catch (err) {
errorRateLimit("[RATE-LIMIT] Failed to load settings:", err.message);
}
}
export async function applyRequestQueueSettings(nextSettings: RequestQueueSettings) {
currentRequestQueueSettings = { ...nextSettings };
const { getProviderConnections } = await import("@/lib/localDb");
const connections = await getProviderConnections();
reconcileEnabledConnections(connections as unknown[], currentRequestQueueSettings);
updateAllLimiterSettings();
}
/**
* Get or create a limiter for a given provider+connection combination
*/
export function enableRateLimitProtection(connectionId) {
enabledConnections.add(connectionId);
}
/**
* Disable rate limit protection for a connection
*/
export function disableRateLimitProtection(connectionId) {
enabledConnections.delete(connectionId);
// Evict limiters for this connection from the cache. Do NOT call limiter.stop() —
// it permanently rejects future .schedule() calls with "This limiter has been stopped",
// and in-flight requests holding a reference to the old instance would fail with 502.
// Call disconnect() (not stop()) to release Bottleneck's internal heartbeat timer
// without permanently poisoning the instance for any remaining in-flight jobs.
// Eviction-only would leak the heartbeat timer until GC; disconnect() releases it
// synchronously so the runtime memory footprint stays flat under heavy connection churn.
// .stop() is reserved exclusively for SIGTERM/SIGINT shutdown (see shutdownLimiters).
for (const [key, limiter] of Array.from(limiters)) {
if (key.includes(connectionId)) {
limiters.delete(key);
lastDispatchAt.delete(key);
limiterLastUsed.delete(key);
trackAsyncOperation(limiter.disconnect());
}
}
}
/**
* Check if rate limit protection is enabled for a connection
*/
export function isRateLimitEnabled(connectionId) {
return enabledConnections.has(connectionId);
}
/**
* Refresh per-connection rate limit overrides.
*
* Called after a PATCH update to `rateLimitOverrides` on a provider connection.
* Updates the in-memory map and evicts existing Bottleneck limiters for the
* connection so the next request gets a fresh limiter with the new settings.
*
* @param {string} connectionId
* @param {Record<string, number> | null} overrides - New overrides (null/undefined clears)
*/
export function refreshConnectionRateLimits(connectionId, overrides) {
if (overrides === null || overrides === undefined) {
connectionRateLimitOverrides.delete(connectionId);
} else {
connectionRateLimitOverrides.set(connectionId, overrides);
}
// Evict limiters referencing this connection so they get recreated on next use
for (const [key, limiter] of Array.from(limiters)) {
if (key.includes(connectionId)) {
limiters.delete(key);
lastDispatchAt.delete(key);
limiterLastUsed.delete(key);
trackAsyncOperation(limiter.disconnect());
}
}
}
/**
* Get or create a limiter for a given provider+connection combination
*/
function getLimiterKey(provider, connectionId, model = null) {
if (provider === "codex" && model) {
return `${provider}:${getCodexRateLimitKey(connectionId, model)}`;
}
// Gemini AI Studio and GitHub Copilot have per-model quotas — use model-scoped
// limiter keys so a 429 on one model doesn't pause requests for other models.
if ((provider === "gemini" || provider === "github") && model) {
return `${provider}:${connectionId}:${model}`;
}
return `${provider}:${connectionId}`;
}
function getLimiter(provider, connectionId, model = null) {
const key = getLimiterKey(provider, connectionId, model);
if (!limiters.has(key)) {
const defaults = buildLimiterDefaults();
const overrides = connectionRateLimitOverrides.get(connectionId);
if (overrides) {
// 0 (or missing) means "no override — fall through to buildLimiterDefaults()".
// Without this guard, an rpm of 0 sets reservoir=0, which Bottleneck treats
// as "depleted" and blocks ALL requests indefinitely. Treating 0 as "use
// default" lets users effectively disable per-connection limits without
// globally raising the system default.
if (typeof overrides.maxConcurrent === "number" && overrides.maxConcurrent > 0) {
defaults.maxConcurrent = overrides.maxConcurrent;
}
if (typeof overrides.minTime === "number" && overrides.minTime > 0) {
defaults.minTime = overrides.minTime;
}
if (typeof overrides.rpm === "number" && overrides.rpm > 0) {
defaults.reservoir = overrides.rpm;
defaults.reservoirRefreshAmount = overrides.rpm;
defaults.reservoirRefreshInterval = 60 * 1000;
}
// TODO: TPM/TPD integration — requires a token-bucket vs request-bucket
// separation (Bottleneck's reservoir is request-count, not token-count).
// When added, treat 0/missing the same way: fall through to system default.
}
const limiter = new Bottleneck({
...defaults,
id: key,
});
// Heartbeat: timestamp every dispatch so the watchdog can tell a healthy
// queue (just dispatched a job) from a wedged one (queue has work but
// nothing has been dispatched in a while).
limiter.on("executing", () => {
lastDispatchAt.set(key, Date.now());
});
limiters.set(key, limiter);
lastDispatchAt.set(key, Date.now());
limiterLastUsed.set(key, Date.now());
}
limiterLastUsed.set(key, Date.now());
return limiters.get(key);
}
/**
* Acquire a rate limit slot before making a request.
* If rate limiting is disabled for this connection, returns immediately.
*
* @param {string} provider - Provider ID
* @param {string} connectionId - Connection ID
* @param {string} model - Model name (optional, for per-model limits)
* @param {Function} fn - The async function to execute (e.g., executor.execute)
* @param {AbortSignal} signal - Optional abort signal to cancel waiting
* @returns {Promise<unknown>} Result of fn()
*/
export async function withRateLimit(provider, connectionId, model, fn, signal = null) {
if (!enabledConnections.has(connectionId)) {
return fn();
}
if (signal?.aborted) {
const reason = signal.reason;
if (reason instanceof Error) throw reason;
const err = new Error(typeof reason === "string" ? reason : "The operation was aborted");
err.name = "AbortError";
throw err;
}
// Proactive sliding-window fallback for header-less providers with a declared cap
// (Fase 8.2). No-op unless PROVIDER_DEFAULT_RATE_LIMITS has an entry for `provider`.
await awaitProviderDefaultSlot(
provider,
connectionId,
signal,
currentRequestQueueSettings.maxWaitMs
);
const limiter = getLimiter(provider, connectionId, model);
const maxWaitMs = currentRequestQueueSettings.maxWaitMs;
const scheduleOpts = maxWaitMs && maxWaitMs > 0 ? { expiration: maxWaitMs } : {};
try {
if (signal) {
let abortListener: (() => void) | undefined;
const abortPromise = new Promise<never>((_, reject) => {
const onAbort = () => {
const reason = signal.reason;
const err =
reason instanceof Error
? reason
: new Error(typeof reason === "string" ? reason : "The operation was aborted");
err.name = "AbortError";
reject(err);
};
if (signal.aborted) {
onAbort();
return;
}
abortListener = onAbort;
signal.addEventListener("abort", abortListener, { once: true });
});
try {
return await Promise.race([limiter.schedule(scheduleOpts, fn), abortPromise]);
} finally {
if (abortListener) {
signal.removeEventListener("abort", abortListener);
}
}
} else {
return await limiter.schedule(scheduleOpts, fn);
}
} catch (err) {
// Bottleneck's raw `This job timed out after <maxWaitMs> ms.` is
// indistinguishable from an upstream gateway timeout, so it leaks into 502
// bodies / call-log `last_error` and gets misdiagnosed as a provider outage
// (#4165). Rewrite it into a clear, OmniRoute-owned error (knob named,
// upstream disclaimed, original kept as `cause`, `code` for classification).
// Behavior is unchanged — the job is still dropped so combo can fall back.
if (err?.message?.includes("This job timed out")) {
const key = getLimiterKey(provider, connectionId, model);
logRateLimit(
`⏰ [RATE-LIMIT] ${key} — job expired after ${Math.ceil((maxWaitMs || 0) / 1000)}s in queue, dropping`
);
const queueErr = new Error(
`Request dropped after exceeding the local rate-limit queue budget maxWaitMs (${maxWaitMs}ms) for ` +
`${model ? `${provider}/${model}` : provider} — this is OmniRoute's request queue ` +
`(resilienceSettings.requestQueue.maxWaitMs), not an upstream timeout. Raise it in ` +
`Settings → Resilience if this is queue saturation rather than a slow provider.`,
{ cause: err }
) as Error & { code?: string };
queueErr.code = "RATE_LIMIT_QUEUE_TIMEOUT";
throw queueErr;
}
throw err;
}
}
/**
* Update rate limiter based on API response headers.
* Called after every successful or failed response from a provider.
*
* @param {string} provider - Provider ID
* @param {string} connectionId - Connection ID
* @param {Headers} headers - Response headers
* @param {number} status - HTTP status code
* @param {string} model - Model name
*/
export function updateFromHeaders(provider, connectionId, headers, status, model = null) {
if (!enabledConnections.has(connectionId)) return;
if (!headers) return;
const plainHeaders = toPlainHeaders(headers);
const limiter = getLimiter(provider, connectionId, model);
const headerMap =
provider === "claude" || provider === "anthropic" ? ANTHROPIC_HEADERS : STANDARD_HEADERS;
// Get header values (handle both Headers object and plain object)
const getHeader = (name: string) => {
return plainHeaders[name.toLowerCase()] || null;
};
const limit = parseInt(getHeader(headerMap.limit));
const remaining = parseInt(getHeader(headerMap.remaining));
const resetStr = getHeader(headerMap.reset);
const retryAfterStr = getHeader(headerMap.retryAfter);
const overLimit = getHeader(STANDARD_HEADERS.overLimit);
// Handle 429 — rate limited
if (status === 429) {
const retryAfterMs = parseResetTime(retryAfterStr) || 60000; // Default 60s
const counts = limiter.counts();
const limiterKey = getLimiterKey(provider, connectionId, model);
logRateLimit(
`🚫 [RATE-LIMIT] ${provider}:${connectionId.slice(0, 8)} — 429 received, pausing for ${Math.ceil(retryAfterMs / 1000)}s, dropping ${counts.QUEUED} queued request(s)`
);
// Evict from the cache so follow-up learning from the same error body
// can materialize a fresh limiter immediately. Do NOT call limiter.stop() —
// it permanently rejects future .schedule() calls with "This limiter has been stopped".
// In-flight requests holding a reference to the evicted instance will fail (they
// were already going to fail — the 429 means the API rejected them), but future
// requests will get a fresh Bottleneck instance via getLimiter().
// Call disconnect() (not stop()) to release Bottleneck's internal heartbeat timer
// without permanently poisoning the instance for any remaining in-flight jobs.
// Without disconnect() here, every 429 leaks a heartbeat timer until GC reclaims
// the abandoned Bottleneck; under sustained quota pressure that is a real leak.
limiters.delete(limiterKey);
lastDispatchAt.delete(limiterKey);
limiterLastUsed.delete(limiterKey);
trackAsyncOperation(limiter.disconnect());
return;
}
// Handle "over limit" soft warning (Fireworks)
if (overLimit === "yes") {
logRateLimit(
`⚠️ [RATE-LIMIT] ${provider}:${connectionId.slice(0, 8)} — near capacity, slowing down`
);
limiter.updateSettings({
minTime: 200, // Add 200ms between requests
});
return;
}
// Normal response — update limiter from headers
if (!isNaN(limit) && limit > 0) {
const resetMs = parseResetTime(resetStr) || 60000;
// Calculate optimal minTime from RPM limit
const minTime = Math.max(0, Math.floor(60000 / limit) - 10); // Small buffer
const updates: LimiterUpdateSettings = { minTime };
// If remaining is low (< 10% of limit), set reservoir to throttle immediately
if (!isNaN(remaining)) {
if (remaining < limit * 0.1) {
updates.reservoir = remaining;
updates.reservoirRefreshAmount = limit;
updates.reservoirRefreshInterval = resetMs;
logRateLimit(
`⚠️ [RATE-LIMIT] ${provider}:${connectionId.slice(0, 8)}${remaining}/${limit} remaining, throttling`
);
} else if (remaining > limit * 0.5) {
// Plenty of headroom — relax the limiter
updates.minTime = 0;
updates.reservoir = null;
updates.reservoirRefreshAmount = null;
updates.reservoirRefreshInterval = null;
}
}
limiter.updateSettings(updates);
// Persist learned limits (debounced)
recordLearnedLimit(
provider,
connectionId,
{ limit, remaining, minTime: updates.minTime },
model
);
}
}
/**
* Get current rate limit status for a provider+connection (for dashboard display)
*/
export function getRateLimitStatus(provider, connectionId) {
const key = `${provider}:${connectionId}`;
const limiter = limiters.get(key);
if (!limiter) {
return {
enabled: enabledConnections.has(connectionId),
active: false,
queued: 0,
running: 0,
};
}
const counts = limiter.counts();
return {
enabled: enabledConnections.has(connectionId),
active: true,
queued: counts.QUEUED || 0,
running: counts.RUNNING || 0,
executing: counts.EXECUTING || 0,
done: counts.DONE || 0,
};
}
/**
* Get all active limiters status (for dashboard overview)
*/
export function getAllRateLimitStatus() {
const result: Record<string, { queued: number; running: number; executing: number }> = {};
for (const [key, limiter] of limiters) {
const counts = limiter.counts();
result[key] = {
queued: counts.QUEUED || 0,
running: counts.RUNNING || 0,
executing: counts.EXECUTING || 0,
};
}
return result;
}
/**
* Get all learned limits (for dashboard display).
*/
export function getLearnedLimits() {
return { ...learnedLimits };
}
// ─── Persistence ────────────────────────────────────────────────────────────
async function persistLearnedLimitsNow() {
try {
const { updateSettings } = await import("@/lib/db/settings");
await updateSettings({ learnedRateLimits: JSON.stringify(learnedLimits) });
logRateLimit(
`💾 [RATE-LIMIT] Persisted learned limits for ${Object.keys(learnedLimits).length} provider(s)`
);
} catch (err) {
errorRateLimit("[RATE-LIMIT] Failed to persist learned limits:", err.message);
}
}
/**
* Record a learned limit for debounced persistence.
*/
function recordLearnedLimit(
provider: string,
connectionId: string,
limits: Partial<Omit<LearnedLimitEntry, "provider" | "connectionId" | "lastUpdated">>,
model: string | null = null
) {
const key = getLimiterKey(provider, connectionId, model);
learnedLimits[key] = {
...limits,
provider,
connectionId,
lastUpdated: Date.now(),
};
// Debounce: save at most once per PERSIST_DEBOUNCE_MS
if (!persistTimer) {
persistTimer = setTimeout(async () => {
persistTimer = null;
await trackAsyncOperation(persistLearnedLimitsNow());
}, PERSIST_DEBOUNCE_MS);
}
}
export async function __flushLearnedLimitsForTests() {
if (persistTimer) {
clearTimeout(persistTimer);
persistTimer = null;
}
await trackAsyncOperation(persistLearnedLimitsNow());
if (pendingAsyncOperations.size > 0) {
await Promise.allSettled(Array.from(pendingAsyncOperations));
}
}
export async function __resetRateLimitManagerForTests() {
if (persistTimer) {
clearTimeout(persistTimer);
persistTimer = null;
}
// Collect and await all disconnect() Promises so Bottleneck's internal
// yieldLoop(0) calls settle before the next test starts. Not awaiting
// these can cause the Node.js test runner IPC channel to receive a
// corrupted message when the pending Promise fires during IPC serialization.
const disconnectPromises: Promise<unknown>[] = [];
for (const limiter of limiters.values()) {
disconnectPromises.push(limiter.disconnect());
}
limiters.clear();
enabledConnections.clear();
initialized = false;
lastDispatchAt.clear();
limiterLastUsed.clear();
shutdownHandlersRegistered = false;
for (const key of Object.keys(learnedLimits)) {
delete learnedLimits[key];
}
if (pendingAsyncOperations.size > 0) {
await Promise.allSettled(Array.from(pendingAsyncOperations));
}
if (disconnectPromises.length > 0) {
await Promise.allSettled(disconnectPromises);
}
}
export async function __getLimiterStateForTests(provider, connectionId, model = null) {
const key = getLimiterKey(provider, connectionId, model);
const limiter = limiters.get(key);
if (!limiter) return null;
const counts = limiter.counts();
const reservoir = await limiter.currentReservoir();
return {
key,
reservoir,
queued: counts.QUEUED || 0,
running: counts.RUNNING || 0,
executing: counts.EXECUTING || 0,
done: counts.DONE || 0,
};
}
/**
* Load persisted learned limits on startup.
*/
async function loadPersistedLimits() {
try {
const { getSettings } = await import("@/lib/db/settings");
const settings = await getSettings();
const raw = settings?.learnedRateLimits;
if (typeof raw !== "string" || raw.trim().length === 0) return;
const parsed = toRecord(JSON.parse(raw) as unknown);
let count = 0;
for (const [key, dataRaw] of Object.entries(parsed)) {
const data = toRecord(dataRaw);
const lastUpdated = toNumber(data.lastUpdated, 0);
// Skip stale entries (older than 24h)
if (lastUpdated > 0 && Date.now() - lastUpdated > 24 * 60 * 60 * 1000) continue;
const connectionId = typeof data.connectionId === "string" ? data.connectionId : "";
const provider = typeof data.provider === "string" ? data.provider : "";
const limit = toNumber(data.limit, 0);
const remaining = toNumber(data.remaining, 0);
const minTime = toNumber(data.minTime, 0);
learnedLimits[key] = {
provider,
connectionId,
lastUpdated,
...(limit > 0 ? { limit } : {}),
...(remaining >= 0 ? { remaining } : {}),
...(minTime >= 0 ? { minTime } : {}),
};
// Apply to limiter if it exists and has rate limit enabled
if (connectionId && enabledConnections.has(connectionId)) {
const limiter = limiters.get(key);
if (limiter && limit > 0) {
const inferredMinTime = minTime || Math.max(0, Math.floor(60000 / limit) - 10);
limiter.updateSettings({ minTime: inferredMinTime });
count++;
}
}
}
if (count > 0) {
logRateLimit(`📥 [RATE-LIMIT] Restored ${count} learned rate limit(s) from persistence`);
}
} catch (err) {
errorRateLimit("[RATE-LIMIT] Failed to load persisted limits:", err.message);
}
}
/**
* Update rate limiter based on API response body (JSON error responses).
* Providers embed retry info in JSON payloads in different formats.
* Should be called alongside updateFromHeaders for 4xx/5xx responses.
*
* @param {string} provider - Provider ID
* @param {string} connectionId - Connection ID
* @param {string|object} responseBody - Response body (string or parsed JSON)
* @param {number} status - HTTP status code
* @param {string} model - Model name (for per-model lockouts)
*/
export function updateFromResponseBody(provider, connectionId, responseBody, status, model = null) {
if (!enabledConnections.has(connectionId)) return;
const { retryAfterMs, reason } = parseRetryAfterFromBody(responseBody);
if (retryAfterMs && retryAfterMs > 0) {
const limiter = getLimiter(provider, connectionId, model);
logRateLimit(
`🚫 [RATE-LIMIT] ${provider}:${connectionId.slice(0, 8)} — body-parsed retry: ${Math.ceil(retryAfterMs / 1000)}s (${reason})`
);
limiter.updateSettings({
reservoir: 0,
reservoirRefreshAmount: 60,
reservoirRefreshInterval: retryAfterMs,
});
}
}