diff --git a/src/shared/users/user-store.ts b/src/shared/users/user-store.ts index 3856892..7665e55 100644 --- a/src/shared/users/user-store.ts +++ b/src/shared/users/user-store.ts @@ -51,16 +51,16 @@ const migrateTokenCountsProperty = ( const configValue = defaultConfigForProperty[family]; if (typeof dbValue === 'number') { - // Case 1: DB has old numeric format - migrate and add legacy_total - result[family] = { input: dbValue, output: 0, legacy_total: dbValue }; + // Case 1: DB has old numeric format - migrate to legacy_total only (no double counting) + result[family] = { input: 0, output: 0, legacy_total: dbValue }; } else if (typeof dbValue === 'object' && dbValue !== null && (typeof dbValue.input === 'number' || typeof dbValue.output === 'number')) { // Case 2: DB has new object format (might or might not have legacy_total from a previous migration) result[family] = { input: dbValue.input ?? 0, output: dbValue.output ?? 0, legacy_total: dbValue.legacy_total }; } else { // Case 3: DB value is missing or invalid, use default from config if (typeof configValue === 'number') { - // Default from config is old numeric format (e.g., config.tokenQuota[family]) - migrate and add legacy_total - result[family] = { input: configValue, output: 0, legacy_total: configValue }; + // Default from config is old numeric format - migrate to legacy_total only + result[family] = { input: 0, output: 0, legacy_total: configValue }; } else if (typeof configValue === 'object' && configValue !== null && (typeof configValue.input === 'number' || typeof configValue.output === 'number')) { // Default from config is new object format (e.g., INITIAL_TOKENS[family]) result[family] = { input: configValue.input ?? 0, output: configValue.output ?? 0, legacy_total: configValue.legacy_total }; @@ -276,9 +276,16 @@ export function incrementTokenCount( if (!user) return; const modelFamily = getModelFamilyForQuotaUsage(model, api); const existingCounts = user.tokenCounts[modelFamily] ?? { input: 0, output: 0 }; + + // Ensure consumption values are non-negative + const safeInput = Math.max(0, consumption.input); + const safeOutput = Math.max(0, consumption.output); + user.tokenCounts[modelFamily] = { - input: (existingCounts.input ?? 0) + consumption.input, - output: (existingCounts.output ?? 0) + consumption.output, + input: (existingCounts.input ?? 0) + safeInput, + output: (existingCounts.output ?? 0) + safeOutput, + // Preserve legacy_total if it exists + legacy_total: existingCounts.legacy_total }; usersToFlush.add(token); } @@ -338,7 +345,26 @@ export function hasAvailableQuota({ const currentUsage = tokenCounts[modelFamily] ?? { input: 0, output: 0 }; // Calculate total tokens consumed so far (including legacy) - const totalConsumed = (currentUsage.input ?? 0) + (currentUsage.output ?? 0) + (currentUsage.legacy_total ?? 0); + // Ensure all values are non-negative to prevent overflow issues + const input = Math.max(0, currentUsage.input ?? 0); + const output = Math.max(0, currentUsage.output ?? 0); + const legacy = Math.max(0, currentUsage.legacy_total ?? 0); + + // Use safe addition to prevent integer overflow + const totalConsumed = input + output + legacy; + + // Sanity check - if total is negative or NaN, something went wrong + if (!Number.isFinite(totalConsumed) || totalConsumed < 0) { + log.error({ + userToken, + modelFamily, + input, + output, + legacy, + totalConsumed + }, "Invalid token consumption calculation"); + return false; + } // Get the quota limit as a single number const limit = tokenLimits[modelFamily] ?? config.tokenQuota[modelFamily] ?? 0; @@ -346,9 +372,11 @@ export function hasAvailableQuota({ // If no limit (0 or undefined), quota is unlimited if (!limit || limit === 0) return true; + // Ensure requested is non-negative + const safeRequested = Math.max(0, requested); + // Check if the request would exceed the limit - // 'requested' is already the sum of input and output tokens from the middleware - return (totalConsumed + requested) <= limit; + return (totalConsumed + safeRequested) <= limit; } /** @@ -381,7 +409,13 @@ export function resetUsage(token: string) { if (!user) return; const { tokenCounts } = user; for (const family of MODEL_FAMILIES) { - tokenCounts[family] = { input: 0, output: 0 }; // legacy_total is implicitly undefined/removed + const existing = tokenCounts[family]; + // Preserve legacy_total when resetting usage + tokenCounts[family] = { + input: 0, + output: 0, + legacy_total: existing?.legacy_total + }; } usersToFlush.add(token); }