From 2aa19e5b09e17af7b8ef003cd2072feeeebf8488 Mon Sep 17 00:00:00 2001
From: nai-degen <khoners@protonmail.com>
Date: Sat, 27 Jul 2024 14:25:53 -0500
Subject: [PATCH] adds user-specific overrides for daily quota refresh

---
 docs/pow-captcha.md                           |  2 +-
 src/admin/web/views/admin_index.ejs           |  2 +-
 src/admin/web/views/admin_view-user.ejs       | 27 ++++++++----
 src/shared/users/schema.ts                    |  2 +
 src/shared/users/user-store.ts                | 42 +++++++++++++------
 src/shared/views/partials/shared_header.ejs   | 29 +++++++++----
 .../views/partials/shared_quota-info.ejs      | 23 ++++++++--
 src/user/web/views/user_lookup.ejs            |  2 +-
 8 files changed, 92 insertions(+), 37 deletions(-)

diff --git a/docs/pow-captcha.md b/docs/pow-captcha.md
index 8fccdfd..0151c6b 100644
--- a/docs/pow-captcha.md
+++ b/docs/pow-captcha.md
@@ -129,7 +129,7 @@ also significantly reduce hash rates on mobile devices.
 - Intel Core i9-13900K (Chrome, in VM limited to 4 cores): 12.2 - 13.0 H/s
 - iPad Pro (M2) (Safari, 6 workers): 8.0 - 10 H/s
   - Thermal throttles early. 8 cores is normal concurrency, but unstable.
-- iPhone 13 Pro (Safari): 4.0 - 4.6 H/s
+- iPhone 15 Pro Max (Safari): 4.0 - 4.6 H/s
 - Samsung Galaxy S10e (Chrome): 3.6 - 3.8 H/s
   - This is a 2019 phone almost matching an iPhone five years newer because of
     bad Safari performance.
diff --git a/src/admin/web/views/admin_index.ejs b/src/admin/web/views/admin_index.ejs
index 339368f..591a9da 100644
--- a/src/admin/web/views/admin_index.ejs
+++ b/src/admin/web/views/admin_index.ejs
@@ -43,7 +43,7 @@
       <legend>Bulk Quota Management</legend>
       <p>
         <button id="refresh-quotas" type="button" onclick="submitForm('resetQuotas')">Refresh All Quotas</button>
-        Resets all users' quotas to the values set in the <code>TOKEN_QUOTA_*</code> environment variables.
+        Immediately refreshes all users' quotas by the configured amounts.
       </p>
       <p>
         <button id="clear-token-counts" type="button" onclick="submitForm('resetCounts')">Clear All Token Counts</button>
diff --git a/src/admin/web/views/admin_view-user.ejs b/src/admin/web/views/admin_view-user.ejs
index d060aec..4e34649 100644
--- a/src/admin/web/views/admin_view-user.ejs
+++ b/src/admin/web/views/admin_view-user.ejs
@@ -101,6 +101,10 @@
   <% ["nickname", "type", "disabledAt", "disabledReason", "maxIps", "adminNote"].forEach(function (key) { %>
   <input type="hidden" name="<%- key %>" value="<%- user[key] %>" />
   <% }); %>
+  <!-- tokenRefresh_ keys are dynamically generated -->
+  <% Object.entries(quota).forEach(([family]) => { %>
+  <input type="hidden" name="tokenRefresh_<%- family %>" value="<%- user.tokenRefresh[family] || quota[family] %>" />
+  <% }); %>
 </form>
 
 <h3>Quota Information</h3>
@@ -111,7 +115,7 @@
   <button type="submit" class="btn btn-primary">Refresh Quotas for User</button>
 </form>
 <% } %>
-<%- include("partials/shared_quota-info", { quota, user }) %>
+<%- include("partials/shared_quota-info", { quota, user, showRefreshEdit: true }) %>
 
 <p><a href="/admin/manage/list-users">Back to User List</a></p>
 
@@ -122,18 +126,25 @@
       const token = a.dataset.token;
       const field = a.dataset.field;
       const existingValue = document.querySelector(`#current-values input[name=${field}]`).value;
-      let value = prompt(`Enter new value for '${field}'':`, existingValue);
+
+      let value = prompt(`Enter new value for '${field}':`, existingValue);
       if (value !== null) {
         if (value === "") {
           value = null;
         }
+
+        const payload = { _csrf: document.querySelector("meta[name=csrf-token]").getAttribute("content") };
+        if (field.startsWith("tokenRefresh_")) {
+          const family = field.slice("tokenRefresh_".length);
+          payload.tokenRefresh = { [family]: Number(value) };
+        } else {
+          payload[field] = value;
+        }
+
         fetch(`/admin/manage/edit-user/${token}`, {
           method: "POST",
           credentials: "same-origin",
-          body: JSON.stringify({
-            [field]: value,
-            _csrf: document.querySelector("meta[name=csrf-token]").getAttribute("content"),
-          }),
+          body: JSON.stringify(payload),
           headers: { "Content-Type": "application/json", Accept: "application/json" },
         })
           .then((res) => Promise.all([res.ok, res.json()]))
@@ -141,9 +152,7 @@
             const url = new URL(window.location.href);
             const params = new URLSearchParams();
             if (!ok) {
-              params.set("flash", `error: ${json.error.message}`);
-            } else {
-              params.set("flash", `success: User's ${field} updated.`);
+              alert(`Failed to edit user: ${json.message}`);
             }
             url.search = params.toString();
             window.location.assign(url);
diff --git a/src/shared/users/schema.ts b/src/shared/users/schema.ts
index fc9ff32..b62936b 100644
--- a/src/shared/users/schema.ts
+++ b/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. */
diff --git a/src/shared/users/user-store.ts b/src/shared/users/user-store.ts
index bca8b46..026933f 100644
--- a/src/shared/users/user-store.ts
+++ b/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);
diff --git a/src/shared/views/partials/shared_header.ejs b/src/shared/views/partials/shared_header.ejs
index 93079f0..2e8f5fb 100644
--- a/src/shared/views/partials/shared_header.ejs
+++ b/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>
diff --git a/src/shared/views/partials/shared_quota-info.ejs b/src/shared/views/partials/shared_quota-info.ejs
index 5580a67..4686896 100644
--- a/src/shared/views/partials/shared_quota-info.ejs
+++ b/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>
     <% }) %>
diff --git a/src/user/web/views/user_lookup.ejs b/src/user/web/views/user_lookup.ejs
index 190d5fb..c0972ee 100644
--- a/src/user/web/views/user_lookup.ejs
+++ b/src/user/web/views/user_lookup.ejs
@@ -64,7 +64,7 @@
 </table>
 
 <h3>Quota Information</h3>
-<%- include("partials/shared_quota-info", { quota, user }) %>
+<%- include("partials/shared_quota-info", { quota, user, showRefreshEdit: false }) %>
 
 <form id="edit-nickname-form" style="display: none" action="/user/edit-nickname" method="post">
   <input type="hidden" name="_csrf" value="<%= csrfToken %>" />