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
This commit is contained in:
oyi77
2026-07-12 06:20:04 +07:00
parent b6ae24306a
commit 9d4cd90e70
3 changed files with 19 additions and 15 deletions

View File

@@ -176,7 +176,7 @@ export async function POST(request: Request) {
}); });
} }
const rateCheck = checkRateLimit(token.id); const rateCheck = checkRateLimit(token.id, token);
if (!rateCheck.allowed) { if (!rateCheck.allowed) {
recordRelayUsage(token.id, { recordRelayUsage(token.id, {
requestId: request.headers.get("x-request-id") || undefined, requestId: request.headers.get("x-request-id") || undefined,

View File

@@ -215,7 +215,7 @@ export async function POST(request: Request) {
} }
// 2b. Per-token rate limit check // 2b. Per-token rate limit check
const rateCheck = checkRateLimit(token.id); const rateCheck = checkRateLimit(token.id, token);
if (!rateCheck.allowed) { if (!rateCheck.allowed) {
recordRelayUsage(token.id, { recordRelayUsage(token.id, {
requestId: request.headers.get("x-request-id") || undefined, requestId: request.headers.get("x-request-id") || undefined,
@@ -303,8 +303,7 @@ export async function POST(request: Request) {
bifrostFallbackReason = bifrostDecision.fallbackReason; bifrostFallbackReason = bifrostDecision.fallbackReason;
} }
if (bifrostDecision.tryBifrost) { if (bifrostDecision.tryBifrost) {
const cooldown = const cooldown = backend === "auto" ? getActiveBifrostCooldown(bifrostConfig.baseUrl) : null;
backend === "auto" ? getActiveBifrostCooldown(bifrostConfig.baseUrl) : null;
if (cooldown) { if (cooldown) {
bifrostFallbackReason = `bifrost-cooldown; remaining=${cooldown.remainingMs}`; bifrostFallbackReason = `bifrost-cooldown; remaining=${cooldown.remainingMs}`;
} else { } else {

View File

@@ -153,8 +153,7 @@ export function getRelayTokens(): RelayToken[] {
export function getRelayToken(id: string): RelayToken | null { export function getRelayToken(id: string): RelayToken | null {
const db = getDbInstance(); const db = getDbInstance();
const row = db.prepare("SELECT * FROM relay_tokens WHERE id = ?").get(id) as const row = db.prepare("SELECT * FROM relay_tokens WHERE id = ?").get(id) as
| RelayTokenRow RelayTokenRow | undefined;
| undefined;
if (!row) return null; if (!row) return null;
return { ...(rowToCamel(row) as unknown as RelayToken), enabled: row.enabled === 1 }; 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 ─────────────────────────────────────────────────────── // ── Usage / Rate Limit ───────────────────────────────────────────────────────
export function checkRateLimit(tokenId: string): { export function checkRateLimit(
tokenId: string,
existingToken?: RelayToken
): {
allowed: boolean; allowed: boolean;
remaining: number; remaining: number;
resetIn: number; resetIn: number;
} { } {
const db = getDbInstance(); const db = getDbInstance();
const token = db.prepare("SELECT * FROM relay_tokens WHERE id = ?").get(tokenId) as let token = existingToken;
| RelayTokenRow if (!token) {
| undefined; const row = db.prepare("SELECT * FROM relay_tokens WHERE id = ?").get(tokenId) as
if (!token) return { allowed: false, remaining: 0, resetIn: 0 }; 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 now = Math.floor(Date.now() / 1000);
const minuteWindow = Math.floor(now / 60) * 60; 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; .get(tokenId, minuteWindow) as { request_count: number; cost: number } | undefined;
const minuteCount = minuteRow?.request_count || 0; 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) }; 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; .get(tokenId, dayWindow) as { total: number } | undefined;
const dayCount = dayRow?.total || 0; 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) }; return { allowed: false, remaining: 0, resetIn: 86400 - (now % 86400) };
} }
const remaining = Math.min( const remaining = Math.min(
token.max_requests_per_minute - minuteCount, token.maxRequestsPerMinute - minuteCount,
token.max_requests_per_day - dayCount token.maxRequestsPerDay - dayCount
); );
return { allowed: true, remaining, resetIn: 60 - (now % 60) }; return { allowed: true, remaining, resetIn: 60 - (now % 60) };