admin ui improvements; adds Force Recheck feature

src/admin/common.ts

@@ -67,12 +67,32 @@ export const UserSchemaWithToken = UserSchema.extend({
   token: z.string(),
 }).strict();
 
-export const injectLocals: RequestHandler = (_req, res, next) => {
+export const injectLocals: RequestHandler = (req, res, next) => {
   const quota = config.tokenQuota;
   res.locals.quotasEnabled =
     quota.turbo > 0 || quota.gpt4 > 0 || quota.claude > 0;
 
   res.locals.persistenceEnabled = config.gatekeeperStore !== "memory";
 
+  if (req.query.flash) {
+    const content = String(req.query.flash)
+      .replace(/</g, "&lt;")
+      .replace(/>/g, "&gt;");
+    const match = content.match(/^([a-z]+):(.*)/);
+    if (match) {
+      res.locals.flash = { type: match[1], message: match[2] };
+    } else {
+      res.locals.flash = { type: "error", message: content };
+    }
+  } else {
+    res.locals.flash = null;
+  }
+
   next();
 };
+
+export class HttpError extends Error {
+  constructor(public status: number, message: string) {
+    super(message);
+  }
+}

src/admin/login.ts

@@ -3,7 +3,11 @@ import { Router } from "express";
 const loginRouter = Router();
 
 loginRouter.get("/login", (req, res) => {
-  res.render("admin/login", { failed: req.query.failed });
+  res.render("admin/login", {
+    flash: req.query.failed
+      ? { type: "error", message: "Invalid admin key" }
+      : null,
+  });
 });
 
 loginRouter.post("/login", (req, res) => {

src/admin/routes.ts

@@ -1,7 +1,7 @@
 import express, { Router } from "express";
 import cookieParser from "cookie-parser";
 import { authorize } from "./auth";
-import { injectLocals } from "./common";
+import { HttpError, injectLocals } from "./common";
 import { injectCsrfToken, checkCsrfToken } from "./csrf";
 import { loginRouter } from "./login";
 import { usersApiRouter as apiRouter } from "./api/users";
@@ -23,4 +23,26 @@ adminRouter.use(injectLocals);
 adminRouter.use("/", loginRouter);
 adminRouter.use("/manage", authorize({ via: "cookie" }), uiRouter);
 
+adminRouter.use(
+  (
+    err: Error,
+    _req: express.Request,
+    res: express.Response,
+    _next: express.NextFunction
+  ) => {
+    const data: any = { message: err.message, stack: err.stack };
+    if (err instanceof HttpError) {
+      data.status = err.status;
+      return res.status(err.status).render("admin/error", data);
+    } else if (err.name === "ForbiddenError") {
+      data.status = 403;
+      if (err.message === "invalid csrf token") {
+        data.message = "Invalid CSRF token; try refreshing the previous page before submitting again.";
+      }
+      return res.status(403).render("admin/error", { ...data, flash: null });
+    }
+    res.status(500).json({ error: data });
+  }
+);
+
 export { adminRouter };

src/admin/ui/users.ts

@@ -9,7 +9,9 @@ import {
   sortBy,
   paginate,
   UserSchema,
+  HttpError,
 } from "../common";
+import { keyPool } from "../../key-management";
 
 const router = Router();
 
@@ -42,8 +44,13 @@ router.post("/create-user", (_req, res) => {
 
 router.get("/view-user/:token", (req, res) => {
   const user = userStore.getUser(req.params.token);
-  if (!user) {
-    return res.status(404).send("User not found");
+  if (!user) throw new HttpError(404, "User not found");
+
+  if (req.query.refreshed) {
+    res.locals.flash = {
+      type: "success",
+      message: "User's quota was refreshed",
+    };
   }
   res.render("admin/view-user", { user });
 });
@@ -71,22 +78,21 @@ router.get("/list-users", (req, res) => {
   });
 });
 
-router.get("/import-users", (req, res) => {
-  const imported = Number(req.query.imported) || 0;
-  res.render("admin/import-users", { imported });
+router.get("/import-users", (_req, res) => {
+  res.render("admin/import-users");
 });
 
 router.post("/import-users", upload.single("users"), (req, res) => {
-  if (!req.file) {
-    return res.status(400).json({ error: "No file uploaded" });
-  }
+  if (!req.file) throw new HttpError(400, "No file uploaded");
+
   const data = JSON.parse(req.file.buffer.toString());
   const result = z.array(UserSchemaWithToken).safeParse(data.users);
-  if (!result.success) {
-    return res.status(400).json({ error: result.error });
-  }
+  if (!result.success) throw new HttpError(400, result.error.toString());
+
   const upserts = result.data.map((user) => userStore.upsertUser(user));
-  res.redirect(`/admin/manage/import-users?imported=${upserts.length}`);
+  res.render("admin/import-users", {
+    flash: { type: "success", message: `${upserts.length} users imported` },
+  });
 });
 
 router.get("/export-users", (_req, res) => {
@@ -106,18 +112,16 @@ router.get("/", (_req, res) => {
 
 router.post("/edit-user/:token", (req, res) => {
   const result = UserSchema.safeParse(req.body);
-  if (!result.success) {
-    return res.status(400).send(result.error);
-  }
+  if (!result.success) throw new HttpError(400, result.error.toString());
+
   userStore.upsertUser({ ...result.data, token: req.params.token });
   return res.sendStatus(204);
 });
 
 router.post("/reactivate-user/:token", (req, res) => {
   const user = userStore.getUser(req.params.token);
-  if (!user) {
-    return res.status(404).send("User not found");
-  }
+  if (!user) throw new HttpError(404, "User not found");
+
   userStore.upsertUser({
     token: user.token,
     disabledAt: 0,
@@ -128,28 +132,44 @@ router.post("/reactivate-user/:token", (req, res) => {
 
 router.post("/disable-user/:token", (req, res) => {
   const user = userStore.getUser(req.params.token);
-  if (!user) {
-    return res.status(404).send("User not found");
-  }
+  if (!user) throw new HttpError(404, "User not found");
+
   userStore.disableUser(req.params.token, req.body.reason);
   return res.sendStatus(204);
 });
 
 router.post("/refresh-user-quota", (req, res) => {
   const user = userStore.getUser(req.body.token);
-  if (!user) {
-    return res.status(404).send("User not found");
-  }
+  if (!user) throw new HttpError(404, "User not found");
+
   userStore.refreshQuota(req.body.token);
-  return res.redirect(`/admin/manage/view-user/${req.body.token}`);
+  return res.redirect(`/admin/manage/view-user/${req.body.token}?refreshed=1`);
 });
 
-router.post("/refresh-all-quotas", (_req, res) => {
-  const users = userStore.getUsers();
-
-  users.forEach((user) => userStore.refreshQuota(user.token));
-
-  return res.send(`Refreshed ${users.length} quotas`);
+router.post("/maintenance", (req, res) => {
+  const action = req.body.action;
+  let message = "";
+  switch (action) {
+    case "recheck": {
+      keyPool.recheck("openai");
+      const size = keyPool
+        .list()
+        .filter((key) => key.service === "openai").length;
+      message = `success: Scheduled recheck of ${size} OpenAI keys.`;
+      break;
+    }
+    case "resetQuotas": {
+      const users = userStore.getUsers();
+      users.forEach((user) => userStore.refreshQuota(user.token));
+      const { claude, gpt4, turbo } = config.tokenQuota;
+      message = `success: All users' token quotas reset to ${turbo} (Turbo), ${gpt4} (GPT-4), ${claude} (Claude).`;
+      break;
+    }
+    default: {
+      throw new HttpError(400, "Invalid action");
+    }
+  }
+  return res.redirect(`/admin/manage?flash=${message}`);
 });
 
 export { router as usersUiRouter };

src/key-management/anthropic/provider.ts

@@ -201,7 +201,7 @@ export class AnthropicKeyProvider implements KeyProvider<AnthropicKey> {
     key.rateLimitedUntil = now + RATE_LIMIT_LOCKOUT;
   }
 
-  public activeLimitInUsd() {
-    return "∞";
+  public recheck() {
+    throw new Error("Method not implemented.");
   }
 }

src/key-management/index.ts

@@ -52,8 +52,8 @@ export interface KeyProvider<T extends Key = Key> {
   anyUnchecked(): boolean;
   incrementPrompt(hash: string): void;
   getLockoutPeriod(model: Model): number;
-  activeLimitInUsd(options?: Record<string, unknown>): string;
   markRateLimited(hash: string): void;
+  recheck(service: AIService): void;
 }
 
 export const keyPool = new KeyPool();

src/key-management/key-pool.ts

@@ -81,11 +81,13 @@ export class KeyPool {
     }
   }
 
-  public activeLimitInUsd(
-    service: AIService,
-    options?: Record<string, unknown>
-  ): string {
-    return this.getKeyProvider(service).activeLimitInUsd(options);
+  public recheck(service: AIService): void {
+    const provider = this.getKeyProvider(service);
+    if (provider instanceof OpenAIKeyProvider) {
+      provider.recheck();
+    } else {
+      throw new Error(`Recheck not implemented for service '${service}'`);
+    }
   }
 
   private getService(model: Model): AIService {

src/key-management/openai/checker.ts

@@ -61,8 +61,9 @@ export class OpenAIKeyChecker {
    * it will schedule a check for the least recently checked key, respecting
    * the minimum check interval.
    **/
-  private scheduleNextCheck() {
+  public scheduleNextCheck() {
     const enabledKeys = this.keys.filter((key) => !key.isDisabled);
+    clearTimeout(this.timeout);
 
     if (enabledKeys.length === 0) {
       this.log.warn("All keys are disabled. Key checker stopping.");

src/key-management/openai/provider.ts

@@ -333,16 +333,9 @@ export class OpenAIKeyProvider implements KeyProvider<OpenAIKey> {
     }
   }
 
-  /**
-   * Returns the total quota limit of all keys in USD. Keys which are disabled
-   * are not included in the total.
-   */
-  public activeLimitInUsd(
-    { gpt4 }: { gpt4: boolean } = { gpt4: false }
-  ): string {
-    const keys = this.keys.filter((k) => !k.isDisabled && k.isGpt4 === gpt4);
-    const totalLimit = keys.reduce((acc, { hardLimit }) => acc + hardLimit, 0);
-    return `$${totalLimit.toFixed(2)}`;
+  public recheck() {
+    this.keys.forEach((key) => (key.lastChecked = 0));
+    this.checker?.scheduleNextCheck();
   }
 
   /** Writes key status to disk. */

src/views/_partials/admin-header.ejs

@@ -59,3 +59,15 @@
     </style>
   </head>
   <body style="font-family: sans-serif; background-color: #f0f0f0; padding: 1em;">
+      <% if (flash && flash.type === "error") { %>
+        <p style="color: red; background-color: #eedddd; padding: 1em">
+          <strong>⚠️ Error:</strong> <%= flash.message %>
+        </p>
+      <% } %>
+      <% if (flash && flash.type === "success") { %>
+        <p style="color: green; background-color: #ddffee; padding: 1em">
+          <strong>✅ Success:</strong> <%= flash.message %>
+        </p>
+      <% } %>
+      
+    

src/views/admin/error.ejs

@@ -0,0 +1,8 @@
+<%- include("../_partials/admin-header", { title: "Error" }) %>
+    <div style="color: red; background-color: #eedddd; padding: 1em">
+      <p><strong>⚠️ Error <%= status %>:</strong> <%= message %></p>
+      <pre><%= stack %></pre>
+      <a href="#" onclick="window.history.back()">Go Back</a> | <a href="/admin">Go Home</a>
+    </div>
+  </body>
+</html>

src/views/admin/import-users.ejs

@@ -16,8 +16,14 @@
     prompt
   </li>
   <li>
-    <code>tokenCount</code> (optional): the number of tokens the user has
-    consumed (not yet implemented)
+    <code>tokenCounts</code> (optional): the number of tokens the user has
+    consumed. This should be an object with keys <code>turbo</code>,
+    <code>gpt4</code>, and <code>claude</code>.
+  </li>
+  <li>
+    <code>tokenLimits</code> (optional): the number of tokens the user can
+    consume. This should be an object with keys <code>turbo</code>,
+    <code>gpt4</code>, and <code>claude</code>.
   </li>
   <li>
     <code>createdAt</code> (optional): the timestamp when the user was created
@@ -38,7 +44,4 @@
   <input type="submit" value="Import" />
 </form>
 </form>
-<% if (imported > 0) { %>
-<p>Imported <code><%= imported %></code> users.</p>
-<% } %>
 <%- include("../_partials/admin-footer") %>

src/views/admin/index.ejs

@@ -1,20 +1,51 @@
-<%- include("../_partials/admin-header", { title: "OAI Reverse Proxy Admin" }) %>
+<%- include("../_partials/admin-header", { title: "OAI Reverse Proxy Admin" })
+%>
 <h1>OAI Reverse Proxy Admin</h1>
 <% if (!persistenceEnabled) { %>
 <p style="color: red; background-color: #eedddd; padding: 1em">
-  <strong>⚠️ Users will be lost when the server restarts because persistence is
-  not configured.</strong><br />
+  <strong
+    >⚠️ Users will be lost when the server restarts because persistence is not
+    configured.</strong
+  ><br />
   <br />Be sure to export your users and import them again after restarting the
   server if you want to keep them.<br />
-  <br /> See the <a target="_blank"
-  href="https://gitgud.io/khanon/oai-reverse-proxy/-/blob/main/docs/user-management.md#firebase-realtime-database">
-  user management documentation</a> to learn how to set up persistence.
+  <br />
+  See the
+  <a
+    target="_blank"
+    href="https://gitgud.io/khanon/oai-reverse-proxy/-/blob/main/docs/user-management.md#firebase-realtime-database"
+  >
+    user management documentation</a
+  >
+  to learn how to set up persistence.
 </p>
 <% } %>
+<h3>Users</h3>
 <ul>
   <li><a href="/admin/manage/list-users">List Users</a></li>
   <li><a href="/admin/manage/create-user">Create User</a></li>
   <li><a href="/admin/manage/import-users">Import Users</a></li>
   <li><a href="/admin/manage/export-users">Export Users</a></li>
 </ul>
+<h3>Maintenance</h3>
+<form id="maintenanceForm" action="/admin/manage/maintenance" method="post">
+  <input id="_csrf" type="hidden" name="_csrf" value="<%= csrfToken %>" />
+  <input id="hiddenAction" type="hidden" name="action" value="" />
+  <button type="button" onclick="submitForm('recheck')">
+    Force OpenAI Key Recheck
+  </button>
+  <% if (quotasEnabled) { %>
+  <button type="button" onclick="submitForm('resetQuotas')">
+    Reset All Users' Quotas
+  </button>
+  <% } %>
+</form>
+
+<script>
+  function submitForm(action) {
+    document.getElementById("hiddenAction").value = action;
+    document.getElementById("maintenanceForm").submit();
+  }
+</script>
+
 <%- include("../_partials/admin-footer") %>

src/views/admin/list-users.ejs

@@ -41,13 +41,6 @@
         </tr>
         <% }); %>
     </table>
-    
-    <% if (quotasEnabled) { %>
-      <form action="/admin/manage/refresh-all-quotas" method="POST">
-        <input type="hidden" name="_csrf" value="<%- csrfToken %>" />
-        <button type="submit" class="btn btn-primary">Refresh All Quotas</button>
-      </form>
-    <% } %>
 
     <ul class="pagination">
       <% if (page > 1) { %>

src/views/admin/login.ejs

@@ -1,8 +1,5 @@
 <%- include("../_partials/admin-header", { title: "Login" }) %>
     <h1>Login</h1>
-    <% if (failed) { %>
-    <p style="color: red;">Please try again.</p>
-    <% } %>
     <form action="/admin/login" method="post">
       <input type="hidden" name="_csrf" value="<%= csrfToken %>" />
       <label for="token">Admin Key</label>