From 9d4cd90e700fe13db2ccc88854627e396737cf9d Mon Sep 17 00:00:00 2001
From: oyi77 <oyi77@users.noreply.github.com>
Date: Sun, 12 Jul 2026 06:20:04 +0700
Subject: [PATCH] perf: thread pre-fetched token to checkRateLimit avoiding
 re-query

getRelayTokenByHash already fetches the full RelayToken row. A few
lines later checkRateLimit(token.id) does a second SELECT * FROM
relay_tokens on a different predicate (id instead of token_hash).

Change:
- checkRateLimit accepts an optional existingToken parameter; when
  provided, skips the re-query entirely.
- Both relay routes (chat completions + bifrost) pass the already-
  fetched token.
- The function now uses RelayToken (camelCase) instead of RelayTokenRow
  (snake_case) when the token is passed in.

PR-URL: fix-relay-thread-token
---
 .../relay/chat/completions/bifrost/route.ts   |  2 +-
 .../api/v1/relay/chat/completions/route.ts    |  5 ++--
 src/lib/db/relayProxies.ts                    | 27 +++++++++++--------
 3 files changed, 19 insertions(+), 15 deletions(-)

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