Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ export interface ApiKeyUsageLimitPayload {
dailySpentUsd: number;
weeklySpentUsd: number;
dailyWindowStartIso: string;
dailyResetAtIso: string;
weeklyWindowStartIso: string;
weeklyResetAtIso: string | null;
dailyExceeded: boolean;
Expand Down
184 changes: 168 additions & 16 deletions src/lib/usage/apiKeyUsageLimits.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import { calculateCost } from "./costCalculator";
import { buildErrorBody, sanitizeErrorMessage } from "@omniroute/open-sse/utils/error.ts";

const FORTALEZA_UTC_OFFSET_MS = 3 * 60 * 60 * 1000;
const DAY_MS = 24 * 60 * 60 * 1000;
const WEEK_MS = 7 * 24 * 60 * 60 * 1000;

export interface ApiKeyUsageLimitMetadata {
Expand All @@ -21,6 +22,7 @@ export interface ApiKeyUsageLimitStatus {
dailySpentUsd: number;
weeklySpentUsd: number;
dailyWindowStartIso: string;
dailyResetAtIso: string;
weeklyWindowStartIso: string;
weeklyResetAtIso: string | null;
dailyExceeded: boolean;
Expand All @@ -46,6 +48,18 @@ interface UsageCostRow {
reasoningTokens: number | null;
}

interface WeeklyResetCandidate {
connectionId: string;
resetAtIso: string;
observedWindowStartIso: string | null;
}

interface QuotaSnapshotRow {
remainingPercentage: number | null;
nextResetAt: string | null;
createdAt: string | null;
}

function toNumber(value: unknown): number {
if (typeof value === "number" && Number.isFinite(value)) return value;
if (typeof value === "string" && value.trim()) {
Expand All @@ -70,11 +84,49 @@ function roundUsd(value: number): number {
return Math.round(value * 1_000_000) / 1_000_000;
}

function clampPercent(value: number): number {
return Math.max(0, Math.min(100, value));
}

function formatUsd(value: number | null): string {
if (value === null || !Number.isFinite(value)) return "Not configured";
return `$${value.toFixed(2)}`;
}

function getUsagePercent(spentUsd: number, limitUsd: number | null): number | null {
if (limitUsd === null || !Number.isFinite(limitUsd) || limitUsd <= 0) return null;
return (spentUsd / limitUsd) * 100;
}

function formatUsagePercent(percent: number | null): string {
if (percent === null || !Number.isFinite(percent)) return "Unavailable";
return `${Math.round(percent)}%`;
}

function formatResetIn(resetAt: string | null, now = Date.now()): string {
if (!resetAt) return "unknown";
const resetMs = Date.parse(resetAt);
if (!Number.isFinite(resetMs)) return "unknown";

const deltaMs = resetMs - now;
if (deltaMs <= 0) return "now";

const minuteMs = 60_000;
const hourMs = 60 * minuteMs;
const dayMs = 24 * hourMs;

if (deltaMs < hourMs) return `${Math.max(1, Math.ceil(deltaMs / minuteMs))}m`;
if (deltaMs < dayMs) return `${Math.max(1, Math.ceil(deltaMs / hourMs))}h`;
return `${Math.max(1, Math.ceil(deltaMs / dayMs))}d`;
}

function resetDay(value: string | null): string | null {
if (!value) return null;
const parsed = Date.parse(value);
if (!Number.isFinite(parsed)) return null;
return new Date(parsed).toISOString().slice(0, 10);
}

export function getFortalezaDayStartIso(nowMs = Date.now()): string {
const fortalezaLocal = new Date(nowMs - FORTALEZA_UTC_OFFSET_MS);
return new Date(
Expand All @@ -90,6 +142,10 @@ export function getFortalezaDayStartIso(nowMs = Date.now()): string {
).toISOString();
}

export function getFortalezaDayResetIso(nowMs = Date.now()): string {
return new Date(Date.parse(getFortalezaDayStartIso(nowMs)) + DAY_MS).toISOString();
}

export function getRollingWeekStartIso(nowMs = Date.now()): string {
return new Date(nowMs - WEEK_MS).toISOString();
}
Expand Down Expand Up @@ -144,6 +200,64 @@ function connectionFromValue(value: unknown): { id: string; provider: string } |
return { id, provider };
}

function isWeeklyQuotaResetSnapshot(row: QuotaSnapshotRow, targetResetAtIso: string): boolean {
const targetDay = resetDay(targetResetAtIso);
if (!targetDay) return false;
return resetDay(row.nextResetAt) === targetDay;
}

function getObservedWeeklyWindowStartIso(
connectionId: string,
targetResetAtIso: string,
nowMs: number
): string | null {
if (!connectionId || !targetResetAtIso) return null;

try {
const rows = getDbInstance()
.prepare(
`
SELECT
remaining_percentage as remainingPercentage,
next_reset_at as nextResetAt,
created_at as createdAt
FROM quota_snapshots
WHERE connection_id = @connectionId
AND LOWER(window_key) LIKE '%weekly%'
AND LOWER(window_key) NOT LIKE '%sonnet%'
AND created_at <= @nowIso
ORDER BY created_at ASC, id ASC
`
)
.all({ connectionId, nowIso: new Date(nowMs).toISOString() }) as QuotaSnapshotRow[];
Comment on lines +209 to +232

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Performance Bottleneck: Unbounded Database Query on Hot Path

The current query in getObservedWeeklyWindowStartIso fetches all weekly snapshots for the given connection from the database without any date range or reset day constraints. Since this function is called on every single API request (via getApiKeyUsageLimitStatus to check usage limits), and snapshots are retained for up to 90 days, this will load, parse, and iterate over thousands of irrelevant rows in memory on every request.

Solution

  1. Filter by next_reset_at LIKE @targetDayPattern to only fetch snapshots that match the target reset day.
  2. Filter by created_at >= @sinceIso (e.g., 8 days ago) to restrict the scan range to the current weekly cycle.
function getObservedWeeklyWindowStartIso(
  connectionId: string, 
  targetResetAtIso: string, 
  nowMs: number
): string | null {
  if (!connectionId || !targetResetAtIso) return null;

  const targetDay = resetDay(targetResetAtIso);
  if (!targetDay) return null;

  const sinceIso = new Date(nowMs - 8 * 24 * 60 * 60 * 1000).toISOString();

  try {
    const rows = getDbInstance()
      .prepare(
        `
        SELECT
          remaining_percentage as remainingPercentage,
          next_reset_at as nextResetAt,
          created_at as createdAt
        FROM quota_snapshots
        WHERE connection_id = @connectionId
          AND LOWER(window_key) LIKE '%weekly%'
          AND LOWER(window_key) NOT LIKE '%sonnet%'
          AND created_at <= @nowIso
          AND created_at >= @sinceIso
          AND next_reset_at LIKE @targetDayPattern
        ORDER BY created_at ASC, id ASC
      `
      )
      .all({
        connectionId,
        nowIso: new Date(nowMs).toISOString(),
        sinceIso,
        targetDayPattern: `${targetDay}%`,
      }) as QuotaSnapshotRow[];


let observedStartIso: string | null = null;
let previousUsedPercent: number | null = null;

for (const row of rows) {
if (!row.createdAt || !isWeeklyQuotaResetSnapshot(row, targetResetAtIso)) continue;
const remaining = toNumber(row.remainingPercentage);
const usedPercent = clampPercent(100 - remaining);

if (!observedStartIso) {
observedStartIso = row.createdAt;
} else if (previousUsedPercent !== null) {
const droppedToResetFloor = usedPercent <= 1 && previousUsedPercent > usedPercent;
const significantDrop = previousUsedPercent - usedPercent >= 5;
if (droppedToResetFloor || significantDrop) {
observedStartIso = row.createdAt;
}
}

previousUsedPercent = usedPercent;
}

return observedStartIso;
} catch {
return null;
}
}

async function resolveDeps(deps: ApiKeyUsageLimitDeps): Promise<Required<ApiKeyUsageLimitDeps>> {
const providers =
deps.getProviderConnectionById && deps.getProviderConnections
Expand All @@ -165,16 +279,16 @@ async function resolveDeps(deps: ApiKeyUsageLimitDeps): Promise<Required<ApiKeyU
};
}

async function getProviderWeeklyResetAt(
async function getProviderWeeklyWindow(
metadata: ApiKeyUsageLimitMetadata,
deps: Required<ApiKeyUsageLimitDeps>,
nowMs: number
): Promise<string | null> {
): Promise<{ resetAtIso: string | null; windowStartIso: string | null }> {
const allowedConnections = Array.isArray(metadata.allowedConnections)
? metadata.allowedConnections.filter((id) => typeof id === "string" && id.trim())
: [];

const resetCandidates: string[] = [];
const resetCandidates: WeeklyResetCandidate[] = [];
if (allowedConnections.length > 0) {
for (const connectionId of allowedConnections) {
const connection = connectionFromValue(await deps.getProviderConnectionById(connectionId));
Expand All @@ -183,7 +297,13 @@ async function getProviderWeeklyResetAt(
deps.getProviderLimitsCache(connection.id)?.quotas,
nowMs
);
if (resetAt) resetCandidates.push(resetAt);
if (resetAt) {
resetCandidates.push({
connectionId: connection.id,
resetAtIso: resetAt,
observedWindowStartIso: getObservedWeeklyWindowStartIso(connection.id, resetAt, nowMs),
});
}
}
} else {
const caches = deps.getAllProviderLimitsCache();
Expand All @@ -192,11 +312,24 @@ async function getProviderWeeklyResetAt(
const connection = connectionFromValue(rawConnection);
if (!connection || connection.provider.toLowerCase() !== "claude") continue;
const resetAt = findWeeklyQuotaResetAt(caches[connection.id]?.quotas, nowMs);
if (resetAt) resetCandidates.push(resetAt);
if (resetAt) {
resetCandidates.push({
connectionId: connection.id,
resetAtIso: resetAt,
observedWindowStartIso: getObservedWeeklyWindowStartIso(connection.id, resetAt, nowMs),
});
}
}
}

return resetCandidates.sort((left, right) => Date.parse(left) - Date.parse(right)).at(0) ?? null;
const selected =
resetCandidates
.sort((left, right) => Date.parse(left.resetAtIso) - Date.parse(right.resetAtIso))
.at(0) ?? null;
return {
resetAtIso: selected?.resetAtIso ?? null,
windowStartIso: selected?.observedWindowStartIso ?? null,
};
}

async function getApiKeyUsdSpendSince(apiKeyId: string, sinceIso: string): Promise<number> {
Expand Down Expand Up @@ -257,10 +390,14 @@ export async function getApiKeyUsageLimitStatus(
const resolvedDeps = await resolveDeps(deps);
const now = resolvedDeps.now();
const dailyWindowStartIso = getFortalezaDayStartIso(now);
const weeklyResetAtIso = await getProviderWeeklyResetAt(metadata, resolvedDeps, now);
const weeklyWindowStartIso = weeklyResetAtIso
? new Date(Date.parse(weeklyResetAtIso) - WEEK_MS).toISOString()
: getRollingWeekStartIso(now);
const dailyResetAtIso = getFortalezaDayResetIso(now);
const weeklyWindow = await getProviderWeeklyWindow(metadata, resolvedDeps, now);
const weeklyResetAtIso = weeklyWindow.resetAtIso;
const weeklyWindowStartIso = weeklyWindow.windowStartIso
? weeklyWindow.windowStartIso
: weeklyResetAtIso
? new Date(Date.parse(weeklyResetAtIso) - WEEK_MS).toISOString()
: getRollingWeekStartIso(now);
const dailyLimitUsd = normalizeLimitUsd(metadata.dailyUsageLimitUsd);
const weeklyLimitUsd = normalizeLimitUsd(metadata.weeklyUsageLimitUsd);
const enabled = metadata.usageLimitEnabled === true;
Expand All @@ -277,33 +414,47 @@ export async function getApiKeyUsageLimitStatus(
dailySpentUsd,
weeklySpentUsd,
dailyWindowStartIso,
dailyResetAtIso,
weeklyWindowStartIso,
weeklyResetAtIso,
dailyExceeded: enabled && dailyLimitUsd !== null && dailySpentUsd >= dailyLimitUsd,
weeklyExceeded: enabled && weeklyLimitUsd !== null && weeklySpentUsd >= weeklyLimitUsd,
};
}

export function buildApiKeyUsageLimitText(status: ApiKeyUsageLimitStatus): string {
export function buildApiKeyUsageLimitText(
status: ApiKeyUsageLimitStatus,
now = Date.now()
): string {
return [
"Cota diaria",
formatUsd(status.dailyLimitUsd),
"Gasto diario",
formatUsd(status.dailySpentUsd),
"Uso diario",
formatUsagePercent(getUsagePercent(status.dailySpentUsd, status.dailyLimitUsd)),
`Resets in ${formatResetIn(status.dailyResetAtIso, now)}`,
"",
"Cota semanal",
formatUsd(status.weeklyLimitUsd),
"Gasto semanal",
formatUsd(status.weeklySpentUsd),
"Uso semanal",
formatUsagePercent(getUsagePercent(status.weeklySpentUsd, status.weeklyLimitUsd)),
`Resets in ${formatResetIn(status.weeklyResetAtIso, now)}`,
].join("\n");
}

function buildUsageLimitExceededMessage(status: ApiKeyUsageLimitStatus): string {
function buildUsageLimitExceededMessage(status: ApiKeyUsageLimitStatus, now = Date.now()): string {
if (status.dailyExceeded && status.dailyLimitUsd !== null) {
return `This API key reached its daily USD usage quota (${formatUsd(status.dailySpentUsd)} of ${formatUsd(status.dailyLimitUsd)}). Choose another allowed model or wait for quota reset.`;
const percent = formatUsagePercent(getUsagePercent(status.dailySpentUsd, status.dailyLimitUsd));
return `This API key reached its daily USD usage quota (${formatUsd(status.dailySpentUsd)} of ${formatUsd(status.dailyLimitUsd)}, ${percent}). Resets in ${formatResetIn(status.dailyResetAtIso, now)}. Choose another allowed model after reset.`;
}
if (status.weeklyExceeded && status.weeklyLimitUsd !== null) {
return `This API key reached its weekly USD usage quota (${formatUsd(status.weeklySpentUsd)} of ${formatUsd(status.weeklyLimitUsd)}). Choose another allowed model or wait for quota reset.`;
const percent = formatUsagePercent(
getUsagePercent(status.weeklySpentUsd, status.weeklyLimitUsd)
);
return `This API key reached its weekly USD usage quota (${formatUsd(status.weeklySpentUsd)} of ${formatUsd(status.weeklyLimitUsd)}, ${percent}). Resets in ${formatResetIn(status.weeklyResetAtIso, now)}. Choose another allowed model after reset.`;
}
return "This API key reached its USD usage quota. Choose another allowed model or wait for quota reset.";
}
Expand All @@ -319,9 +470,10 @@ function isAnthropicMessagesRequest(request: Request): boolean {

export function buildApiKeyUsageLimitRejection(
request: Request,
status: ApiKeyUsageLimitStatus
status: ApiKeyUsageLimitStatus,
now = Date.now()
): Response {
const message = sanitizeErrorMessage(buildUsageLimitExceededMessage(status));
const message = sanitizeErrorMessage(buildUsageLimitExceededMessage(status, now));
if (isAnthropicMessagesRequest(request)) {
return new Response(
JSON.stringify({
Expand Down
3 changes: 2 additions & 1 deletion src/lib/usage/internalUsageCommand.ts
Original file line number Diff line number Diff line change
Expand Up @@ -376,7 +376,8 @@ export async function buildUsageCommandText(
const resolvedDeps = await normalizeDeps(deps);
if (metadata.usageLimitEnabled === true) {
return buildApiKeyUsageLimitText(
await resolvedDeps.getApiKeyUsageLimitStatus(metadata, { now: resolvedDeps.now })
await resolvedDeps.getApiKeyUsageLimitStatus(metadata, { now: resolvedDeps.now }),
resolvedDeps.now()
);
}

Expand Down
Loading
Loading