mirror of
https://github.com/diegosouzapw/OmniRoute.git
synced 2026-07-26 09:52:11 +03:00
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:
@@ -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,
|
||||||
|
|||||||
@@ -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 {
|
||||||
|
|||||||
@@ -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) };
|
||||||
|
|||||||
Reference in New Issue
Block a user