key routing persists across restarts

.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.

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

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];
+    },
+  };
+}

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 } = {}) {
+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 = createKeyRotator(keys);
+  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()}`,
           },
         });
 

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', () => {

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');
+  });
 });