diff --git a/.env.example b/.env.example index 95f2b31..6504728 100644 --- a/.env.example +++ b/.env.example @@ -12,6 +12,8 @@ DATABASE_URL=postgres://user:password@localhost:5432/opencode_proxy # REDIS_URL=redis://localhost:6379 # Optional Redis key namespace, useful when sharing a Redis database. # REDIS_RATE_LIMIT_PREFIX=opencode-proxy:rate-limit +# Optional Redis counter key for upstream API-key round-robin routing. +# REDIS_KEY_ROTATION_KEY=opencode-proxy:key-rotation # 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. diff --git a/README.md b/README.md index eea1fc0..59d65e2 100644 --- a/README.md +++ b/README.md @@ -37,7 +37,7 @@ 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. -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. +Set `REDIS_URL` to persist and share these limits across process restarts and multiple proxy instances. The rate-limit backend stores only the active five-hour window and uses atomic operations, so concurrent instances enforce one shared limit. Redis also keeps the upstream API-key round-robin cursor in a counter, so routing continues across restarts and is shared by every proxy instance. Without `REDIS_URL`, both limits and key routing remain in memory as before. Use `REDIS_KEY_ROTATION_KEY` to change that counter's key when sharing a Redis database. ```sh REDIS_URL=redis://localhost:6379 diff --git a/src/key-store.js b/src/key-store.js index 9ea96b3..bfe77ea 100644 --- a/src/key-store.js +++ b/src/key-store.js @@ -52,3 +52,24 @@ export function createKeyRotator(keys) { }, }; } + +/** + * Creates a Redis-backed round-robin key selector. INCR is atomic, so proxy + * instances sharing Redis also share one selection sequence. + */ +export function createRedisKeyRotator({ client, keys, key = 'opencode-proxy:key-rotation' } = {}) { + if (!client?.incr) throw new Error('A connected Redis client is required for Redis key rotation'); + if (!Array.isArray(keys) || keys.length === 0) { + throw new Error('At least one OpenCode Go API key is required'); + } + + return { + async next() { + const sequence = Number(await client.incr(key)); + if (!Number.isSafeInteger(sequence) || sequence < 1) { + throw new Error('Redis returned an invalid key-rotation sequence'); + } + return keys[(sequence - 1) % keys.length]; + }, + }; +} diff --git a/src/proxy.js b/src/proxy.js index 2a7d058..3779e91 100644 --- a/src/proxy.js +++ b/src/proxy.js @@ -187,8 +187,8 @@ function validateAllowedModel(model, allowedModels) { } } -export function createProxyApp({ keys, fetchImpl = globalThis.fetch, upstreamBaseUrl = DEFAULT_UPSTREAM_BASE_URL, maxContextTokens = maxContextFromEnv(), rateLimiter = createRateLimiter(), unlimitedKeys = unlimitedKeysFromEnv(), allowedModels = allowedModelsFromEnv(), pagesDirectory = defaultPagesDirectory, requestLogger } = {}) { - const rotator = createKeyRotator(keys); +export function createProxyApp({ keys, keyRotator = createKeyRotator(keys), fetchImpl = globalThis.fetch, upstreamBaseUrl = DEFAULT_UPSTREAM_BASE_URL, maxContextTokens = maxContextFromEnv(), rateLimiter = createRateLimiter(), unlimitedKeys = unlimitedKeysFromEnv(), allowedModels = allowedModelsFromEnv(), pagesDirectory = defaultPagesDirectory, requestLogger } = {}) { + const rotator = keyRotator; const app = express(); if (requestLogger) app.post(CHAT_COMPLETION_ENDPOINTS, collectRequestData(requestLogger)); app.use(express.json({ limit: '10mb' })); @@ -211,7 +211,7 @@ export function createProxyApp({ keys, fetchImpl = globalThis.fetch, upstreamBas ...options, headers: { ...options.headers, - authorization: `Bearer ${rotator.next()}`, + authorization: `Bearer ${await rotator.next()}`, }, }); diff --git a/src/server.js b/src/server.js index b39677c..c72acc8 100644 --- a/src/server.js +++ b/src/server.js @@ -3,7 +3,7 @@ import { createClient } from 'redis'; import { createProxyApp } from './proxy.js'; import { projectRoot } from './env.js'; import { createPool, createRequestLogger, ensureSchema, removeNonChatCompletionRequests } from './db.js'; -import { loadConfiguredKeys } from './key-store.js'; +import { createRedisKeyRotator, loadConfiguredKeys } from './key-store.js'; import { createRedisRateLimiter } from './rate-limit.js'; const keys = await loadConfiguredKeys({ path: path.join(projectRoot, 'keys.txt') }); @@ -12,13 +12,15 @@ await ensureSchema(pool); await removeNonChatCompletionRequests(pool); let redis; let rateLimiter; +let keyRotator; 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 }); + keyRotator = createRedisKeyRotator({ client: redis, keys, key: process.env.REDIS_KEY_ROTATION_KEY || undefined }); } -const app = createProxyApp({ keys, requestLogger: createRequestLogger(pool), rateLimiter }); +const app = createProxyApp({ keys, keyRotator, requestLogger: createRequestLogger(pool), rateLimiter }); const port = Number(process.env.PORT || 4005); const server = app.listen(port, '0.0.0.0', () => { diff --git a/test/key-store.test.js b/test/key-store.test.js index e4ca0a5..8dd0524 100644 --- a/test/key-store.test.js +++ b/test/key-store.test.js @@ -1,5 +1,5 @@ import { describe, expect, it } from 'vitest'; -import { createKeyRotator, loadConfiguredKeys, parseKeys, parseKeysJson } from '../src/key-store.js'; +import { createKeyRotator, createRedisKeyRotator, loadConfiguredKeys, parseKeys, parseKeysJson } from '../src/key-store.js'; describe('key store', () => { it('parses JSON keys and prefers the environment variable', async () => { @@ -16,4 +16,21 @@ describe('key store', () => { const rotator = createKeyRotator(['a', 'b']); expect([rotator.next(), rotator.next(), rotator.next()]).toEqual(['a', 'b', 'a']); }); + + it('shares a Redis-backed rotation sequence across rotators', async () => { + const counters = new Map(); + const client = { + incr: async (key) => { + const next = (counters.get(key) || 0) + 1; + counters.set(key, next); + return String(next); + }, + }; + const firstInstance = createRedisKeyRotator({ client, keys: ['a', 'b'], key: 'test:key-rotation' }); + const restartedInstance = createRedisKeyRotator({ client, keys: ['a', 'b'], key: 'test:key-rotation' }); + + await expect(firstInstance.next()).resolves.toBe('a'); + await expect(firstInstance.next()).resolves.toBe('b'); + await expect(restartedInstance.next()).resolves.toBe('a'); + }); });