Per-user token quotas and automatic quota refreshing (khanon/oai-reverse-proxy!37)

src/admin/common.ts

@@ -1,5 +1,7 @@
 import { z } from "zod";
+import { RequestHandler } from "express";
 import { Query } from "express-serve-static-core";
+import { config } from "../config";
 
 export function parseSort(sort: Query["sort"]) {
   if (!sort) return null;
@@ -45,7 +47,15 @@ export const UserSchema = z
     ip: z.array(z.string()).optional(),
     type: z.enum(["normal", "special"]).optional(),
     promptCount: z.number().optional(),
-    tokenCount: z.number().optional(),
+    tokenCount: z.any().optional(), // never used, but remains for compatibility
+    tokenCounts: z
+      .object({ turbo: z.number(), gpt4: z.number(), claude: z.number() })
+      .strict()
+      .optional(),
+    tokenLimits: z
+      .object({ turbo: z.number(), gpt4: z.number(), claude: z.number() })
+      .strict()
+      .optional(),
     createdAt: z.number().optional(),
     lastUsedAt: z.number().optional(),
     disabledAt: z.number().optional(),
@@ -56,3 +66,13 @@ export const UserSchema = z
 export const UserSchemaWithToken = UserSchema.extend({
   token: z.string(),
 }).strict();
+
+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";
+
+  next();
+};

src/admin/routes.ts

@@ -1,10 +1,11 @@
 import express, { Router } from "express";
 import cookieParser from "cookie-parser";
 import { authorize } from "./auth";
+import { injectLocals } from "./common";
 import { injectCsrfToken, checkCsrfToken } from "./csrf";
+import { loginRouter } from "./login";
 import { usersApiRouter as apiRouter } from "./api/users";
 import { usersUiRouter as uiRouter } from "./ui/users";
-import { loginRouter } from "./login";
 
 const adminRouter = Router();
 
@@ -18,6 +19,7 @@ adminRouter.use(injectCsrfToken);
 adminRouter.use("/users", authorize({ via: "header" }), apiRouter);
 
 adminRouter.use(checkCsrfToken); // All UI routes require CSRF token
+adminRouter.use(injectLocals);
 adminRouter.use("/", loginRouter);
 adminRouter.use("/manage", authorize({ via: "cookie" }), uiRouter);
 

src/admin/ui/users.ts

@@ -53,7 +53,13 @@ router.get("/list-users", (req, res) => {
   const requestedPageSize =
     Number(req.query.perPage) || Number(req.cookies.perPage) || 20;
   const perPage = Math.max(1, Math.min(1000, requestedPageSize));
-  const users = userStore.getUsers().sort(sortBy(sort, false));
+  const users = userStore
+    .getUsers()
+    .map((user) => {
+      const sum = Object.values(user.tokenCounts).reduce((a, b) => a + b, 0); // TODO: cache
+      return { ...user, sumTokenCounts: sum };
+    })
+    .sort(sortBy(sort, false));
 
   const page = Number(req.query.page) || 1;
   const { items, ...pagination } = paginate(users, page, perPage);
@@ -95,9 +101,7 @@ router.get("/export-users.json", (_req, res) => {
 });
 
 router.get("/", (_req, res) => {
-  res.render("admin/index", {
-    isPersistenceEnabled: config.gatekeeperStore !== "memory",
-  });
+  res.render("admin/index");
 });
 
 router.post("/edit-user/:token", (req, res) => {
@@ -129,7 +133,23 @@ router.post("/disable-user/:token", (req, res) => {
   }
   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");
+  }
+  userStore.refreshQuota(req.body.token);
+  return res.redirect(`/admin/manage/view-user/${req.body.token}`);
+});
+
+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`);
+});
 
 export { router as usersUiRouter };