mirror of
https://github.com/diegosouzapw/OmniRoute.git
synced 2026-08-23 23:52:18 +03:00
Merged — locally validated (11/11 focused tests across both new test files, typecheck:core clean after a 1-char fix pushed to this branch: ComboLogger.error is optional in combo/types.ts so the defensive race-catch needed log.error?.(...) — TS2722 otherwise). Solid production diagnosis (47 unhandledRejections traced to the orphaned race loser). Thanks!
1021 lines
37 KiB
TypeScript
1021 lines
37 KiB
TypeScript
/**
|
|
* 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 { applyBottleneckDoExpirePatch, applyBottleneckHeartbeatPatch } from "./bottleneckPatch.ts";
|
|
import { parseRetryAfterFromBody } from "./accountFallback.ts";
|
|
import { getAntigravityQuotaFamily } from "./antigravityQuotaFamily.ts";
|
|
import { getProviderCategory } from "../config/providerRegistry.ts";
|
|
import { getCodexRateLimitKey } from "../executors/codex.ts";
|
|
import { awaitProviderDefaultSlot, setProviderQuotaOverrides } 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";
|
|
import { checkQueueAdmission } from "./rateLimitManager/admission";
|
|
import {
|
|
markLocalRateLimitError,
|
|
RATE_LIMIT_EXECUTION_TIMEOUT_CODE,
|
|
RATE_LIMIT_QUEUE_WEDGED_CODE,
|
|
} from "./rateLimitManager/errors";
|
|
import { LimiterWedgeWatchdog, WATCHDOG_INTERVAL_MS } from "./rateLimitManager/wedgeWatchdog";
|
|
import { toNumber } from "@/shared/utils/numeric";
|
|
|
|
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 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 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;
|
|
export const ZAI_WEB_REQUEST_QUEUE_MAX_WAIT_MS = 60_000;
|
|
|
|
const limiterEffectiveSettings = new WeakMap<Bottleneck, Bottleneck.ConstructorOptions>();
|
|
const preservedReplacementSettings = new Map<string, Bottleneck.ConstructorOptions>();
|
|
const limiterWatchdog = new LimiterWedgeWatchdog({
|
|
limiters,
|
|
limiterLastUsed,
|
|
limiterEffectiveSettings,
|
|
preservedReplacementSettings,
|
|
trackBackground: (promise) => {
|
|
trackAsyncOperation(promise);
|
|
},
|
|
log: logRateLimit,
|
|
warn: warnRateLimit,
|
|
});
|
|
let watchdogInterval: ReturnType<typeof setInterval> | null = null;
|
|
|
|
type LimiterFactory = (options: Bottleneck.ConstructorOptions) => Bottleneck;
|
|
const defaultLimiterFactory: LimiterFactory = (options) => new Bottleneck(options);
|
|
let limiterFactory: LimiterFactory = defaultLimiterFactory;
|
|
|
|
/**
|
|
* 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;
|
|
}
|
|
|
|
export function resolveRequestQueueMaxWaitMs(
|
|
provider: string,
|
|
configuredMaxWaitMs: number = currentRequestQueueSettings.maxWaitMs
|
|
): number {
|
|
return provider.trim().toLowerCase() === "zai-web"
|
|
? Math.max(configuredMaxWaitMs, ZAI_WEB_REQUEST_QUEUE_MAX_WAIT_MS)
|
|
: configuredMaxWaitMs;
|
|
}
|
|
|
|
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 updateLimiterSettings(
|
|
limiter: Bottleneck,
|
|
updates: Bottleneck.ConstructorOptions
|
|
): Bottleneck {
|
|
const effective = limiterEffectiveSettings.get(limiter) ?? {};
|
|
limiterEffectiveSettings.set(limiter, { ...effective, ...updates });
|
|
return limiter.updateSettings(updates);
|
|
}
|
|
|
|
function updateAllLimiterSettings() {
|
|
const defaults = buildLimiterDefaults();
|
|
for (const limiter of limiters.values()) {
|
|
updateLimiterSettings(limiter, defaults);
|
|
}
|
|
}
|
|
|
|
function clearPreservedReplacementSettings(connectionId: string): void {
|
|
for (const key of preservedReplacementSettings.keys()) {
|
|
if (key.includes(connectionId)) preservedReplacementSettings.delete(key);
|
|
}
|
|
}
|
|
|
|
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 queue-progress listeners are wired up.
|
|
// Otherwise a limiter created here could not be evaluated safely by the watchdog.
|
|
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,
|
|
};
|
|
}
|
|
|
|
let shutdownHandlersRegistered = false;
|
|
|
|
export function startRateLimitWatchdog(): void {
|
|
if (watchdogInterval) return;
|
|
watchdogInterval = setInterval(() => {
|
|
const run = trackAsyncOperation(limiterWatchdog.run());
|
|
void run.then(undefined, (error) => {
|
|
errorRateLimit("[RATE-LIMIT] Watchdog scan failed:", error);
|
|
});
|
|
}, 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.
|
|
* Runtime wedge recovery also uses stop(), but only after synchronously
|
|
* removing that limiter from the cache so it can never accept new work.
|
|
*/
|
|
function shutdownLimiters(): void {
|
|
for (const limiter of limiters.values()) {
|
|
limiter.stop({ dropWaitingJobs: false });
|
|
}
|
|
limiters.clear();
|
|
limiterLastUsed.clear();
|
|
preservedReplacementSettings.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;
|
|
// Fix Bottleneck v2.19.5 doExpire bug before any limiter is created.
|
|
applyBottleneckDoExpirePatch();
|
|
applyBottleneckHeartbeatPatch();
|
|
|
|
try {
|
|
const { getCachedProviderConnections, getSettings } = await import("@/lib/localDb");
|
|
const [connections, settings] = await Promise.all([
|
|
getCachedProviderConnections(),
|
|
getSettings(),
|
|
]);
|
|
const resilience = resolveResilienceSettings(settings);
|
|
currentRequestQueueSettings = { ...resilience.requestQueue };
|
|
// #6846 Phase 1: operator overrides for header-less providers' static RPM
|
|
// budget + concurrency cap (nvidia today). No-op for every provider without
|
|
// an entry in either providerQuotaOverrides or PROVIDER_DEFAULT_RATE_LIMITS.
|
|
setProviderQuotaOverrides(resilience.providerQuotaOverrides);
|
|
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 };
|
|
// Global policy changes invalidate snapshots from the previous generation.
|
|
preservedReplacementSettings.clear();
|
|
const { getCachedProviderConnections } = await import("@/lib/localDb");
|
|
const connections = await getCachedProviderConnections();
|
|
// Also discard any snapshot created while the asynchronous DB read yielded.
|
|
preservedReplacementSettings.clear();
|
|
reconcileEnabledConnections(connections as unknown[], currentRequestQueueSettings);
|
|
updateAllLimiterSettings();
|
|
}
|
|
|
|
/**
|
|
* Get or create a limiter for a given provider+connection combination
|
|
*/
|
|
export function enableRateLimitProtection(connectionId) {
|
|
if (!enabledConnections.has(connectionId)) clearPreservedReplacementSettings(connectionId);
|
|
enabledConnections.add(connectionId);
|
|
}
|
|
|
|
/**
|
|
* Disable rate limit protection for a connection
|
|
*/
|
|
export function disableRateLimitProtection(connectionId) {
|
|
enabledConnections.delete(connectionId);
|
|
clearPreservedReplacementSettings(connectionId);
|
|
// Ordinary administrative eviction uses disconnect(), not stop(), so
|
|
// in-flight requests can finish. Wedge recovery is the deliberate exception:
|
|
// it removes the limiter from the cache first, then stops it to settle jobs
|
|
// that were already proven stranded.
|
|
for (const [key, limiter] of Array.from(limiters)) {
|
|
if (key.includes(connectionId)) {
|
|
limiters.delete(key);
|
|
limiterWatchdog.forget(limiter);
|
|
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);
|
|
}
|
|
clearPreservedReplacementSettings(connectionId);
|
|
// 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);
|
|
limiterWatchdog.forget(limiter);
|
|
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)}`;
|
|
}
|
|
if ((provider === "antigravity" || provider === "agy") && model) {
|
|
const family = getAntigravityQuotaFamily(model);
|
|
const scope = family === "other" ? model : family;
|
|
return `${provider}:${connectionId}:${scope}`;
|
|
}
|
|
// 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)) {
|
|
// Idempotent — covers callers (and tests) that reach limiter creation
|
|
// without going through initializeRateLimits().
|
|
applyBottleneckDoExpirePatch();
|
|
applyBottleneckHeartbeatPatch();
|
|
const preserved = preservedReplacementSettings.get(key);
|
|
let options: Bottleneck.ConstructorOptions;
|
|
if (preserved) {
|
|
preservedReplacementSettings.delete(key);
|
|
options = { ...preserved, id: key };
|
|
} else {
|
|
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.
|
|
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 separate token and request buckets.
|
|
}
|
|
options = { ...defaults, id: key };
|
|
}
|
|
const limiter = limiterFactory(options);
|
|
limiterEffectiveSettings.set(limiter, { ...options });
|
|
limiter.on("queued", () => {
|
|
limiterWatchdog.noteQueued(key, limiter);
|
|
});
|
|
const markQueueProgress = () => {
|
|
limiterWatchdog.noteProgress(key, limiter);
|
|
};
|
|
limiter.on("executing", markQueueProgress);
|
|
// A long-running job can leave older work queued. Start the idle grace
|
|
// from its completion, not from when that waiting work first arrived.
|
|
limiter.on("done", markQueueProgress);
|
|
|
|
limiters.set(key, limiter);
|
|
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`.
|
|
const maxWaitMs = resolveRequestQueueMaxWaitMs(provider);
|
|
await awaitProviderDefaultSlot(provider, connectionId, signal, maxWaitMs);
|
|
|
|
const limiter = getLimiter(provider, connectionId, model);
|
|
// Bottleneck's `expiration` starts only after a job leaves QUEUED. The
|
|
// legacy maxWaitMs setting therefore bounds limiter-managed execution; it
|
|
// is not a queue-wait deadline.
|
|
const executionExpirationMs = maxWaitMs;
|
|
const scheduleOpts =
|
|
executionExpirationMs && executionExpirationMs > 0 ? { expiration: executionExpirationMs } : {};
|
|
|
|
// Issue #6593: opt-in admission cap — fast-reject before Bottleneck's
|
|
// schedule() (and before any downstream compression/prompt work runs) when
|
|
// the queue is already at/over maxQueueDepth. Default 0 = disabled.
|
|
const admissionErr = checkQueueAdmission(
|
|
limiter.counts().QUEUED,
|
|
currentRequestQueueSettings.maxQueueDepth,
|
|
model ? `${provider}/${model}` : provider
|
|
);
|
|
if (admissionErr) {
|
|
logRateLimit(
|
|
`🚧 [RATE-LIMIT] ${getLimiterKey(provider, connectionId, model)} — queue full, rejecting fast (maxQueueDepth=${currentRequestQueueSettings.maxQueueDepth})`
|
|
);
|
|
throw admissionErr;
|
|
}
|
|
|
|
try {
|
|
if (signal) {
|
|
let abortListener: (() => void) | undefined;
|
|
const { promise: abortPromise, reject: rejectAbort } = Promise.withResolvers<never>();
|
|
const onAbort = () => {
|
|
const reason = signal.reason;
|
|
// Preserve native Error reasons (including AbortController's
|
|
// read-only DOMException) instead of mutating or wrapping them.
|
|
if (reason instanceof Error) {
|
|
rejectAbort(reason);
|
|
return;
|
|
}
|
|
const err = new Error(typeof reason === "string" ? reason : "The operation was aborted");
|
|
err.name = "AbortError";
|
|
if (reason !== undefined) {
|
|
(err as Error & { cause?: unknown }).cause = reason;
|
|
}
|
|
rejectAbort(err);
|
|
};
|
|
if (signal.aborted) {
|
|
onAbort();
|
|
} else {
|
|
abortListener = onAbort;
|
|
signal.addEventListener("abort", abortListener, { once: true });
|
|
}
|
|
|
|
try {
|
|
// Race the work against the abort signal. When abort wins, fn is still
|
|
// running inside Bottleneck's limiter — its eventual rejection must not
|
|
// surface as an unhandledRejection. The .catch(noop) silences only the
|
|
// orphaned branch; the real rejection comes from abortPromise.
|
|
const scheduled = limiter.schedule(scheduleOpts, fn);
|
|
scheduled.catch(() => {}); // prevent unhandledRejection when abort wins
|
|
abortPromise.catch(() => {}); // prevent unhandledRejection when scheduled wins
|
|
return await Promise.race([scheduled, abortPromise]);
|
|
} finally {
|
|
if (abortListener) {
|
|
signal.removeEventListener("abort", abortListener);
|
|
}
|
|
}
|
|
} else {
|
|
return await limiter.schedule(scheduleOpts, fn);
|
|
}
|
|
} catch (err) {
|
|
// Only Bottleneck-owned failures are rewritten. Application code can throw
|
|
// the same text and must retain its original identity and semantics.
|
|
if (
|
|
err instanceof Bottleneck.BottleneckError &&
|
|
/^This job timed out after \d+ ms\.$/.test(err.message)
|
|
) {
|
|
const key = getLimiterKey(provider, connectionId, model);
|
|
logRateLimit(
|
|
`⏰ [RATE-LIMIT] ${key} — limiter-managed execution expired after ${Math.ceil((executionExpirationMs || 0) / 1000)}s`
|
|
);
|
|
throw markLocalRateLimitError(
|
|
new Error(
|
|
`Request exceeded OmniRoute's local rate-limit execution expiration ` +
|
|
`(legacy resilienceSettings.requestQueue.maxWaitMs=${executionExpirationMs}ms) for ` +
|
|
`${model ? `${provider}/${model}` : provider}. Bottleneck applies this deadline only ` +
|
|
`after dispatch; it does not bound queue wait and is not an upstream-generated timeout.`,
|
|
{ cause: err }
|
|
),
|
|
RATE_LIMIT_EXECUTION_TIMEOUT_CODE
|
|
);
|
|
}
|
|
|
|
if (
|
|
err instanceof Bottleneck.BottleneckError &&
|
|
err.message === "rate-limit-watchdog-wedge-reset"
|
|
) {
|
|
const cleanup = limiterWatchdog.getEviction(limiter);
|
|
if (!cleanup) throw err;
|
|
|
|
let cleanupError: unknown;
|
|
try {
|
|
await cleanup;
|
|
} catch (error) {
|
|
cleanupError = error;
|
|
errorRateLimit("[RATE-LIMIT] Wedge cleanup failed:", error);
|
|
}
|
|
|
|
const key = getLimiterKey(provider, connectionId, model);
|
|
logRateLimit(`↪️ [RATE-LIMIT] ${key} — surfacing local wedge; caller will not be replayed`);
|
|
const wedgeErr = new Error(
|
|
`Request dropped: the local rate-limit queue for ${model ? `${provider}/${model}` : provider} ` +
|
|
`was detected as wedged (stalled with nothing executing) and force-reset. OmniRoute does ` +
|
|
`not replay dropped work automatically; combo routing may fall back to another target.`,
|
|
{ cause: err }
|
|
) as Error & { cleanupError?: unknown };
|
|
if (cleanupError !== undefined) wedgeErr.cleanupError = cleanupError;
|
|
throw markLocalRateLimitError(wedgeErr, RATE_LIMIT_QUEUE_WEDGED_CODE);
|
|
}
|
|
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);
|
|
limiterWatchdog.forget(limiter);
|
|
limiterLastUsed.delete(limiterKey);
|
|
preservedReplacementSettings.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`
|
|
);
|
|
updateLimiterSettings(limiter, {
|
|
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;
|
|
}
|
|
}
|
|
|
|
updateLimiterSettings(limiter, 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 function __setLimiterFactoryForTests(factory: LimiterFactory): void {
|
|
limiterFactory = factory;
|
|
}
|
|
|
|
export async function __runLimiterWatchdogForTests(now = Date.now()): Promise<void> {
|
|
await limiterWatchdog.run(now);
|
|
}
|
|
|
|
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;
|
|
limiterLastUsed.clear();
|
|
preservedReplacementSettings.clear();
|
|
limiterFactory = defaultLimiterFactory;
|
|
limiterWatchdog.reset();
|
|
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);
|
|
updateLimiterSettings(limiter, { 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})`
|
|
);
|
|
|
|
updateLimiterSettings(limiter, {
|
|
reservoir: 0,
|
|
reservoirRefreshAmount: 60,
|
|
reservoirRefreshInterval: retryAfterMs,
|
|
});
|
|
}
|
|
}
|