Merge branch 'main' into 'main'

more quota fixes

See merge request reanon/nonono!5

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,31 +372,40 @@ 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;
 }
 
 /**
  * For the given user, refreshes token limits for each model family. The new limit
- * is set to the configured quota value (either from user's tokenRefresh or global config).
- * This replaces the current limits entirely, not adding to them.
+ * is set to the current usage + the refresh amount, ensuring users get their full
+ * refresh allocation regardless of current usage.
  */
 export function refreshQuota(token: string) {
   const user = users.get(token);
   if (!user) return;
-  const { tokenLimits, tokenRefresh } = user;
+  const { tokenCounts, tokenLimits, tokenRefresh } = user;
 
   for (const family of MODEL_FAMILIES) {
-    // Get the quota value to set (from user refresh config or global default)
+    // Get the quota value to add (from user refresh config or global default)
     const userQuota = tokenRefresh[family] ?? 0;
     const globalQuota = config.tokenQuota[family] ?? 0;
-    
-    const quotaToSet = userQuota || globalQuota;
+    const quotaToAdd = userQuota || globalQuota;
 
-    // Only update if we have a valid quota
-    if (quotaToSet > 0) {
-      tokenLimits[family] = quotaToSet;
+    if (quotaToAdd > 0) {
+      // Calculate current usage including legacy
+      const currentUsage = tokenCounts[family] ?? { input: 0, output: 0 };
+      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);
+      const totalUsage = input + output + legacy;
+      
+      // Set new limit to current usage + refresh amount
+      // This ensures users always get their full refresh allocation
+      tokenLimits[family] = totalUsage + quotaToAdd;
     }
   }
   usersToFlush.add(token);
@@ -381,7 +416,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);
 }