diff --git a/changelog.d/fixes/6930-relay-checkratelimit-token-reuse.md b/changelog.d/fixes/6930-relay-checkratelimit-token-reuse.md new file mode 100644 index 0000000000..e52bb9d46a --- /dev/null +++ b/changelog.d/fixes/6930-relay-checkratelimit-token-reuse.md @@ -0,0 +1 @@ +- perf(api): relay chat-completions routes now thread the already-fetched `RelayToken` into `checkRateLimit`, skipping a redundant `SELECT * FROM relay_tokens WHERE id = ?` re-query on every request (#6930) diff --git a/src/app/api/v1/relay/chat/completions/bifrost/route.ts b/src/app/api/v1/relay/chat/completions/bifrost/route.ts index da8722d11c..b0df931b0e 100644 --- a/src/app/api/v1/relay/chat/completions/bifrost/route.ts +++ b/src/app/api/v1/relay/chat/completions/bifrost/route.ts @@ -176,7 +176,7 @@ export async function POST(request: Request) { }); } - const rateCheck = checkRateLimit(token.id); + const rateCheck = checkRateLimit(token.id, token); if (!rateCheck.allowed) { recordRelayUsage(token.id, { requestId: request.headers.get("x-request-id") || undefined, diff --git a/src/app/api/v1/relay/chat/completions/route.ts b/src/app/api/v1/relay/chat/completions/route.ts index f674327f62..8ee4a25a2d 100644 --- a/src/app/api/v1/relay/chat/completions/route.ts +++ b/src/app/api/v1/relay/chat/completions/route.ts @@ -215,7 +215,7 @@ export async function POST(request: Request) { } // 2b. Per-token rate limit check - const rateCheck = checkRateLimit(token.id); + const rateCheck = checkRateLimit(token.id, token); if (!rateCheck.allowed) { recordRelayUsage(token.id, { requestId: request.headers.get("x-request-id") || undefined, @@ -303,8 +303,7 @@ export async function POST(request: Request) { bifrostFallbackReason = bifrostDecision.fallbackReason; } if (bifrostDecision.tryBifrost) { - const cooldown = - backend === "auto" ? getActiveBifrostCooldown(bifrostConfig.baseUrl) : null; + const cooldown = backend === "auto" ? getActiveBifrostCooldown(bifrostConfig.baseUrl) : null; if (cooldown) { bifrostFallbackReason = `bifrost-cooldown; remaining=${cooldown.remainingMs}`; } else { diff --git a/src/lib/db/relayProxies.ts b/src/lib/db/relayProxies.ts index b88c6db945..5c9ffabe37 100644 --- a/src/lib/db/relayProxies.ts +++ b/src/lib/db/relayProxies.ts @@ -153,8 +153,7 @@ export function getRelayTokens(): RelayToken[] { export function getRelayToken(id: string): RelayToken | null { const db = getDbInstance(); const row = db.prepare("SELECT * FROM relay_tokens WHERE id = ?").get(id) as - | RelayTokenRow - | undefined; + RelayTokenRow | undefined; if (!row) return null; return { ...(rowToCamel(row) as unknown as RelayToken), enabled: row.enabled === 1 }; } @@ -235,16 +234,22 @@ export function toggleRelayToken(id: string, enabled: boolean): RelayToken | nul // ── Usage / Rate Limit ─────────────────────────────────────────────────────── -export function checkRateLimit(tokenId: string): { +export function checkRateLimit( + tokenId: string, + existingToken?: RelayToken +): { allowed: boolean; remaining: number; resetIn: number; } { const db = getDbInstance(); - const token = db.prepare("SELECT * FROM relay_tokens WHERE id = ?").get(tokenId) as - | RelayTokenRow - | undefined; - if (!token) return { allowed: false, remaining: 0, resetIn: 0 }; + let token = existingToken; + if (!token) { + const row = db.prepare("SELECT * FROM relay_tokens WHERE id = ?").get(tokenId) as + RelayTokenRow | undefined; + if (!row) return { allowed: false, remaining: 0, resetIn: 0 }; + token = rowToCamel(row) as unknown as RelayToken; + } const now = Math.floor(Date.now() / 1000); const minuteWindow = Math.floor(now / 60) * 60; @@ -258,7 +263,7 @@ export function checkRateLimit(tokenId: string): { .get(tokenId, minuteWindow) as { request_count: number; cost: number } | undefined; const minuteCount = minuteRow?.request_count || 0; - if (minuteCount >= token.max_requests_per_minute) { + if (minuteCount >= token.maxRequestsPerMinute) { return { allowed: false, remaining: 0, resetIn: 60 - (now % 60) }; } @@ -270,13 +275,13 @@ export function checkRateLimit(tokenId: string): { .get(tokenId, dayWindow) as { total: number } | undefined; const dayCount = dayRow?.total || 0; - if (dayCount >= token.max_requests_per_day) { + if (dayCount >= token.maxRequestsPerDay) { return { allowed: false, remaining: 0, resetIn: 86400 - (now % 86400) }; } const remaining = Math.min( - token.max_requests_per_minute - minuteCount, - token.max_requests_per_day - dayCount + token.maxRequestsPerMinute - minuteCount, + token.maxRequestsPerDay - dayCount ); return { allowed: true, remaining, resetIn: 60 - (now % 60) }; diff --git a/tests/unit/relay-check-rate-limit-existing-token.test.ts b/tests/unit/relay-check-rate-limit-existing-token.test.ts new file mode 100644 index 0000000000..2843c15a25 --- /dev/null +++ b/tests/unit/relay-check-rate-limit-existing-token.test.ts @@ -0,0 +1,129 @@ +import test from "node:test"; +import assert from "node:assert/strict"; +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; + +// Regression test for src/lib/db/relayProxies.ts::checkRateLimit. +// +// Covers the perf change that threads an already-fetched RelayToken into +// checkRateLimit to avoid a redundant `SELECT * FROM relay_tokens WHERE id = ?` +// re-query: +// - the `existingToken` fast-path must agree with the legacy re-query path +// (same allowed/remaining for identical DB state) +// - the legacy re-query path (no token passed) must still work unmodified +// - the per-minute cap must still be enforced correctly via the fast-path + +const TEST_DATA_DIR = fs.mkdtempSync( + path.join(os.tmpdir(), "omniroute-relay-check-rate-limit-") +); +process.env.DATA_DIR = TEST_DATA_DIR; + +const core = await import("../../src/lib/db/core.ts"); +const relayProxies = await import("../../src/lib/db/relayProxies.ts"); + +async function resetStorage() { + core.resetDbInstance(); + + for (let attempt = 0; attempt < 10; attempt++) { + try { + if (fs.existsSync(TEST_DATA_DIR)) { + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); + } + break; + } catch (error) { + const code = (error as NodeJS.ErrnoException)?.code; + if ((code === "EBUSY" || code === "EPERM") && attempt < 9) { + await new Promise((resolve) => setTimeout(resolve, 50 * (attempt + 1))); + } else { + throw error; + } + } + } + + fs.mkdirSync(TEST_DATA_DIR, { recursive: true }); +} + +test.beforeEach(async () => { + await resetStorage(); +}); + +test.after(async () => { + core.resetDbInstance(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); +}); + +// Inserts a relay_tokens row directly (bypassing createRelayToken, which uses +// a CommonJS `require("node:crypto")` that is unavailable under this ESM test +// runner — a pre-existing, unrelated issue) and returns the RelayToken as +// checkRateLimit's existingToken param expects it (camelCase, via getRelayToken). +function insertRelayToken(overrides: { + id: string; + name: string; + maxRequestsPerMinute: number; + maxRequestsPerDay: number; +}) { + const db = core.getDbInstance(); + const now = Math.floor(Date.now() / 1000); + db.prepare( + ` + INSERT INTO relay_tokens (id, name, token_hash, token_prefix, description, combo_id, allowed_models, + max_tokens_per_request, max_requests_per_minute, max_requests_per_day, max_cost_per_day, + enabled, created_at, updated_at, expires_at, metadata) + VALUES (?, ?, ?, ?, '', NULL, '["*"]', 128000, ?, ?, 0, 1, ?, ?, NULL, '{}') + ` + ).run( + overrides.id, + overrides.name, + `hash-${overrides.id}`, + `rl_${overrides.id}`, + overrides.maxRequestsPerMinute, + overrides.maxRequestsPerDay, + now, + now + ); + const token = relayProxies.getRelayToken(overrides.id); + if (!token) throw new Error("failed to insert test relay token"); + return token; +} + +test("checkRateLimit: existingToken fast-path agrees with the legacy re-query path", () => { + const token = insertRelayToken({ + id: "rl_fastpath1", + name: "fast-path-token", + maxRequestsPerMinute: 10, + maxRequestsPerDay: 1000, + }); + + const legacy = relayProxies.checkRateLimit(token.id); + const fastPath = relayProxies.checkRateLimit(token.id, token); + + assert.deepEqual(fastPath.allowed, legacy.allowed); + assert.deepEqual(fastPath.remaining, legacy.remaining); +}); + +test("checkRateLimit: legacy re-query path (no token passed) still works when the token does not exist", () => { + const result = relayProxies.checkRateLimit("does-not-exist"); + assert.equal(result.allowed, false); + assert.equal(result.remaining, 0); +}); + +test("checkRateLimit: existingToken fast-path still enforces the per-minute cap", () => { + const token = insertRelayToken({ + id: "rl_captoken1", + name: "cap-token", + maxRequestsPerMinute: 2, + maxRequestsPerDay: 1000, + }); + + // Record 2 requests in the current minute window — matches the cap. + relayProxies.recordRelayUsage(token.id, { model: "test-model", cost: 0 }); + relayProxies.recordRelayUsage(token.id, { model: "test-model", cost: 0 }); + + const fastPath = relayProxies.checkRateLimit(token.id, token); + const legacy = relayProxies.checkRateLimit(token.id); + + assert.equal(fastPath.allowed, false); + assert.equal(fastPath.remaining, 0); + assert.deepEqual(fastPath, legacy); +});