Redis ratelimits

This commit is contained in:
2026-07-21 15:55:32 -05:00
parent 571f3e9fd8
commit d59673e0c5
8 changed files with 290 additions and 34 deletions
+5
View File
@@ -8,6 +8,11 @@ DATABASE_URL=postgres://user:password@localhost:5432/opencode_proxy
# Send one in Authorization: Bearer <key> or x-api-key. # Send one in Authorization: Bearer <key> or x-api-key.
# UNLIMITED_API_KEYS=["personal-unlimited-key"] # UNLIMITED_API_KEYS=["personal-unlimited-key"]
# Optional: persist shared rate limits across restarts and proxy instances.
# REDIS_URL=redis://localhost:6379
# Optional Redis key namespace, useful when sharing a Redis database.
# REDIS_RATE_LIMIT_PREFIX=opencode-proxy:rate-limit
# Optional: override the built-in comma-separated model allowlist for ordinary clients. # Optional: override the built-in comma-separated model allowlist for ordinary clients.
# Unlimited API keys bypass this restriction. Set this to an empty value to allow every upstream model. # Unlimited API keys bypass this restriction. Set this to an empty value to allow every upstream model.
# ALLOWED_MODELS=kimi-k3,gpt-4.1 # ALLOWED_MODELS=kimi-k3,gpt-4.1
+6
View File
@@ -37,6 +37,12 @@ The global context limit defaults to 256k tokens and can be changed with `MAX_CO
Inference requests are limited per client IP to 2 requests per second, 100 requests per five hours, and $15 of reported upstream cost per five hours. Railway's `X-Real-IP` header is used to identify clients. Inference requests are limited per client IP to 2 requests per second, 100 requests per five hours, and $15 of reported upstream cost per five hours. Railway's `X-Real-IP` header is used to identify clients.
Set `REDIS_URL` to persist and share these limits across process restarts and multiple proxy instances. The Redis backend stores only the active five-hour window and uses atomic operations, so concurrent instances enforce one shared limit. Without `REDIS_URL`, limits remain in memory as before.
```sh
REDIS_URL=redis://localhost:6379
```
To exempt trusted clients from those proxy limits, set `UNLIMITED_API_KEYS` to a JSON array and have the client send a configured key using `Authorization: Bearer <key>` or `x-api-key`. These keys only bypass this proxy's rate and spend limits; they do not bypass upstream OpenCode Go limits. To exempt trusted clients from those proxy limits, set `UNLIMITED_API_KEYS` to a JSON array and have the client send a configured key using `Authorization: Bearer <key>` or `x-api-key`. These keys only bypass this proxy's rate and spend limits; they do not bypass upstream OpenCode Go limits.
```sh ```sh
+102 -1
View File
@@ -10,7 +10,8 @@
"dependencies": { "dependencies": {
"dotenv": "^16.6.1", "dotenv": "^16.6.1",
"express": "^5.1.0", "express": "^5.1.0",
"pg": "^8.22.0" "pg": "^8.22.0",
"redis": "^4.7.1"
}, },
"devDependencies": { "devDependencies": {
"supertest": "^7.1.1", "supertest": "^7.1.1",
@@ -492,6 +493,65 @@
"@noble/hashes": "^1.1.5" "@noble/hashes": "^1.1.5"
} }
}, },
"node_modules/@redis/bloom": {
"version": "1.2.0",
"resolved": "https://registry.npmjs.org/@redis/bloom/-/bloom-1.2.0.tgz",
"integrity": "sha512-HG2DFjYKbpNmVXsa0keLHp/3leGJz1mjh09f2RLGGLQZzSHpkmZWuwJbAvo3QcRY8p80m5+ZdXZdYOSBLlp7Cg==",
"license": "MIT",
"peerDependencies": {
"@redis/client": "^1.0.0"
}
},
"node_modules/@redis/client": {
"version": "1.6.1",
"resolved": "https://registry.npmjs.org/@redis/client/-/client-1.6.1.tgz",
"integrity": "sha512-/KCsg3xSlR+nCK8/8ZYSknYxvXHwubJrU82F3Lm1Fp6789VQ0/3RJKfsmRXjqfaTA++23CvC3hqmqe/2GEt6Kw==",
"license": "MIT",
"dependencies": {
"cluster-key-slot": "1.1.2",
"generic-pool": "3.9.0",
"yallist": "4.0.0"
},
"engines": {
"node": ">=14"
}
},
"node_modules/@redis/graph": {
"version": "1.1.1",
"resolved": "https://registry.npmjs.org/@redis/graph/-/graph-1.1.1.tgz",
"integrity": "sha512-FEMTcTHZozZciLRl6GiiIB4zGm5z5F3F6a6FZCyrfxdKOhFlGkiAqlexWMBzCi4DcRoyiOsuLfW+cjlGWyExOw==",
"license": "MIT",
"peerDependencies": {
"@redis/client": "^1.0.0"
}
},
"node_modules/@redis/json": {
"version": "1.0.7",
"resolved": "https://registry.npmjs.org/@redis/json/-/json-1.0.7.tgz",
"integrity": "sha512-6UyXfjVaTBTJtKNG4/9Z8PSpKE6XgSyEb8iwaqDcy+uKrd/DGYHTWkUdnQDyzm727V7p21WUMhsqz5oy65kPcQ==",
"license": "MIT",
"peerDependencies": {
"@redis/client": "^1.0.0"
}
},
"node_modules/@redis/search": {
"version": "1.2.0",
"resolved": "https://registry.npmjs.org/@redis/search/-/search-1.2.0.tgz",
"integrity": "sha512-tYoDBbtqOVigEDMAcTGsRlMycIIjwMCgD8eR2t0NANeQmgK/lvxNAvYyb6bZDD4frHRhIHkJu2TBRvB0ERkOmw==",
"license": "MIT",
"peerDependencies": {
"@redis/client": "^1.0.0"
}
},
"node_modules/@redis/time-series": {
"version": "1.1.0",
"resolved": "https://registry.npmjs.org/@redis/time-series/-/time-series-1.1.0.tgz",
"integrity": "sha512-c1Q99M5ljsIuc4YdaCwfUEXsofakb9c8+Zse2qxTadu8TalLXuAESzLvFAvNVbkmSlvlzIQOLpBCmWI9wTOt+g==",
"license": "MIT",
"peerDependencies": {
"@redis/client": "^1.0.0"
}
},
"node_modules/@rollup/rollup-android-arm-eabi": { "node_modules/@rollup/rollup-android-arm-eabi": {
"version": "4.62.2", "version": "4.62.2",
"resolved": "https://registry.npmjs.org/@rollup/rollup-android-arm-eabi/-/rollup-android-arm-eabi-4.62.2.tgz", "resolved": "https://registry.npmjs.org/@rollup/rollup-android-arm-eabi/-/rollup-android-arm-eabi-4.62.2.tgz",
@@ -1131,6 +1191,15 @@
"node": ">= 16" "node": ">= 16"
} }
}, },
"node_modules/cluster-key-slot": {
"version": "1.1.2",
"resolved": "https://registry.npmjs.org/cluster-key-slot/-/cluster-key-slot-1.1.2.tgz",
"integrity": "sha512-RMr0FhtfXemyinomL4hrWcYJxmX6deFdCxpJzhDttxgO1+bcCnkk+9drydLVDmAMG7NE6aN/fl4F7ucU/90gAA==",
"license": "Apache-2.0",
"engines": {
"node": ">=0.10.0"
}
},
"node_modules/combined-stream": { "node_modules/combined-stream": {
"version": "1.0.8", "version": "1.0.8",
"resolved": "https://registry.npmjs.org/combined-stream/-/combined-stream-1.0.8.tgz", "resolved": "https://registry.npmjs.org/combined-stream/-/combined-stream-1.0.8.tgz",
@@ -1618,6 +1687,15 @@
"url": "https://github.com/sponsors/ljharb" "url": "https://github.com/sponsors/ljharb"
} }
}, },
"node_modules/generic-pool": {
"version": "3.9.0",
"resolved": "https://registry.npmjs.org/generic-pool/-/generic-pool-3.9.0.tgz",
"integrity": "sha512-hymDOu5B53XvN4QT9dBmZxPX4CWhBPPLguTZ9MMFeFa/Kg0xWVfylOVNlJji/E7yTZWFd/q9GO5TxDLq156D7g==",
"license": "MIT",
"engines": {
"node": ">= 4"
}
},
"node_modules/get-intrinsic": { "node_modules/get-intrinsic": {
"version": "1.3.0", "version": "1.3.0",
"resolved": "https://registry.npmjs.org/get-intrinsic/-/get-intrinsic-1.3.0.tgz", "resolved": "https://registry.npmjs.org/get-intrinsic/-/get-intrinsic-1.3.0.tgz",
@@ -2203,6 +2281,23 @@
"node": ">= 0.10" "node": ">= 0.10"
} }
}, },
"node_modules/redis": {
"version": "4.7.1",
"resolved": "https://registry.npmjs.org/redis/-/redis-4.7.1.tgz",
"integrity": "sha512-S1bJDnqLftzHXHP8JsT5II/CtHWQrASX5K96REjWjlmWKrviSOLWmM7QnRLstAWsu1VBBV1ffV6DzCvxNP0UJQ==",
"license": "MIT",
"workspaces": [
"./packages/*"
],
"dependencies": {
"@redis/bloom": "1.2.0",
"@redis/client": "1.6.1",
"@redis/graph": "1.1.1",
"@redis/json": "1.0.7",
"@redis/search": "1.2.0",
"@redis/time-series": "1.1.0"
}
},
"node_modules/rollup": { "node_modules/rollup": {
"version": "4.62.2", "version": "4.62.2",
"resolved": "https://registry.npmjs.org/rollup/-/rollup-4.62.2.tgz", "resolved": "https://registry.npmjs.org/rollup/-/rollup-4.62.2.tgz",
@@ -2812,6 +2907,12 @@
"engines": { "engines": {
"node": ">=0.4" "node": ">=0.4"
} }
},
"node_modules/yallist": {
"version": "4.0.0",
"resolved": "https://registry.npmjs.org/yallist/-/yallist-4.0.0.tgz",
"integrity": "sha512-3wdGidZyq5PB084XLES5TpOSRA3wjXAlIWMhum2kRcv/41Sn2emQ0dycQW4uZXLejwKvg6EsvbdlVL+FYEct7A==",
"license": "ISC"
} }
} }
} }
+2 -1
View File
@@ -17,7 +17,8 @@
"dependencies": { "dependencies": {
"dotenv": "^16.6.1", "dotenv": "^16.6.1",
"express": "^5.1.0", "express": "^5.1.0",
"pg": "^8.22.0" "pg": "^8.22.0",
"redis": "^4.7.1"
}, },
"devDependencies": { "devDependencies": {
"supertest": "^7.1.1", "supertest": "^7.1.1",
+24 -25
View File
@@ -257,7 +257,7 @@ export function createProxyApp({ keys, fetchImpl = globalThis.fetch, upstreamBas
const unlimited = hasUnlimitedKey(request, unlimitedKeys); const unlimited = hasUnlimitedKey(request, unlimitedKeys);
if (!unlimited) validateAllowedModel(request.body?.model, allowedModels); if (!unlimited) validateAllowedModel(request.body?.model, allowedModels);
if (!unlimited) { if (!unlimited) {
const limit = rateLimiter.check(ip); const limit = await rateLimiter.check(ip);
if (!limit.allowed) return rateLimitError(response, false, limit); if (!limit.allowed) return rateLimitError(response, false, limit);
} }
const upstream = await fetchWithKeyRetries(`${upstreamBaseUrl}/chat/completions`, { const upstream = await fetchWithKeyRetries(`${upstreamBaseUrl}/chat/completions`, {
@@ -285,23 +285,24 @@ export function createProxyApp({ keys, fetchImpl = globalThis.fetch, upstreamBas
for (const eventText of events) { for (const eventText of events) {
const data = eventText.split('\n').find((line) => line.startsWith('data:'))?.slice(5).trim(); const data = eventText.split('\n').find((line) => line.startsWith('data:'))?.slice(5).trim();
if (!recordedCost && data && data !== '[DONE]') { if (!recordedCost && data && data !== '[DONE]') {
try { let cost = null;
const cost = numericCost(JSON.parse(data)); try { cost = numericCost(JSON.parse(data)); } catch { /* ignore malformed stream events */ }
if (cost !== null) { if (!unlimited) rateLimiter.recordCost(ip, cost); recordedCost = true; } if (cost !== null) { if (!unlimited) await rateLimiter.recordCost(ip, cost); recordedCost = true; }
} catch { /* ignore malformed stream events */ }
} }
} }
} }
if (!recordedCost) { if (!recordedCost) {
const data = streamBuffer.split('\n').find((line) => line.startsWith('data:'))?.slice(5).trim(); const data = streamBuffer.split('\n').find((line) => line.startsWith('data:'))?.slice(5).trim();
try { const cost = numericCost(JSON.parse(data)); if (cost !== null && !unlimited) rateLimiter.recordCost(ip, cost); } catch { /* ignore */ } let cost = null;
try { cost = numericCost(JSON.parse(data)); } catch { /* ignore malformed stream event */ }
if (cost !== null && !unlimited) await rateLimiter.recordCost(ip, cost);
} }
response.end(); response.end();
return; return;
} }
const body = await upstream.text(); const body = await upstream.text();
try { if (!unlimited) rateLimiter.recordCost(ip, numericCost(JSON.parse(body))); } catch { /* non-JSON upstream response */ } try { if (!unlimited) await rateLimiter.recordCost(ip, numericCost(JSON.parse(body))); } catch { /* non-JSON upstream response */ }
response.status(upstream.status).send(body); response.status(upstream.status).send(body);
} catch (error) { } catch (error) {
if (!response.headersSent) { if (!response.headersSent) {
@@ -326,7 +327,7 @@ export function createProxyApp({ keys, fetchImpl = globalThis.fetch, upstreamBas
const unlimited = hasUnlimitedKey(request, unlimitedKeys); const unlimited = hasUnlimitedKey(request, unlimitedKeys);
if (!unlimited) validateAllowedModel(translated.model, allowedModels); if (!unlimited) validateAllowedModel(translated.model, allowedModels);
if (!unlimited) { if (!unlimited) {
const limit = rateLimiter.check(ip); const limit = await rateLimiter.check(ip);
if (!limit.allowed) return rateLimitError(response, false, limit); if (!limit.allowed) return rateLimitError(response, false, limit);
} }
const upstream = await fetchWithKeyRetries(`${upstreamBaseUrl}/chat/completions`, { const upstream = await fetchWithKeyRetries(`${upstreamBaseUrl}/chat/completions`, {
@@ -340,7 +341,7 @@ export function createProxyApp({ keys, fetchImpl = globalThis.fetch, upstreamBas
if (!translated.stream) { if (!translated.stream) {
const body = await upstream.text(); const body = await upstream.text();
const payload = JSON.parse(body); const payload = JSON.parse(body);
if (!unlimited) rateLimiter.recordCost(ip, numericCost(payload)); if (!unlimited) await rateLimiter.recordCost(ip, numericCost(payload));
return response.type('application/json').send(JSON.stringify(translateResponsesResponse(payload, translated.model, responseId))); return response.type('application/json').send(JSON.stringify(translateResponsesResponse(payload, translated.model, responseId)));
} }
if (!upstream.body) return response.status(502).json({ error: { type: 'upstream_error', message: 'Upstream returned no streaming body' } }); if (!upstream.body) return response.status(502).json({ error: { type: 'upstream_error', message: 'Upstream returned no streaming body' } });
@@ -351,17 +352,16 @@ export function createProxyApp({ keys, fetchImpl = globalThis.fetch, upstreamBas
let buffer = ''; let buffer = '';
let recordedCost = false; let recordedCost = false;
const state = { id: responseId, model: translated.model }; const state = { id: responseId, model: translated.model };
const writeEvents = (text) => { const writeEvents = async (text) => {
for (const eventText of text.split(/\n\n/)) { for (const eventText of text.split(/\n\n/)) {
if (!eventText.trim()) continue; if (!eventText.trim()) continue;
const data = eventText.split('\n').find((line) => line.startsWith('data:'))?.slice(5).trim(); const data = eventText.split('\n').find((line) => line.startsWith('data:'))?.slice(5).trim();
if (!data || data === '[DONE]') continue; if (!data || data === '[DONE]') continue;
try { let payload;
const payload = JSON.parse(data); try { payload = JSON.parse(data); } catch { continue; /* ignore malformed upstream events */ }
const cost = numericCost(payload); const cost = numericCost(payload);
if (!recordedCost && cost !== null) { if (!unlimited) rateLimiter.recordCost(ip, cost); recordedCost = true; } if (!recordedCost && cost !== null) { if (!unlimited) await rateLimiter.recordCost(ip, cost); recordedCost = true; }
response.write(translateResponsesChunk(payload, state)); response.write(translateResponsesChunk(payload, state));
} catch { /* ignore malformed upstream events */ }
} }
}; };
while (true) { while (true) {
@@ -369,10 +369,10 @@ export function createProxyApp({ keys, fetchImpl = globalThis.fetch, upstreamBas
buffer += decoder.decode(value || new Uint8Array(), { stream: !done }); buffer += decoder.decode(value || new Uint8Array(), { stream: !done });
const parts = buffer.split('\n\n'); const parts = buffer.split('\n\n');
buffer = parts.pop() || ''; buffer = parts.pop() || '';
writeEvents(parts.join('\n\n')); await writeEvents(parts.join('\n\n'));
if (done) break; if (done) break;
} }
writeEvents(buffer); await writeEvents(buffer);
response.end(); response.end();
} catch (error) { } catch (error) {
if (!response.headersSent) response.status(error.status || 502).json({ error: { message: error.status === 400 ? error.message : `Unable to reach OpenCode Go: ${error.message}`, type: error.status === 400 ? 'invalid_request_error' : 'upstream_error' } }); if (!response.headersSent) response.status(error.status || 502).json({ error: { message: error.status === 400 ? error.message : `Unable to reach OpenCode Go: ${error.message}`, type: error.status === 400 ? 'invalid_request_error' : 'upstream_error' } });
@@ -395,7 +395,7 @@ export function createProxyApp({ keys, fetchImpl = globalThis.fetch, upstreamBas
const unlimited = hasUnlimitedKey(request, unlimitedKeys); const unlimited = hasUnlimitedKey(request, unlimitedKeys);
if (!unlimited) validateAllowedModel(translated.model, allowedModels); if (!unlimited) validateAllowedModel(translated.model, allowedModels);
if (!unlimited) { if (!unlimited) {
const limit = rateLimiter.check(ip); const limit = await rateLimiter.check(ip);
if (!limit.allowed) return rateLimitError(response, true, limit); if (!limit.allowed) return rateLimitError(response, true, limit);
} }
const upstream = await fetchWithKeyRetries(`${upstreamBaseUrl}/chat/completions`, { const upstream = await fetchWithKeyRetries(`${upstreamBaseUrl}/chat/completions`, {
@@ -406,7 +406,7 @@ export function createProxyApp({ keys, fetchImpl = globalThis.fetch, upstreamBas
if (!upstream.ok) return forwardError(upstream, response); if (!upstream.ok) return forwardError(upstream, response);
if (!translated.stream) { if (!translated.stream) {
const payload = JSON.parse(await upstream.text()); const payload = JSON.parse(await upstream.text());
if (!unlimited) rateLimiter.recordCost(ip, numericCost(payload)); if (!unlimited) await rateLimiter.recordCost(ip, numericCost(payload));
return response.type('application/json').send(JSON.stringify(translateAnthropicResponse(payload, translated.model))); return response.type('application/json').send(JSON.stringify(translateAnthropicResponse(payload, translated.model)));
} }
if (!upstream.body) return response.status(502).json({ type: 'error', error: { type: 'upstream_error', message: 'Upstream returned no streaming body' } }); if (!upstream.body) return response.status(502).json({ type: 'error', error: { type: 'upstream_error', message: 'Upstream returned no streaming body' } });
@@ -417,17 +417,16 @@ export function createProxyApp({ keys, fetchImpl = globalThis.fetch, upstreamBas
let buffer = ''; let buffer = '';
const state = { id: `msg_${Date.now()}`, model: translated.model, nextBlockIndex: 0 }; const state = { id: `msg_${Date.now()}`, model: translated.model, nextBlockIndex: 0 };
let recordedCost = false; let recordedCost = false;
const writeEvents = (text) => { const writeEvents = async (text) => {
for (const eventText of text.split('\n\n')) { for (const eventText of text.split('\n\n')) {
if (!eventText.trim()) continue; if (!eventText.trim()) continue;
const data = eventText.split('\n').find((line) => line.startsWith('data:'))?.slice(5).trim(); const data = eventText.split('\n').find((line) => line.startsWith('data:'))?.slice(5).trim();
if (!data || data === '[DONE]') continue; if (!data || data === '[DONE]') continue;
try { let payload;
const payload = JSON.parse(data); try { payload = JSON.parse(data); } catch { continue; /* ignore malformed upstream event */ }
const cost = numericCost(payload); const cost = numericCost(payload);
if (!recordedCost && cost !== null) { if (!unlimited) rateLimiter.recordCost(ip, cost); recordedCost = true; } if (!recordedCost && cost !== null) { if (!unlimited) await rateLimiter.recordCost(ip, cost); recordedCost = true; }
response.write(translateOpenAIChunk(payload, state)); response.write(translateOpenAIChunk(payload, state));
} catch { /* ignore malformed upstream event */ }
} }
}; };
while (true) { while (true) {
@@ -435,10 +434,10 @@ export function createProxyApp({ keys, fetchImpl = globalThis.fetch, upstreamBas
buffer += decoder.decode(value || new Uint8Array(), { stream: !done }); buffer += decoder.decode(value || new Uint8Array(), { stream: !done });
const parts = buffer.split('\n\n'); const parts = buffer.split('\n\n');
buffer = parts.pop() || ''; buffer = parts.pop() || '';
writeEvents(parts.join('\n\n')); await writeEvents(parts.join('\n\n'));
if (done) break; if (done) break;
} }
writeEvents(buffer); await writeEvents(buffer);
response.end(); response.end();
} catch (error) { } catch (error) {
if (!response.headersSent) response.status(error.status || 502).json(error.status === 400 if (!response.headersSent) response.status(error.status || 502).json(error.status === 400
+111
View File
@@ -1,6 +1,66 @@
const SECOND = 1_000; const SECOND = 1_000;
const FIVE_HOURS = 5 * 60 * 60 * 1_000; const FIVE_HOURS = 5 * 60 * 60 * 1_000;
// The request check and insertion need to happen together. Keeping this in
// Lua means separate proxy processes cannot both accept a request based on the
// same stale count.
const REDIS_CHECK_SCRIPT = `
local now = tonumber(ARGV[1])
local secondStart = tonumber(ARGV[2])
local fiveHoursStart = tonumber(ARGV[3])
local perSecond = tonumber(ARGV[4])
local perFiveHours = tonumber(ARGV[5])
local spendLimit = tonumber(ARGV[6])
local window = tonumber(ARGV[7])
redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', fiveHoursStart)
redis.call('ZREMRANGEBYSCORE', KEYS[2], '-inf', fiveHoursStart)
local recentRequests = redis.call('ZCOUNT', KEYS[1], secondStart, '+inf')
if recentRequests >= perSecond then
local first = redis.call('ZRANGEBYSCORE', KEYS[1], secondStart, '+inf', 'WITHSCORES', 'LIMIT', 0, 1)
return { 0, 'rate', tonumber(first[2]) + 1000 - now }
end
local requestCount = redis.call('ZCARD', KEYS[1])
if requestCount >= perFiveHours then
local first = redis.call('ZRANGE', KEYS[1], 0, 0, 'WITHSCORES')
return { 0, 'rate', tonumber(first[2]) + window - now }
end
local spend = 0
local entries = redis.call('ZRANGE', KEYS[2], 0, -1, 'WITHSCORES')
for i = 1, #entries, 2 do
local amount = string.match(entries[i], '([^:]+)$')
spend = spend + (tonumber(amount) or 0)
end
if spend >= spendLimit then
local first = redis.call('ZRANGE', KEYS[2], 0, 0, 'WITHSCORES')
return { 0, 'spend', tonumber(first[2]) + window - now }
end
local sequence = redis.call('INCR', KEYS[3])
redis.call('ZADD', KEYS[1], now, tostring(now) .. ':' .. tostring(sequence))
redis.call('PEXPIRE', KEYS[1], window)
redis.call('PEXPIRE', KEYS[3], window)
if redis.call('EXISTS', KEYS[2]) == 1 then redis.call('PEXPIRE', KEYS[2], window) end
return { 1 }
`;
const REDIS_RECORD_COST_SCRIPT = `
local now = tonumber(ARGV[1])
local fiveHoursStart = tonumber(ARGV[2])
local amount = ARGV[3]
local window = tonumber(ARGV[4])
redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', fiveHoursStart)
local sequence = redis.call('INCR', KEYS[2])
redis.call('ZADD', KEYS[1], now, tostring(now) .. ':' .. tostring(sequence) .. ':' .. amount)
redis.call('PEXPIRE', KEYS[1], window)
redis.call('PEXPIRE', KEYS[2], window)
return 1
`;
export const RATE_LIMITS = { export const RATE_LIMITS = {
requestsPerSecond: 2, requestsPerSecond: 2,
requestsPerFiveHours: 100, requestsPerFiveHours: 100,
@@ -59,6 +119,57 @@ export function createRateLimiter({ now = () => Date.now() } = {}) {
}; };
} }
/**
* Creates a Redis-backed limiter. `client` is a connected node-redis client.
* Its sorted sets expire after the five-hour window, so Redis retains no rate
* limit data beyond the period in which it can affect a decision.
*/
export function createRedisRateLimiter({ client, now = () => Date.now(), keyPrefix = 'opencode-proxy:rate-limit' } = {}) {
if (!client?.eval) throw new Error('A connected Redis client is required for Redis rate limiting');
function keysFor(ip) {
// The hash tag keeps the three keys in the same slot when REDIS_URL
// points at a Redis Cluster.
const prefix = `${keyPrefix}:{${encodeURIComponent(ip)}}`;
return { requests: `${prefix}:requests`, spend: `${prefix}:spend`, sequence: `${prefix}:sequence` };
}
return {
async check(ip) {
const timestamp = now();
const keys = keysFor(ip);
const result = await client.eval(REDIS_CHECK_SCRIPT, {
keys: [keys.requests, keys.spend, keys.sequence],
arguments: [
String(timestamp),
String(timestamp - SECOND),
String(timestamp - FIVE_HOURS),
String(RATE_LIMITS.requestsPerSecond),
String(RATE_LIMITS.requestsPerFiveHours),
String(RATE_LIMITS.spendPerFiveHours),
String(FIVE_HOURS),
],
});
if (Number(result[0]) === 1) return { allowed: true };
return {
allowed: false,
reason: result[1],
retryAfter: Math.max(1, Math.ceil(Number(result[2]) / 1_000)),
};
},
async recordCost(ip, amount) {
if (!Number.isFinite(amount) || amount < 0) return;
const timestamp = now();
const keys = keysFor(ip);
await client.eval(REDIS_RECORD_COST_SCRIPT, {
keys: [keys.spend, keys.sequence],
arguments: [String(timestamp), String(timestamp - FIVE_HOURS), String(amount), String(FIVE_HOURS)],
});
},
};
}
export function numericCost(payload) { export function numericCost(payload) {
const cost = payload?.cost ?? payload?.usage?.cost; const cost = payload?.cost ?? payload?.usage?.cost;
return typeof cost === 'number' && Number.isFinite(cost) ? cost : null; return typeof cost === 'number' && Number.isFinite(cost) ? cost : null;
+12 -1
View File
@@ -1,14 +1,24 @@
import path from 'node:path'; import path from 'node:path';
import { createClient } from 'redis';
import { createProxyApp } from './proxy.js'; import { createProxyApp } from './proxy.js';
import { projectRoot } from './env.js'; import { projectRoot } from './env.js';
import { createPool, createRequestLogger, ensureSchema, removeNonChatCompletionRequests } from './db.js'; import { createPool, createRequestLogger, ensureSchema, removeNonChatCompletionRequests } from './db.js';
import { loadConfiguredKeys } from './key-store.js'; import { loadConfiguredKeys } from './key-store.js';
import { createRedisRateLimiter } from './rate-limit.js';
const keys = await loadConfiguredKeys({ path: path.join(projectRoot, 'keys.txt') }); const keys = await loadConfiguredKeys({ path: path.join(projectRoot, 'keys.txt') });
const pool = createPool(); const pool = createPool();
await ensureSchema(pool); await ensureSchema(pool);
await removeNonChatCompletionRequests(pool); await removeNonChatCompletionRequests(pool);
const app = createProxyApp({ keys, requestLogger: createRequestLogger(pool) }); let redis;
let rateLimiter;
if (process.env.REDIS_URL) {
redis = createClient({ url: process.env.REDIS_URL });
redis.on('error', (error) => console.error('Redis rate limiter error:', error.message));
await redis.connect();
rateLimiter = createRedisRateLimiter({ client: redis, keyPrefix: process.env.REDIS_RATE_LIMIT_PREFIX || undefined });
}
const app = createProxyApp({ keys, requestLogger: createRequestLogger(pool), rateLimiter });
const port = Number(process.env.PORT || 4005); const port = Number(process.env.PORT || 4005);
const server = app.listen(port, '0.0.0.0', () => { const server = app.listen(port, '0.0.0.0', () => {
@@ -17,6 +27,7 @@ const server = app.listen(port, '0.0.0.0', () => {
async function shutdown() { async function shutdown() {
server.close(async () => { server.close(async () => {
if (redis?.isOpen) await redis.quit();
await pool.end(); await pool.end();
process.exit(0); process.exit(0);
}); });
+24 -2
View File
@@ -1,5 +1,5 @@
import { describe, expect, it } from 'vitest'; import { describe, expect, it, vi } from 'vitest';
import { createRateLimiter } from '../src/rate-limit.js'; import { createRateLimiter, createRedisRateLimiter } from '../src/rate-limit.js';
describe('rate limiter', () => { describe('rate limiter', () => {
it('enforces request and spend windows', () => { it('enforces request and spend windows', () => {
@@ -23,4 +23,26 @@ describe('rate limiter', () => {
expect(limiter.check('a').allowed).toBe(false); expect(limiter.check('a').allowed).toBe(false);
expect(limiter.check('b').allowed).toBe(true); expect(limiter.check('b').allowed).toBe(true);
}); });
it('uses Redis scripts and namespaces keys by client IP', async () => {
const client = { eval: vi.fn().mockResolvedValueOnce([1]).mockResolvedValueOnce(1) };
const limiter = createRedisRateLimiter({ client, now: () => 10_000, keyPrefix: 'test:limits' });
await expect(limiter.check('2001:db8::1')).resolves.toEqual({ allowed: true });
await limiter.recordCost('2001:db8::1', 1.25);
expect(client.eval).toHaveBeenCalledTimes(2);
expect(client.eval.mock.calls[0][1].keys).toEqual([
'test:limits:{2001%3Adb8%3A%3A1}:requests',
'test:limits:{2001%3Adb8%3A%3A1}:spend',
'test:limits:{2001%3Adb8%3A%3A1}:sequence',
]);
expect(client.eval.mock.calls[1][1].arguments).toEqual(['10000', '-17990000', '1.25', '18000000']);
});
it('converts Redis denials into rate-limit responses', async () => {
const client = { eval: vi.fn().mockResolvedValue([0, 'spend', 1_001]) };
const limiter = createRedisRateLimiter({ client, now: () => 10_000 });
await expect(limiter.check('ip')).resolves.toEqual({ allowed: false, reason: 'spend', retryAfter: 2 });
});
}); });