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