diff --git a/src/shared/users/user-store.ts b/src/shared/users/user-store.ts index 7665e55..9a3c5a4 100644 --- a/src/shared/users/user-store.ts +++ b/src/shared/users/user-store.ts @@ -381,24 +381,31 @@ export function hasAvailableQuota({ /** * 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);