@limitrate/core
Version:
Core rate limiting and cost control engine for LimitRate
1,726 lines (1,674 loc) • 72 kB
JavaScript
// 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
import Redis from "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 Redis) {
this.client = options.client;
} else if (typeof options.client === "string") {
this.client = new Redis(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
import { Redis as Redis2 } from "@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 Redis2({
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) {
await this.events.emit(event);
}
/**
* Get event emitter (for external handlers)
*/
getEventEmitter() {
return this.events;
}
};
// src/validation.ts
var ValidationError = class extends Error {
constructor(message) {
super(`[LimitRate] Invalid configuration: ${message}`);
this.name = "ValidationError";
}
};
function isPositiveNumber(value) {
return typeof value === "number" && value > 0 && !isNaN(value);
}
function isFunction(value) {
return typeof value === "function";
}
function validateRateRule(rule, path) {
if (rule.maxPerSecond !== void 0 && !isPositiveNumber(rule.maxPerSecond)) {
throw new ValidationError(`${path}.maxPerSecond must be a positive number, got: ${rule.maxPerSecond}`);
}
if (rule.maxPerMinute !== void 0 && !isPositiveNumber(rule.maxPerMinute)) {
throw new ValidationError(`${path}.maxPerMinute must be a positive number, got: ${rule.maxPerMinute}`);
}
if (rule.maxPerHour !== void 0 && !isPositiveNumber(rule.maxPerHour)) {
throw new ValidationError(`${path}.maxPerHour must be a positive number, got: ${rule.maxPerHour}`);
}
if (rule.maxPerDay !== void 0 && !isPositiveNumber(rule.maxPerDay)) {
throw new ValidationError(`${path}.maxPerDay must be a positive number, got: ${rule.maxPerDay}`);
}
const timeWindows = [
rule.maxPerSecond,
rule.maxPerMinute,
rule.maxPerHour,
rule.maxPerDay
].filter((w) => w !== void 0);
if (timeWindows.length === 0) {
throw new ValidationError(
`${path} must specify exactly one time window (maxPerSecond, maxPerMinute, maxPerHour, or maxPerDay)`
);
}
if (timeWindows.length > 1) {
throw new ValidationError(
`${path} can only specify one time window, but found ${timeWindows.length}. Use separate endpoint policies for multiple limits.`
);
}
if (rule.burst !== void 0 && !isPositiveNumber(rule.burst)) {
throw new ValidationError(`${path}.burst must be a positive number, got: ${rule.burst}`);
}
if (rule.slowdownMs !== void 0) {
if (!isPositiveNumber(rule.slowdownMs)) {
throw new ValidationError(`${path}.slowdownMs must be a positive number, got: ${rule.slowdownMs}`);
}
if (rule.slowdownMs > 6e4) {
throw new ValidationError(`${path}.slowdownMs must be <= 60000 (60 seconds), got: ${rule.slowdownMs}`);
}
}
const validActions = ["allow", "block", "slowdown", "allow-and-log"];
if (!validActions.includes(rule.actionOnExceed)) {
throw new ValidationError(
`${path}.actionOnExceed must be one of: ${validActions.join(", ")}, got: ${rule.actionOnExceed}`
);
}
if (rule.actionOnExceed === "slowdown" && !rule.slowdownMs) {
throw new ValidationError(`${path}.slowdownMs is required when actionOnExceed is 'slowdown'`);
}
}
function validateCostRule(rule, path) {
if (!isFunction(rule.estimateCost)) {
throw new ValidationError(`${path}.estimateCost must be a function`);
}
if (rule.hourlyCap !== void 0 && !isPositiveNumber(rule.hourlyCap)) {
throw new ValidationError(`${path}.hourlyCap must be a positive number, got: ${rule.hourlyCap}`);
}
if (rule.dailyCap !== void 0 && !isPositiveNumber(rule.dailyCap)) {
throw new ValidationError(`${path}.dailyCap must be a positive number, got: ${rule.dailyCap}`);
}
const validActions = ["allow", "block", "slowdown", "allow-and-log"];
if (!validActions.includes(rule.actionOnExceed)) {
throw new ValidationError(
`${path}.actionOnExceed must be one of: ${validActions.join(", ")}, got: ${rule.actionOnExceed}`
);
}
}
function validatePolicyConfig(config) {
if (!config || typeof config !== "object") {
throw new ValidationError("PolicyConfig must be an object");
}
for (const [plan, planConfig] of Object.entries(config)) {
if (!planConfig || typeof planConfig !== "object") {
throw new ValidationError(`Plan '${plan}' config must be an object`);
}
const hasEndpoints = planConfig.endpoints && typeof planConfig.endpoints === "object";
const hasDefaults = planConfig.defaults && typeof planConfig.defaults === "object";
if (!hasEndpoints && !hasDefaults) {
throw new ValidationError(`Plan '${plan}' must have 'endpoints' or 'defaults' object`);
}
if (hasEndpoints) {
for (const [endpoint, policy] of Object.entries(planConfig.endpoints)) {
const basePath = `policies.${plan}.endpoints['${endpoint}']`;
if (policy.rate) {
validateRateRule(policy.rate, `${basePath}.rate`);
}
if (policy.cost) {
validateCostRule(policy.cost, `${basePath}.cost`);
}
if (!policy.rate && !policy.cost && !policy.concurrency) {
throw new ValidationError(`${basePath} must have at least one of: rate, cost, concurrency`);
}
}
}
if (planConfig.defaults) {
const basePath = `policies.${plan}.defaults`;
if (planConfig.defaults.rate) {
validateRateRule(planConfig.defaults.rate, `${basePath}.rate`);
}
if (planConfig.defaults.cost) {
validateCostRule(planConfig.defaults.cost, `${basePath}.cost`);
}
}
}
}
function validateStoreConfig(config) {
if (!config || typeof config !== "object") {
throw new ValidationError("StoreConfig must be an object");
}
const validTypes = ["memory", "redis", "upstash"];
if (!validTypes.includes(config.type)) {
throw new ValidationError(`store.type must be one of: ${validTypes.join(", ")}, got: ${config.type}`);
}
if (config.type === "redis" && !config.url) {
throw new ValidationError('store.url is required when type is "redis"');
}
if (config.type === "upstash") {
if (!config.url) {
throw new ValidationError('store.url is required when type is "upstash"');
}
if (!config.token) {
throw new ValidationError('store.token is required when type is "upstash"');
}
}
}
function validateIPAddress(ip) {
const ipv4Regex = /^(\d{1,3}\.){3}\d{1,3}$/;
if (ipv4Regex.test(ip)) {
const parts = ip.split(".");
return parts.every((part) => {
const num = parseInt(part, 10);
return num >= 0 && num <= 255;
});
}
const ipv4MappedRegex = /^::ffff:(\d{1,3}\.){3}\d{1,3}$/i;
if (ipv4MappedRegex.test(ip)) {
const ipv4Part = ip.substring(7);
const parts = ipv4Part.split(".");
return parts.every((part) => {
const num = parseInt(part, 10);