@limitrate/core
Version:
Core rate limiting and cost control engine for LimitRate
1,733 lines (1,678 loc) • 78.6 kB
JavaScript
"use strict";
var __create = Object.create;
var __defProp = Object.defineProperty;
var __getOwnPropDesc = Object.getOwnPropertyDescriptor;
var __getOwnPropNames = Object.getOwnPropertyNames;
var __getProtoOf = Object.getPrototypeOf;
var __hasOwnProp = Object.prototype.hasOwnProperty;
var __esm = (fn, res) => function __init() {
return fn && (res = (0, fn[__getOwnPropNames(fn)[0]])(fn = 0)), res;
};
var __export = (target, all) => {
for (var name in all)
__defProp(target, name, { get: all[name], enumerable: true });
};
var __copyProps = (to, from, except, desc) => {
if (from && typeof from === "object" || typeof from === "function") {
for (let key of __getOwnPropNames(from))
if (!__hasOwnProp.call(to, key) && key !== except)
__defProp(to, key, { get: () => from[key], enumerable: !(desc = __getOwnPropDesc(from, key)) || desc.enumerable });
}
return to;
};
var __toESM = (mod, isNodeMode, target) => (target = mod != null ? __create(__getProtoOf(mod)) : {}, __copyProps(
// If the importer is in node compatibility mode or this is not an ESM
// file that has been converted to a CommonJS file using a Babel-
// compatible transform (i.e. "__esModule" has not been set), then set
// "default" to the CommonJS "module.exports" for node compatibility.
isNodeMode || !mod || !mod.__esModule ? __defProp(target, "default", { value: mod, enumerable: true }) : target,
mod
));
var __toCommonJS = (mod) => __copyProps(__defProp({}, "__esModule", { value: true }), mod);
// src/tokenizers/openai.ts
var openai_exports = {};
__export(openai_exports, {
createOpenAITokenizer: () => createOpenAITokenizer
});
async function createOpenAITokenizer(model) {
const tiktoken = await import("tiktoken");
const modelMap = {
"gpt-3.5-turbo": "gpt-3.5-turbo",
"gpt-3.5-turbo-0301": "gpt-3.5-turbo-0301",
"gpt-3.5-turbo-0613": "gpt-3.5-turbo-0613",
"gpt-3.5-turbo-1106": "gpt-3.5-turbo-1106",
"gpt-3.5-turbo-16k": "gpt-3.5-turbo-16k",
"gpt-4": "gpt-4",
"gpt-4-0314": "gpt-4-0314",
"gpt-4-0613": "gpt-4-0613",
"gpt-4-32k": "gpt-4-32k",
"gpt-4-turbo": "gpt-4-turbo",
"gpt-4-turbo-preview": "gpt-4-turbo-preview",
"gpt-4o": "gpt-4o",
"gpt-4o-mini": "gpt-4o-mini"
};
const encodingModel = modelMap[model] || "gpt-3.5-turbo";
let encoding;
try {
encoding = tiktoken.encoding_for_model(encodingModel);
} catch (error) {
console.warn(`[LimitRate] Unknown OpenAI model "${model}", using cl100k_base encoding`);
encoding = tiktoken.get_encoding("cl100k_base");
}
return {
count: async (text) => {
const combined = Array.isArray(text) ? text.join("\n") : text;
const tokens = encoding.encode(combined);
return tokens.length;
},
model,
isFallback: false
};
}
var init_openai = __esm({
"src/tokenizers/openai.ts"() {
"use strict";
}
});
// src/tokenizers/anthropic.ts
var anthropic_exports = {};
__export(anthropic_exports, {
createAnthropicTokenizer: () => createAnthropicTokenizer
});
async function createAnthropicTokenizer(model) {
const AnthropicModule = await import("@anthropic-ai/sdk");
const Anthropic = AnthropicModule.default || AnthropicModule;
const client = new Anthropic({
apiKey: "dummy-key-for-counting"
// API key not needed for countTokens
});
const modelMap = {
"claude-3-opus": "claude-3-opus-20240229",
"claude-3-opus-20240229": "claude-3-opus-20240229",
"claude-3-sonnet": "claude-3-sonnet-20240229",
"claude-3-sonnet-20240229": "claude-3-sonnet-20240229",
"claude-3-haiku": "claude-3-haiku-20240307",
"claude-3-haiku-20240307": "claude-3-haiku-20240307",
"claude-3-5-sonnet": "claude-3-5-sonnet-20240620",
"claude-3-5-sonnet-20240620": "claude-3-5-sonnet-20240620",
"claude-3-5-sonnet-20241022": "claude-3-5-sonnet-20241022"
};
const fullModel = modelMap[model] || model;
return {
count: async (text) => {
const combined = Array.isArray(text) ? text.join("\n") : text;
try {
const result = await client.beta.messages.countTokens({
model: fullModel,
messages: [{ role: "user", content: combined }]
});
return result.input_tokens || 0;
} catch (error) {
if (error.message?.includes("Could not resolve")) {
console.warn(`[LimitRate] Anthropic token counting failed (${error.message}), using fallback`);
return Math.ceil(combined.length / 4);
}
throw error;
}
},
model,
isFallback: false
};
}
var init_anthropic = __esm({
"src/tokenizers/anthropic.ts"() {
"use strict";
}
});
// src/index.ts
var index_exports = {};
__export(index_exports, {
ConcurrencyLimiter: () => ConcurrencyLimiter,
EndpointTracker: () => EndpointTracker,
EventEmitter: () => EventEmitter,
MODEL_LIMITS: () => MODEL_LIMITS,
MemoryStore: () => MemoryStore,
PenaltyManager: () => PenaltyManager,
PolicyEngine: () => PolicyEngine,
RedisStore: () => RedisStore,
StreamingTracker: () => StreamingTracker,
UpstashStore: () => UpstashStore,
ValidationError: () => ValidationError,
clearAllLimiters: () => clearAllLimiters,
clearTokenizerCache: () => clearTokenizerCache,
createEndpointKey: () => createEndpointKey,
createSharedMemoryStore: () => createSharedMemoryStore,
createSharedRedisStore: () => createSharedRedisStore,
createSharedUpstashStore: () => createSharedUpstashStore,
createStore: () => createStore,
createTokenizer: () => createTokenizer,
estimateTokens: () => estimateTokens,
extractIP: () => extractIP,
formatValidationError: () => formatValidationError,
getConcurrencyLimiter: () => getConcurrencyLimiter,
getGlobalEndpointTracker: () => getGlobalEndpointTracker,
getModelLimits: () => getModelLimits,
getSuggestedAlternatives: () => getSuggestedAlternatives,
isIPInList: () => isIPInList,
normalizeRoutePath: () => normalizeRoutePath,
parseAnthropicChunk: () => parseAnthropicChunk,
parseOpenAIChunk: () => parseOpenAIChunk,
setGlobalEndpointTracker: () => setGlobalEndpointTracker,
validateIPList: () => validateIPList,
validatePolicyConfig: () => validatePolicyConfig,
validatePrompt: () => validatePrompt,
validateStoreConfig: () => validateStoreConfig
});
module.exports = __toCommonJS(index_exports);
// src/stores/memory.ts
function isCacheEntry(entry) {
return "count" in entry;
}
var MemoryStore = class {
cache;
maxKeys;
cleanupInterval;
constructor(options = {}) {
this.cache = /* @__PURE__ */ new Map();
this.maxKeys = options.maxKeys ?? 1e4;
this.cleanupInterval = null;
const intervalMs = options.cleanupIntervalMs ?? 6e4;
this.cleanupInterval = setInterval(() => this.cleanup(), intervalMs);
}
async checkRate(key, limit, windowSeconds, burst) {
const now = Date.now();
const entry = this.cache.get(key);
if (entry && !isCacheEntry(entry)) {
const expiresAt = now + windowSeconds * 1e3;
const burstTokens = burst;
const cacheEntry = { count: 1, expiresAt, burstTokens };
this.cache.set(key, cacheEntry);
this.evictIfNeeded();
return {
allowed: true,
current: 1,
remaining: limit - 1,
resetInSeconds: windowSeconds,
limit,
burstTokens
};
}
if (!entry || entry.expiresAt <= now) {
const expiresAt = now + windowSeconds * 1e3;
const burstTokens = burst;
this.cache.set(key, { count: 1, expiresAt, burstTokens });
this.evictIfNeeded();
return {
allowed: true,
current: 1,
remaining: limit - 1,
resetInSeconds: windowSeconds,
limit,
burstTokens
};
}
const current = entry.count;
const resetInSeconds = Math.ceil((entry.expiresAt - now) / 1e3);
if (current < limit) {
entry.count += 1;
return {
allowed: true,
current: entry.count,
remaining: limit - entry.count,
resetInSeconds,
limit,
burstTokens: entry.burstTokens
};
}
if (burst !== void 0 && entry.burstTokens !== void 0 && entry.burstTokens > 0) {
entry.burstTokens -= 1;
entry.count += 1;
return {
allowed: true,
current: entry.count,
remaining: 0,
resetInSeconds,
limit,
burstTokens: entry.burstTokens
};
}
return {
allowed: false,
current,
remaining: 0,
resetInSeconds,
limit,
burstTokens: entry.burstTokens ?? 0
};
}
async peekRate(key, limit, windowSeconds) {
const now = Date.now();
const entry = this.cache.get(key);
if (entry && !isCacheEntry(entry)) {
return {
allowed: true,
current: 0,
remaining: limit,
resetInSeconds: windowSeconds,
limit
};
}
if (!entry || entry.expiresAt <= now) {
return {
allowed: true,
current: 0,
remaining: limit,
resetInSeconds: windowSeconds,
limit
};
}
const current = entry.count;
const resetInSeconds = Math.ceil((entry.expiresAt - now) / 1e3);
const remaining = Math.max(0, limit - current);
return {
allowed: current < limit,
current,
remaining,
resetInSeconds,
limit,
burstTokens: entry.burstTokens
};
}
async incrementCost(key, cost, windowSeconds, cap) {
const now = Date.now();
const entry = this.cache.get(key);
if (entry && !isCacheEntry(entry)) {
const expiresAt = now + windowSeconds * 1e3;
const cacheEntry = { count: cost, expiresAt };
this.cache.set(key, cacheEntry);
this.evictIfNeeded();
return {
allowed: cost <= cap,
current: cost,
remaining: cap - cost,
resetInSeconds: windowSeconds,
cap
};
}
if (!entry || entry.expiresAt <= now) {
const expiresAt = now + windowSeconds * 1e3;
this.cache.set(key, { count: cost, expiresAt });
this.evictIfNeeded();
return {
allowed: cost <= cap,
current: cost,
remaining: cap - cost,
resetInSeconds: windowSeconds,
cap
};
}
const newCost = entry.count + cost;
const resetInSeconds = Math.ceil((entry.expiresAt - now) / 1e3);
if (newCost > cap) {
return {
allowed: false,
current: entry.count,
remaining: cap - entry.count,
resetInSeconds,
cap
};
}
entry.count = newCost;
return {
allowed: true,
current: newCost,
remaining: cap - newCost,
resetInSeconds,
cap
};
}
async incrementTokens(key, tokens, windowSeconds, limit) {
const now = Date.now();
const entry = this.cache.get(key);
if (entry && !isCacheEntry(entry)) {
const expiresAt = now + windowSeconds * 1e3;
const cacheEntry = { count: tokens, expiresAt };
this.cache.set(key, cacheEntry);
this.evictIfNeeded();
return {
allowed: tokens <= limit,
current: tokens,
remaining: Math.max(0, limit - tokens),
resetInSeconds: windowSeconds,
limit
};
}
if (!entry || entry.expiresAt <= now) {
const expiresAt = now + windowSeconds * 1e3;
this.cache.set(key, { count: tokens, expiresAt });
this.evictIfNeeded();
return {
allowed: tokens <= limit,
current: tokens,
remaining: Math.max(0, limit - tokens),
resetInSeconds: windowSeconds,
limit
};
}
const newTokens = entry.count + tokens;
const resetInSeconds = Math.ceil((entry.expiresAt - now) / 1e3);
if (newTokens > limit) {
return {
allowed: false,
current: entry.count,
remaining: Math.max(0, limit - entry.count),
resetInSeconds,
limit
};
}
entry.count = newTokens;
return {
allowed: true,
current: newTokens,
remaining: limit - newTokens,
resetInSeconds,
limit
};
}
async ping() {
return true;
}
async close() {
if (this.cleanupInterval) {
clearInterval(this.cleanupInterval);
this.cleanupInterval = null;
}
this.cache.clear();
}
/**
* Remove expired entries
*/
cleanup() {
const now = Date.now();
for (const [key, entry] of this.cache.entries()) {
if (entry.expiresAt <= now) {
this.cache.delete(key);
}
}
}
/**
* Evict oldest entry if cache is full (simple LRU)
*/
evictIfNeeded() {
if (this.cache.size > this.maxKeys) {
const firstKey = this.cache.keys().next().value;
if (firstKey) {
this.cache.delete(firstKey);
}
}
}
/**
* Generic get method for arbitrary data (v2.0.0 - D4)
*/
async get(key) {
const now = Date.now();
const entry = this.cache.get(key);
if (!entry) {
return null;
}
if ("value" in entry) {
const genericEntry = entry;
if (genericEntry.expiresAt <= now) {
this.cache.delete(key);
return null;
}
return genericEntry.value;
}
return null;
}
/**
* Generic set method for arbitrary data (v2.0.0 - D4)
*/
async set(key, value, ttl) {
const now = Date.now();
const expiresAt = ttl ? now + ttl * 1e3 : now + 86400 * 1e3;
const entry = {
value,
expiresAt
};
this.cache.set(key, entry);
this.evictIfNeeded();
}
/**
* Generic delete method (v2.0.0 - D4)
*/
async delete(key) {
this.cache.delete(key);
}
/**
* Get current cache size (for testing/debugging)
*/
getCacheSize() {
return this.cache.size;
}
};
// src/stores/redis.ts
var import_ioredis = __toESM(require("ioredis"));
var RATE_PEEK_SCRIPT = `
local key = KEYS[1]
local limit = tonumber(ARGV[1])
local window = tonumber(ARGV[2])
local current = redis.call('GET', key)
if not current then
return {0, limit, window, limit}
end
current = tonumber(current)
local ttl = redis.call('TTL', key)
local remaining = math.max(0, limit - current)
return {current, remaining, ttl, limit}
`;
var RATE_CHECK_SCRIPT = `
local key = KEYS[1]
local limit = tonumber(ARGV[1])
local window = tonumber(ARGV[2])
local now = tonumber(ARGV[3])
local current = redis.call('GET', key)
if not current then
redis.call('SETEX', key, window, 1)
return {1, limit - 1, window, limit}
end
current = tonumber(current)
if current >= limit then
local ttl = redis.call('TTL', key)
return {current, 0, ttl, limit}
end
redis.call('INCR', key)
local ttl = redis.call('TTL', key)
return {current + 1, limit - current - 1, ttl, limit}
`;
var RATE_CHECK_BURST_SCRIPT = `
local rateKey = KEYS[1]
local burstKey = KEYS[2]
local limit = tonumber(ARGV[1])
local window = tonumber(ARGV[2])
local now = tonumber(ARGV[3])
local burst = tonumber(ARGV[4])
local current = redis.call('GET', rateKey)
-- No entry: create new with initial burst tokens
if not current then
redis.call('SETEX', rateKey, window, 1)
redis.call('SETEX', burstKey, window, burst)
return {1, limit - 1, window, limit, burst}
end
current = tonumber(current)
-- Within limit: increment and return
if current < limit then
redis.call('INCR', rateKey)
local ttl = redis.call('TTL', rateKey)
local burstTokens = tonumber(redis.call('GET', burstKey) or 0)
return {current + 1, limit - current - 1, ttl, limit, burstTokens}
end
-- Limit reached: try burst
local burstTokens = tonumber(redis.call('GET', burstKey) or 0)
if burstTokens > 0 then
redis.call('DECR', burstKey)
redis.call('INCR', rateKey)
local ttl = redis.call('TTL', rateKey)
return {current + 1, 0, ttl, limit, burstTokens - 1}
end
-- Both exhausted
local ttl = redis.call('TTL', rateKey)
return {current, 0, ttl, limit, 0}
`;
var COST_INCREMENT_SCRIPT = `
local key = KEYS[1]
local cost = tonumber(ARGV[1])
local window = tonumber(ARGV[2])
local cap = tonumber(ARGV[3])
local current = redis.call('GET', key)
if not current then
if cost > cap then
redis.call('SETEX', key, window, 0)
return {0, false, cap, window, cap}
end
redis.call('SETEX', key, window, cost)
return {cost, true, cap - cost, window, cap}
end
current = tonumber(current)
local newCost = current + cost
if newCost > cap then
local ttl = redis.call('TTL', key)
return {current, false, cap - current, ttl, cap}
end
redis.call('SET', key, newCost, 'KEEPTTL')
local ttl = redis.call('TTL', key)
return {newCost, true, cap - newCost, ttl, cap}
`;
var TOKEN_INCREMENT_SCRIPT = `
local key = KEYS[1]
local tokens = tonumber(ARGV[1])
local window = tonumber(ARGV[2])
local limit = tonumber(ARGV[3])
local current = redis.call('GET', key)
if not current then
if tokens > limit then
redis.call('SETEX', key, window, 0)
return {0, false, limit, window, limit}
end
redis.call('SETEX', key, window, tokens)
return {tokens, true, limit - tokens, window, limit}
end
current = tonumber(current)
local newTokens = current + tokens
if newTokens > limit then
local ttl = redis.call('TTL', key)
return {current, false, math.max(0, limit - current), ttl, limit}
end
redis.call('SET', key, newTokens, 'KEEPTTL')
local ttl = redis.call('TTL', key)
return {newTokens, true, limit - newTokens, ttl, limit}
`;
var RedisStore = class {
client;
ownClient;
keyPrefix;
constructor(options = {}) {
this.keyPrefix = options.keyPrefix ?? "limitrate:";
this.ownClient = false;
if (options.client instanceof import_ioredis.default) {
this.client = options.client;
} else if (typeof options.client === "string") {
this.client = new import_ioredis.default(options.client, options.redisOptions);
this.ownClient = true;
} else {
throw new Error("RedisStore requires either a Redis URL or ioredis instance");
}
}
async checkRate(key, limit, windowSeconds, burst) {
const now = Math.floor(Date.now() / 1e3);
try {
if (burst !== void 0) {
const rateKey = `${this.keyPrefix}rate:${key}`;
const burstKey = `${this.keyPrefix}burst:${key}`;
const result2 = await this.client.eval(
RATE_CHECK_BURST_SCRIPT,
2,
rateKey,
burstKey,
limit.toString(),
windowSeconds.toString(),
now.toString(),
burst.toString()
);
const [current2, remaining2, resetInSeconds2, returnedLimit2, burstTokens] = result2;
return {
allowed: current2 <= limit + burst,
current: current2,
remaining: Math.max(0, remaining2),
resetInSeconds: Math.max(1, resetInSeconds2),
limit: returnedLimit2,
burstTokens
};
}
const prefixedKey = `${this.keyPrefix}rate:${key}`;
const result = await this.client.eval(
RATE_CHECK_SCRIPT,
1,
prefixedKey,
limit.toString(),
windowSeconds.toString(),
now.toString()
);
const [current, remaining, resetInSeconds, returnedLimit] = result;
return {
allowed: current <= limit,
current,
remaining: Math.max(0, remaining),
resetInSeconds: Math.max(1, resetInSeconds),
limit: returnedLimit
};
} catch (error) {
console.error("[LimitRate] Redis rate check error:", error);
throw error;
}
}
async peekRate(key, limit, windowSeconds) {
const prefixedKey = `${this.keyPrefix}rate:${key}`;
try {
const result = await this.client.eval(
RATE_PEEK_SCRIPT,
1,
prefixedKey,
limit.toString(),
windowSeconds.toString()
);
const [current, remaining, resetInSeconds, returnedLimit] = result;
return {
allowed: current < limit,
current,
remaining: Math.max(0, remaining),
resetInSeconds: Math.max(1, resetInSeconds),
limit: returnedLimit
};
} catch (error) {
console.error("[LimitRate] Redis peek rate error:", error);
throw error;
}
}
async incrementCost(key, cost, windowSeconds, cap) {
const prefixedKey = `${this.keyPrefix}cost:${key}`;
try {
const result = await this.client.eval(
COST_INCREMENT_SCRIPT,
1,
prefixedKey,
cost.toString(),
windowSeconds.toString(),
cap.toString()
);
const [current, allowed, remaining, resetInSeconds, returnedCap] = result;
return {
allowed: allowed === 1,
current,
remaining: Math.max(0, remaining),
resetInSeconds: Math.max(1, resetInSeconds),
cap: returnedCap
};
} catch (error) {
console.error("[LimitRate] Redis cost increment error:", error);
throw error;
}
}
async incrementTokens(key, tokens, windowSeconds, limit) {
const prefixedKey = `${this.keyPrefix}tokens:${key}`;
try {
const result = await this.client.eval(
TOKEN_INCREMENT_SCRIPT,
1,
prefixedKey,
tokens.toString(),
windowSeconds.toString(),
limit.toString()
);
const [current, allowed, remaining, resetInSeconds, returnedLimit] = result;
return {
allowed: allowed === 1,
current,
remaining: Math.max(0, remaining),
resetInSeconds: Math.max(1, resetInSeconds),
limit: returnedLimit
};
} catch (error) {
console.error("[LimitRate] Redis token increment error:", error);
throw error;
}
}
async ping() {
try {
const result = await this.client.ping();
return result === "PONG";
} catch {
return false;
}
}
async close() {
if (this.ownClient) {
await this.client.quit();
}
}
/**
* Generic get method for arbitrary data (v2.0.0 - D4)
*/
async get(key) {
try {
const prefixedKey = `${this.keyPrefix}generic:${key}`;
const value = await this.client.get(prefixedKey);
if (!value) {
return null;
}
return JSON.parse(value);
} catch (error) {
console.error("[LimitRate] Redis get error:", error);
throw error;
}
}
/**
* Generic set method for arbitrary data (v2.0.0 - D4)
*/
async set(key, value, ttl) {
try {
const prefixedKey = `${this.keyPrefix}generic:${key}`;
const serialized = JSON.stringify(value);
if (ttl) {
await this.client.setex(prefixedKey, ttl, serialized);
} else {
await this.client.setex(prefixedKey, 86400, serialized);
}
} catch (error) {
console.error("[LimitRate] Redis set error:", error);
throw error;
}
}
/**
* Generic delete method (v2.0.0 - D4)
*/
async delete(key) {
try {
const prefixedKey = `${this.keyPrefix}generic:${key}`;
await this.client.del(prefixedKey);
} catch (error) {
console.error("[LimitRate] Redis delete error:", error);
throw error;
}
}
/**
* Get underlying Redis client (for advanced use cases)
*/
getClient() {
return this.client;
}
};
// src/stores/upstash.ts
var import_redis = require("@upstash/redis");
var RATE_CHECK_SCRIPT2 = `
local key = KEYS[1]
local limit = tonumber(ARGV[1])
local window = tonumber(ARGV[2])
local now = tonumber(ARGV[3])
local current = redis.call('GET', key)
if not current then
redis.call('SETEX', key, window, 1)
return {1, limit - 1, window, limit}
end
current = tonumber(current)
if current >= limit then
local ttl = redis.call('TTL', key)
return {current, 0, ttl, limit}
end
redis.call('INCR', key)
local ttl = redis.call('TTL', key)
return {current + 1, limit - current - 1, ttl, limit}
`;
var RATE_CHECK_BURST_SCRIPT2 = `
local rateKey = KEYS[1]
local burstKey = KEYS[2]
local limit = tonumber(ARGV[1])
local window = tonumber(ARGV[2])
local now = tonumber(ARGV[3])
local burst = tonumber(ARGV[4])
local current = redis.call('GET', rateKey)
if not current then
redis.call('SETEX', rateKey, window, 1)
redis.call('SETEX', burstKey, window, burst)
return {1, limit - 1, window, limit, burst}
end
current = tonumber(current)
if current < limit then
redis.call('INCR', rateKey)
local ttl = redis.call('TTL', rateKey)
local burstTokens = tonumber(redis.call('GET', burstKey) or 0)
return {current + 1, limit - current - 1, ttl, limit, burstTokens}
end
local burstTokens = tonumber(redis.call('GET', burstKey) or 0)
if burstTokens > 0 then
redis.call('DECR', burstKey)
redis.call('INCR', rateKey)
local ttl = redis.call('TTL', rateKey)
return {current + 1, 0, ttl, limit, burstTokens - 1}
end
local ttl = redis.call('TTL', rateKey)
return {current, 0, ttl, limit, 0}
`;
var COST_INCREMENT_SCRIPT2 = `
local key = KEYS[1]
local cost = tonumber(ARGV[1])
local window = tonumber(ARGV[2])
local cap = tonumber(ARGV[3])
local current = redis.call('GET', key)
if not current then
if cost > cap then
redis.call('SETEX', key, window, 0)
return {0, false, cap, window, cap}
end
redis.call('SETEX', key, window, cost)
return {cost, true, cap - cost, window, cap}
end
current = tonumber(current)
local newCost = current + cost
if newCost > cap then
local ttl = redis.call('TTL', key)
return {current, false, cap - current, ttl, cap}
end
redis.call('SET', key, newCost, 'KEEPTTL')
local ttl = redis.call('TTL', key)
return {newCost, true, cap - newCost, ttl, cap}
`;
var TOKEN_INCREMENT_SCRIPT2 = `
local key = KEYS[1]
local tokens = tonumber(ARGV[1])
local window = tonumber(ARGV[2])
local limit = tonumber(ARGV[3])
local current = redis.call('GET', key)
if not current then
if tokens > limit then
redis.call('SETEX', key, window, 0)
return {0, false, limit, window, limit}
end
redis.call('SETEX', key, window, tokens)
return {tokens, true, limit - tokens, window, limit}
end
current = tonumber(current)
local newTokens = current + tokens
if newTokens > limit then
local ttl = redis.call('TTL', key)
return {current, false, math.max(0, limit - current), ttl, limit}
end
redis.call('SET', key, newTokens, 'KEEPTTL')
local ttl = redis.call('TTL', key)
return {newTokens, true, limit - newTokens, ttl, limit}
`;
var UpstashStore = class {
client;
keyPrefix;
url;
token;
constructor(options) {
if (!options.url || !options.token) {
throw new Error("UpstashStore requires both url and token");
}
this.keyPrefix = options.keyPrefix ?? "limitrate:";
this.url = options.url;
this.token = options.token;
this.client = new import_redis.Redis({
url: options.url,
token: options.token
});
}
async checkRate(key, limit, windowSeconds, burst) {
const now = Math.floor(Date.now() / 1e3);
try {
if (burst !== void 0) {
const rateKey = `${this.keyPrefix}rate:${key}`;
const burstKey = `${this.keyPrefix}burst:${key}`;
const result2 = await this.client.eval(
RATE_CHECK_BURST_SCRIPT2,
[rateKey, burstKey],
[limit.toString(), windowSeconds.toString(), now.toString(), burst.toString()]
);
const [current2, remaining2, resetInSeconds2, returnedLimit2, burstTokens] = result2;
return {
allowed: current2 <= limit + burst,
current: current2,
remaining: Math.max(0, remaining2),
resetInSeconds: Math.max(1, resetInSeconds2),
limit: returnedLimit2,
burstTokens
};
}
const prefixedKey = `${this.keyPrefix}rate:${key}`;
const result = await this.client.eval(
RATE_CHECK_SCRIPT2,
[prefixedKey],
[limit.toString(), windowSeconds.toString(), now.toString()]
);
const [current, remaining, resetInSeconds, returnedLimit] = result;
return {
allowed: current <= limit,
current,
remaining: Math.max(0, remaining),
resetInSeconds: Math.max(1, resetInSeconds),
limit: returnedLimit
};
} catch (error) {
console.error("[LimitRate] Upstash rate check error:", error);
throw error;
}
}
async peekRate(key, limit, windowSeconds) {
const prefixedKey = `${this.keyPrefix}rate:${key}`;
try {
const response = await fetch(`${this.url}/get/${encodeURIComponent(prefixedKey)}`, {
headers: { "Authorization": `Bearer ${this.token}` }
});
if (!response.ok && response.status !== 404) {
throw new Error(`Upstash API error: ${response.status}`);
}
const data = await response.json();
const current = data.result ? parseInt(String(data.result), 10) : 0;
const ttlResponse = await fetch(`${this.url}/ttl/${encodeURIComponent(prefixedKey)}`, {
headers: { "Authorization": `Bearer ${this.token}` }
});
const ttlData = await ttlResponse.json();
const resetInSeconds = ttlData.result && ttlData.result > 0 ? ttlData.result : windowSeconds;
const remaining = Math.max(0, limit - current);
return {
allowed: current < limit,
current,
remaining,
resetInSeconds,
limit
};
} catch (error) {
console.error("[LimitRate] Upstash peek rate error:", error);
throw error;
}
}
async incrementCost(key, cost, windowSeconds, cap) {
const prefixedKey = `${this.keyPrefix}cost:${key}`;
try {
const result = await this.client.eval(
COST_INCREMENT_SCRIPT2,
[prefixedKey],
[cost.toString(), windowSeconds.toString(), cap.toString()]
);
const [current, allowed, remaining, resetInSeconds, returnedCap] = result;
return {
allowed: allowed === 1,
current,
remaining: Math.max(0, remaining),
resetInSeconds: Math.max(1, resetInSeconds),
cap: returnedCap
};
} catch (error) {
console.error("[LimitRate] Upstash cost increment error:", error);
throw error;
}
}
async incrementTokens(key, tokens, windowSeconds, limit) {
const prefixedKey = `${this.keyPrefix}tokens:${key}`;
try {
const result = await this.client.eval(
TOKEN_INCREMENT_SCRIPT2,
[prefixedKey],
[tokens.toString(), windowSeconds.toString(), limit.toString()]
);
const [current, allowed, remaining, resetInSeconds, returnedLimit] = result;
return {
allowed: allowed === 1,
current,
remaining: Math.max(0, remaining),
resetInSeconds: Math.max(1, resetInSeconds),
limit: returnedLimit
};
} catch (error) {
console.error("[LimitRate] Upstash token increment error:", error);
throw error;
}
}
async ping() {
try {
const result = await this.client.ping();
return result === "PONG";
} catch {
return false;
}
}
async close() {
return Promise.resolve();
}
/**
* Generic get method for arbitrary data (v2.0.0 - D4)
*/
async get(key) {
try {
const prefixedKey = `${this.keyPrefix}generic:${key}`;
const value = await this.client.get(prefixedKey);
if (!value) {
return null;
}
return value;
} catch (error) {
console.error("[LimitRate] Upstash get error:", error);
throw error;
}
}
/**
* Generic set method for arbitrary data (v2.0.0 - D4)
*/
async set(key, value, ttl) {
try {
const prefixedKey = `${this.keyPrefix}generic:${key}`;
const expirySeconds = ttl || 86400;
await this.client.setex(prefixedKey, expirySeconds, value);
} catch (error) {
console.error("[LimitRate] Upstash set error:", error);
throw error;
}
}
/**
* Generic delete method (v2.0.0 - D4)
*/
async delete(key) {
try {
const prefixedKey = `${this.keyPrefix}generic:${key}`;
await this.client.del(prefixedKey);
} catch (error) {
console.error("[LimitRate] Upstash delete error:", error);
throw error;
}
}
/**
* Get underlying Upstash client (for advanced use cases)
*/
getClient() {
return this.client;
}
};
// src/stores/index.ts
function createStore(config) {
switch (config.type) {
case "memory":
return new MemoryStore();
case "redis":
if (!config.url) {
throw new Error("Redis store requires url");
}
return new RedisStore({ client: config.url, redisOptions: config.options });
case "upstash":
if (!config.url || !config.token) {
throw new Error("Upstash store requires url and token");
}
return new UpstashStore({ url: config.url, token: config.token });
default:
throw new Error(`Unknown store type: ${config.type}`);
}
}
function createSharedMemoryStore(options) {
return new MemoryStore(options);
}
function createSharedRedisStore(options) {
if ("url" in options) {
return new RedisStore({ client: options.url, keyPrefix: options.keyPrefix, redisOptions: options.redisOptions });
}
return new RedisStore({ client: options.client, keyPrefix: options.keyPrefix });
}
function createSharedUpstashStore(options) {
return new UpstashStore(options);
}
// src/utils/events.ts
var EventEmitter = class {
handlers;
constructor() {
this.handlers = /* @__PURE__ */ new Set();
}
/**
* Register an event handler
*/
on(handler) {
this.handlers.add(handler);
}
/**
* Unregister an event handler
*/
off(handler) {
this.handlers.delete(handler);
}
/**
* Emit an event to all handlers
*/
async emit(event) {
const promises = [];
for (const handler of this.handlers) {
try {
const result = handler(event);
if (result instanceof Promise) {
promises.push(result);
}
} catch (error) {
console.error("[LimitRate] Event handler error:", error);
}
}
if (promises.length > 0) {
await Promise.allSettled(promises);
}
}
/**
* Get number of registered handlers
*/
getHandlerCount() {
return this.handlers.size;
}
/**
* Clear all handlers
*/
clear() {
this.handlers.clear();
}
};
// src/penalty/manager.ts
var PenaltyManager = class {
store;
constructor(store) {
this.store = store;
}
/**
* Get penalty state key for a user/endpoint
*/
getPenaltyKey(user, endpoint) {
return `penalty:${user}:${endpoint}`;
}
/**
* Get current penalty/reward multiplier for a user/endpoint
* Returns 1.0 if no penalty/reward is active
*/
async getMultiplier(user, endpoint) {
try {
const key = this.getPenaltyKey(user, endpoint);
const state = await this.store.get(key);
if (!state) {
return 1;
}
if (Date.now() > state.expiresAt) {
await this.store.delete(key);
return 1;
}
return state.multiplier;
} catch (error) {
console.warn("[LimitRate] Failed to get penalty multiplier:", error);
return 1;
}
}
/**
* Apply a penalty to a user/endpoint
*/
async applyPenalty(user, endpoint, config) {
try {
const key = this.getPenaltyKey(user, endpoint);
const state = {
multiplier: config.multiplier,
expiresAt: Date.now() + config.duration * 1e3,
reason: "violation"
};
await this.store.set(key, state, config.duration);
} catch (error) {
console.error("[LimitRate] Failed to apply penalty:", error);
}
}
/**
* Apply a reward to a user/endpoint
*/
async applyReward(user, endpoint, config) {
try {
const key = this.getPenaltyKey(user, endpoint);
const state = {
multiplier: config.multiplier,
expiresAt: Date.now() + config.duration * 1e3,
reason: "reward"
};
await this.store.set(key, state, config.duration);
} catch (error) {
console.error("[LimitRate] Failed to apply reward:", error);
}
}
/**
* Check if a reward should be granted based on usage
*/
shouldGrantReward(currentUsage, limit, rewardConfig) {
const usagePercent = currentUsage / limit * 100;
switch (rewardConfig.trigger) {
case "below_10_percent":
return usagePercent < 10;
case "below_25_percent":
return usagePercent < 25;
case "below_50_percent":
return usagePercent < 50;
default:
return false;
}
}
/**
* Clear penalty/reward for a user/endpoint
*/
async clear(user, endpoint) {
try {
const key = this.getPenaltyKey(user, endpoint);
await this.store.delete(key);
} catch (error) {
console.error("[LimitRate] Failed to clear penalty:", error);
}
}
};
// src/engine.ts
var PolicyEngine = class {
store;
policies;
events;
penaltyManager;
constructor(store, policies) {
this.store = store;
this.policies = policies;
this.events = new EventEmitter();
this.penaltyManager = new PenaltyManager(store);
}
/**
* Register event handler
*/
onEvent(handler) {
this.events.on(handler);
}
/**
* Check if request should be allowed
*/
async check(context) {
const policy = context.policyOverride || this.resolvePolicy(context.plan, context.endpoint);
if (!policy) {
await this.emitEvent({
timestamp: Date.now(),
user: context.user,
plan: context.plan,
endpoint: context.endpoint,
type: "allowed"
});
return {
allowed: true,
action: "allow",
details: { used: 0, limit: Infinity, remaining: Infinity, resetInSeconds: 0 }
};
}
let finalDetails = { used: 0, limit: 0, remaining: 0, resetInSeconds: 0 };
if (policy.rate) {
const rateResult = await this.checkRate(context, policy);
if (!rateResult.allowed || rateResult.action !== "allow") {
return rateResult;
}
finalDetails = rateResult.details;
}
if (policy.rate && context.tokens !== void 0 && context.tokens > 0) {
const tokenResult = await this.checkTokens(context, policy);
if (!tokenResult.allowed || tokenResult.action !== "allow") {
return tokenResult;
}
}
if (policy.cost) {
const costResult = await this.checkCost(context, policy);
if (!costResult.allowed || costResult.action !== "allow") {
return costResult;
}
}
await this.emitEvent({
timestamp: Date.now(),
user: context.user,
plan: context.plan,
endpoint: context.endpoint,
type: "allowed"
});
return {
allowed: true,
action: "allow",
details: finalDetails
};
}
/**
* Check rate limit
*/
async checkRate(context, policy) {
if (!policy.rate) {
return { allowed: true, action: "allow", details: { used: 0, limit: 0, remaining: 0, resetInSeconds: 0 } };
}
let { maxPerSecond, maxPerMinute, maxPerHour, maxPerDay, burst, actionOnExceed, slowdownMs } = policy.rate;
if (context.userOverride) {
const override = context.userOverride;
const isValidLimit = (val) => {
return typeof val === "number" && val > 0 && !isNaN(val) && isFinite(val);
};
const endpointOverride = override.endpoints?.[context.endpoint];
if (endpointOverride) {
if (isValidLimit(endpointOverride.maxPerSecond)) maxPerSecond = endpointOverride.maxPerSecond;
if (isValidLimit(endpointOverride.maxPerMinute)) maxPerMinute = endpointOverride.maxPerMinute;
if (isValidLimit(endpointOverride.maxPerHour)) maxPerHour = endpointOverride.maxPerHour;
if (isValidLimit(endpointOverride.maxPerDay)) maxPerDay = endpointOverride.maxPerDay;
} else {
if (isValidLimit(override.maxPerSecond)) maxPerSecond = override.maxPerSecond;
if (isValidLimit(override.maxPerMinute)) maxPerMinute = override.maxPerMinute;
if (isValidLimit(override.maxPerHour)) maxPerHour = override.maxPerHour;
if (isValidLimit(override.maxPerDay)) maxPerDay = override.maxPerDay;
}
}
let limit;
let windowSeconds;
let windowLabel;
if (maxPerSecond !== void 0) {
limit = maxPerSecond;
windowSeconds = 1;
windowLabel = "1s";
} else if (maxPerMinute !== void 0) {
limit = maxPerMinute;
windowSeconds = 60;
windowLabel = "1m";
} else if (maxPerHour !== void 0) {
limit = maxPerHour;
windowSeconds = 3600;
windowLabel = "1h";
} else if (maxPerDay !== void 0) {
limit = maxPerDay;
windowSeconds = 86400;
windowLabel = "1d";
} else {
throw new Error("No time window specified in rate rule");
}
const originalLimit = limit;
if (policy.penalty?.enabled) {
const multiplier = await this.penaltyManager.getMultiplier(context.user, context.endpoint);
limit = Math.floor(originalLimit * multiplier);
if (limit < 1) limit = 1;
}
const rateKey = `${context.user}:${context.endpoint}`;
const result = await this.store.checkRate(rateKey, limit, windowSeconds, burst);
if (result.allowed) {
if (policy.penalty?.enabled && policy.penalty.rewards) {
const shouldReward = this.penaltyManager.shouldGrantReward(
result.current,
originalLimit,
policy.penalty.rewards
);
if (shouldReward) {
await this.penaltyManager.applyReward(
context.user,
context.endpoint,
policy.penalty.rewards
);
}
}
return {
allowed: true,
action: "allow",
details: {
used: result.current,
limit: result.limit,
remaining: result.remaining,
resetInSeconds: result.resetInSeconds,
burstTokens: result.burstTokens
}
};
}
if (policy.penalty?.enabled && policy.penalty.onViolation) {
await this.penaltyManager.applyPenalty(
context.user,
context.endpoint,
policy.penalty.onViolation
);
}
await this.emitEvent({
timestamp: Date.now(),
user: context.user,
plan: context.plan,
endpoint: context.endpoint,
type: "rate_exceeded",
window: windowLabel,
value: result.current,
threshold: limit
});
if (actionOnExceed === "block") {
return {
allowed: false,
action: "block",
reason: "rate_exceeded",
retryAfterSeconds: result.resetInSeconds,
details: {
used: result.current,
limit: result.limit,
remaining: 0,
resetInSeconds: result.resetInSeconds,
burstTokens: result.burstTokens
}
};
}
if (actionOnExceed === "slowdown") {
await this.emitEvent({
timestamp: Date.now(),
user: context.user,
plan: context.plan,
endpoint: context.endpoint,
type: "slowdown_applied",
value: slowdownMs
});
return {
allowed: true,
action: "slowdown",
slowdownMs,
details: {
used: result.current,
limit: result.limit,
remaining: 0,
resetInSeconds: result.resetInSeconds,
burstTokens: result.burstTokens
}
};
}
if (actionOnExceed === "allow-and-log") {
return {
allowed: true,
action: "allow-and-log",
details: {
used: result.current,
limit: result.limit,
remaining: 0,
resetInSeconds: result.resetInSeconds,
burstTokens: result.burstTokens
}
};
}
return {
allowed: true,
action: "allow",
details: {
used: result.current,
limit: result.limit,
remaining: 0,
resetInSeconds: result.resetInSeconds,
burstTokens: result.burstTokens
}
};
}
/**
* Check cost limit
*/
async checkCost(context, policy) {
if (!policy.cost) {
return { allowed: true, action: "allow", details: { used: 0, limit: 0, remaining: 0, resetInSeconds: 0 } };
}
const { estimateCost, hourlyCap, dailyCap, actionOnExceed } = policy.cost;
const cost = estimateCost(context.costContext);
const cap = dailyCap ?? hourlyCap;
const windowSeconds = dailyCap ? 86400 : 3600;
const costKey = `${context.user}:${context.endpoint}:cost`;
const result = await this.store.incrementCost(costKey, cost, windowSeconds, cap);
if (result.allowed) {
return {
allowed: true,
action: "allow",
details: {
used: result.current,
limit: cap,
remaining: result.remaining,
resetInSeconds: result.resetInSeconds
}
};
}
await this.emitEvent({
timestamp: Date.now(),
user: context.user,
plan: context.plan,
endpoint: context.endpoint,
type: "cost_exceeded",
window: dailyCap ? "1d" : "1h",
value: result.current + cost,
threshold: cap
});
if (actionOnExceed === "block") {
return {
allowed: false,
action: "block",
reason: "cost_exceeded",
retryAfterSeconds: result.resetInSeconds,
details: {
used: result.current,
limit: cap,
remaining: 0,
resetInSeconds: result.resetInSeconds
}
};
}
if (actionOnExceed === "slowdown") {
return {
allowed: false,
action: "block",
reason: "cost_exceeded",
retryAfterSeconds: result.resetInSeconds,
details: {
used: result.current,
limit: cap,
remaining: 0,
resetInSeconds: result.resetInSeconds
}
};
}
if (actionOnExceed === "allow-and-log") {
return {
allowed: true,
action: "allow-and-log",
details: {
used: result.current,
limit: cap,
remaining: 0,
resetInSeconds: result.resetInSeconds
}
};
}
return {
allowed: true,
action: "allow",
details: {
used: result.current,
limit: cap,
remaining: 0,
resetInSeconds: result.resetInSeconds
}
};
}
/**
* Check token limits (v1.4.0 - AI feature)
*/
async checkTokens(context, policy) {
if (!policy.rate || !context.tokens) {
return { allowed: true, action: "allow", details: { used: 0, limit: 0, remaining: 0, resetInSeconds: 0 } };
}
const { maxTokensPerMinute, maxTokensPerHour, maxTokensPerDay, actionOnExceed, slowdownMs } = policy.rate;
if (!maxTokensPerMinute && !maxTokensPerHour && !maxTokensPerDay) {
return { allowed: true, action: "allow", details: { used: 0, limit: 0, remaining: 0, resetInSeconds: 0 } };
}
const checks = [];
if (maxTokensPerMinute) {
checks.push({ limit: maxTokensPerMinute, windowSeconds: 60, windowLabel: "1m" });
}
if (maxTokensPerHour) {
checks.push({ limit: maxTokensPerHour, windowSeconds: 3600, windowLabel: "1h" });
}
if (maxTokensPerDay) {
checks.push({ limit: maxTokensPerDay, windowSeconds: 86400, windowLabel: "1d" });
}
for (const check of checks) {
const tokenKey = `${context.user}:${context.endpoint}:tokens:${check.windowLabel}`;
const result = await this.store.incrementTokens(
tokenKey,
context.tokens,
check.windowSeconds,
check.limit
);
if (!result.allowed) {
await this.emitEvent({
timestamp: Date.now(),
user: context.user,
plan: context.plan,
endpoint: context.endpoint,
type: "token_limit_exceeded",
window: check.windowLabel,
value: result.current,
threshold: check.limit,
tokens: context.tokens
});
if (actionOnExceed === "block") {
return {
allowed: false,
action: "block",
reason: "token_limit_exceeded",
retryAfterSeconds: result.resetInSeconds,
details: {
used: result.current,
limit: result.limit,
remaining: 0,
resetInSeconds: result.resetInSeconds
}
};
}
if (actionOnExceed === "slowdown") {
await this.emitEvent({
timestamp: Date.now(),
user: context.user,
plan: context.plan,
endpoint: context.endpoint,
type: "slowdown_applied",
value: slowdownMs
});
return {
allowed: true,
action: "slowdown",
slowdownMs,
details: {
used: result.current,
limit: result.limit,
remaining: 0,
resetInSeconds: result.resetInSeconds
}
};
}
if (actionOnExceed === "allow-and-log") {
return {
allowed: true,
action: "allow-and-log",
details: {
used: result.current,
limit: result.limit,
remaining: 0,
resetInSeconds: result.resetInSeconds
}
};
}
return {
allowed: true,
action: "allow",
details: {
used: result.current,
limit: result.limit,
remaining: 0,
resetInSeconds: result.resetInSeconds
}
};
}
}
await this.emitEvent({
timestamp: Date.now(),
user: context.user,
plan: context.plan,
endpoint: context.endpoint,
type: "token_usage_tracked",
tokens: context.tokens
});
return {
allowed: true,
action: "allow",
details: { used: 0, limit: 0, remaining: 0, resetInSeconds: 0 }
};
}
/**
* Resolve policy for plan and endpoint
*/
resolvePolicy(plan, endpoint) {
const planConfig = this.policies[plan];
if (!planConfig) {
return null;
}
if (planConfig.endpoints?.[endpoint]) {
return planConfig.endpoints[endpoint];
}
return planConfig.defaults ?? null;
}
/**
* Emit event
*/
async emitEvent(event) {