adds user-specific overrides for daily quota refresh

src/shared/users/schema.ts

@@ -37,6 +37,8 @@ export const UserSchema = z
     tokenCounts: tokenCountsSchema,
     /** Maximum number of tokens the user can consume, by model family. */
     tokenLimits: tokenCountsSchema,
+    /** User-specific token refresh amount, by model family. */
+    tokenRefresh: tokenCountsSchema,
     /** Time at which the user was created. */
     createdAt: z.number(),
     /** Time at which the user last connected. */

src/shared/users/user-store.ts

@@ -70,6 +70,7 @@ export function createUser(createOptions?: {
   type?: User["type"];
   expiresAt?: number;
   tokenLimits?: User["tokenLimits"];
+  tokenRefresh?: User["tokenRefresh"];
 }) {
   const token = uuid();
   const newUser: User = {
@@ -79,6 +80,7 @@ export function createUser(createOptions?: {
     promptCount: 0,
     tokenCounts: { ...INITIAL_TOKENS },
     tokenLimits: createOptions?.tokenLimits ?? { ...config.tokenQuota },
+    tokenRefresh: createOptions?.tokenRefresh ?? { ...INITIAL_TOKENS },
     createdAt: Date.now(),
     meta: {},
   };
@@ -123,6 +125,7 @@ export function upsertUser(user: UserUpdate) {
     promptCount: 0,
     tokenCounts: { ...INITIAL_TOKENS },
     tokenLimits: { ...config.tokenQuota },
+    tokenRefresh: { ...INITIAL_TOKENS },
     createdAt: Date.now(),
     meta: {},
   };
@@ -139,7 +142,6 @@ export function upsertUser(user: UserUpdate) {
     }
   }
 
-  // TODO: Write firebase migration to backfill new fields
   if (updates.tokenCounts) {
     for (const family of MODEL_FAMILIES) {
       updates.tokenCounts[family] ??= 0;
@@ -150,6 +152,11 @@ export function upsertUser(user: UserUpdate) {
       updates.tokenLimits[family] ??= 0;
     }
   }
+  if (updates.tokenRefresh) {
+    for (const family of MODEL_FAMILIES) {
+      updates.tokenRefresh[family] ??= 0;
+    }
+  }
 
   users.set(user.token, Object.assign(existing, updates));
   usersToFlush.add(user.token);
@@ -245,19 +252,30 @@ export function hasAvailableQuota({
   return tokensConsumed < tokenLimit;
 }
 
+/**
+ * For the given user, sets token limits for each model family to the sum of the
+ * current count and the refresh amount, up to the default limit. If a quota is
+ * not specified for a model family, it is not touched.
+ */
 export function refreshQuota(token: string) {
   const user = users.get(token);
   if (!user) return;
-  const { tokenCounts, tokenLimits } = user;
-  const quotas = Object.entries(config.tokenQuota) as [ModelFamily, number][];
-  quotas
-    // If a quota is not configured, don't touch any existing limits a user may
-    // already have been assigned manually.
-    .filter(([, quota]) => quota > 0)
-    .forEach(
-      ([model, quota]) =>
-        (tokenLimits[model] = (tokenCounts[model] ?? 0) + quota)
-    );
+  const { tokenQuota } = config;
+  const { tokenCounts, tokenLimits, tokenRefresh } = user;
+
+  // Get default quotas for each model family.
+  const defaultQuotas = Object.entries(tokenQuota) as [ModelFamily, number][];
+  // If any user-specific refresh quotas are present, override default quotas.
+  const userQuotas = defaultQuotas.map(([f, q]) => [
+    f,
+    (tokenRefresh[f] ?? 0) || q,
+  ] as const /* narrow to tuple */);
+
+  userQuotas
+    // Ignore families with no global or user-specific refresh quota.
+    .filter(([, q]) => q > 0)
+    // Increase family token limit by the family's refresh amount.
+    .forEach(([f, q]) => (tokenLimits[f] = (tokenCounts[f] ?? 0) + q));
   usersToFlush.add(token);
 }
 
@@ -307,7 +325,7 @@ function cleanupExpiredTokens() {
       user.meta.refreshable = config.captchaMode !== "none";
       disabled++;
     }
-    const purgeTimeout =  config.powTokenPurgeHours * 60 * 60 * 1000;
+    const purgeTimeout = config.powTokenPurgeHours * 60 * 60 * 1000;
     if (user.disabledAt && user.disabledAt + purgeTimeout < now) {
       users.delete(user.token);
       usersToFlush.add(user.token);

src/shared/views/partials/shared_header.ejs

@@ -33,7 +33,7 @@
       .pagination li a {
         display: block;
         padding: 0.5em 1em;
-          border-bottom: none;
+        border-bottom: none;
         text-decoration: none;
       }
       .pagination li.active a {
@@ -71,20 +71,24 @@
       td.actions:hover {
         background-color: #e0e6f6;
       }
+      tr > td,
+      tr > th {
+          border-right: 1px solid #dedede;
+      }
 
       @media (max-width: 800px) {
         body {
           padding: 0.5em;
         }
 
-          table.full-width {
-            width: 100%;
-            position: static;
-            left: auto;
-            right: auto;
-            margin-left: 0;
-            margin-right: 0;
-          }
+        table.full-width {
+          width: 100%;
+          position: static;
+          left: auto;
+          right: auto;
+          margin-left: 0;
+          margin-right: 0;
+        }
       }
 
       @media (prefers-color-scheme: dark) {
@@ -95,6 +99,13 @@
         th.active {
           background-color: #446;
         }
+        td.actions:hover {
+          background-color: #446;
+        }
+        tr > td,
+        tr > th {
+          border-right: 1px solid #444;
+        }
       }
     </style>
   </head>

src/shared/views/partials/shared_quota-info.ejs

@@ -1,4 +1,6 @@
-<p>Next refresh: <time><%- nextQuotaRefresh %></time></p>
+<p>
+  Next refresh: <time><%- nextQuotaRefresh %></time>
+</p>
 <table class="striped">
   <thead>
     <tr>
@@ -9,7 +11,7 @@
       <% } %>
       <th scope="col">Limit</th>
       <th scope="col">Remaining</th>
-      <th scope="col">Refresh Amount</th>
+      <th scope="col" colspan="<%= showRefreshEdit ? 2 : 1 %>">Refresh Amount</th>
     </tr>
   </thead>
   <tbody>
@@ -19,7 +21,7 @@
       <td><%- prettyTokens(user.tokenCounts[key]) %></td>
       <% if (showTokenCosts) { %>
       <td>$<%- tokenCost(key, user.tokenCounts[key]).toFixed(2) %></td>
-      <% } %>      
+      <% } %>
       <% if (!user.tokenLimits[key]) { %>
       <td colspan="2" style="text-align: center">unlimited</td>
       <% } else { %>
@@ -29,7 +31,20 @@
       <% if (user.type === "temporary") { %>
       <td>N/A</td>
       <% } else { %>
-      <td><%- prettyTokens(quota[key]) %></td>
+      <td><%- prettyTokens(user.tokenRefresh[key] || quota[key]) %></td>
+      <% } %>
+      <% if (showRefreshEdit) { %>
+      <td class="actions">
+        <a
+          title="Edit"
+          id="edit-refresh"
+          href="#"
+          data-field="tokenRefresh_<%= key %>"
+          data-token="<%= user.token %>"
+          data-modelFamily="<%= key %>"
+          >✏️</a
+        >
+      </td>
       <% } %>
     </tr>
     <% }) %>