🔄 fix: Ensure lastRefill Date for Existing Users & Refactor Balance Middleware (#9086)

- Deleted setBalanceConfig middleware and its associated file.
- Introduced createSetBalanceConfig factory function to create middleware for synchronizing user balance settings.
- Updated auth and oauth routes to use the new balance configuration middleware.
- Added comprehensive tests for the new balance middleware functionality.
- Updated package versions and dependencies in package.json and package-lock.json.
- Added balance types and updated middleware index to export new balance middleware.

api/server/middleware/index.js

@@ -9,7 +9,6 @@ const validateEndpoint = require('./validateEndpoint');
 const requireLocalAuth = require('./requireLocalAuth');
 const canDeleteAccount = require('./canDeleteAccount');
 const accessResources = require('./accessResources');
-const setBalanceConfig = require('./setBalanceConfig');
 const requireLdapAuth = require('./requireLdapAuth');
 const abortMiddleware = require('./abortMiddleware');
 const checkInviteUser = require('./checkInviteUser');
@@ -44,7 +43,6 @@ module.exports = {
   requireLocalAuth,
   canDeleteAccount,
   validateEndpoint,
-  setBalanceConfig,
   concurrentLimiter,
   checkDomainAllowed,
   validateMessageReq,

api/server/middleware/setBalanceConfig.js

@@ -1,91 +0,0 @@
-const { logger } = require('@librechat/data-schemas');
-const { getBalanceConfig } = require('~/server/services/Config');
-const { Balance } = require('~/db/models');
-
-/**
- * Middleware to synchronize user balance settings with current balance configuration.
- * @function
- * @param {Object} req - Express request object containing user information.
- * @param {Object} res - Express response object.
- * @param {import('express').NextFunction} next - Next middleware function.
- */
-const setBalanceConfig = async (req, res, next) => {
-  try {
-    const balanceConfig = await getBalanceConfig();
-    if (!balanceConfig?.enabled) {
-      return next();
-    }
-    if (balanceConfig.startBalance == null) {
-      return next();
-    }
-
-    const userId = req.user._id;
-    const userBalanceRecord = await Balance.findOne({ user: userId }).lean();
-    const updateFields = buildUpdateFields(balanceConfig, userBalanceRecord);
-
-    if (Object.keys(updateFields).length === 0) {
-      return next();
-    }
-
-    await Balance.findOneAndUpdate(
-      { user: userId },
-      { $set: updateFields },
-      { upsert: true, new: true },
-    );
-
-    next();
-  } catch (error) {
-    logger.error('Error setting user balance:', error);
-    next(error);
-  }
-};
-
-/**
- * Build an object containing fields that need updating
- * @param {Object} config - The balance configuration
- * @param {Object|null} userRecord - The user's current balance record, if any
- * @returns {Object} Fields that need updating
- */
-function buildUpdateFields(config, userRecord) {
-  const updateFields = {};
-
-  // Ensure user record has the required fields
-  if (!userRecord) {
-    updateFields.user = userRecord?.user;
-    updateFields.tokenCredits = config.startBalance;
-  }
-
-  if (userRecord?.tokenCredits == null && config.startBalance != null) {
-    updateFields.tokenCredits = config.startBalance;
-  }
-
-  const isAutoRefillConfigValid =
-    config.autoRefillEnabled &&
-    config.refillIntervalValue != null &&
-    config.refillIntervalUnit != null &&
-    config.refillAmount != null;
-
-  if (!isAutoRefillConfigValid) {
-    return updateFields;
-  }
-
-  if (userRecord?.autoRefillEnabled !== config.autoRefillEnabled) {
-    updateFields.autoRefillEnabled = config.autoRefillEnabled;
-  }
-
-  if (userRecord?.refillIntervalValue !== config.refillIntervalValue) {
-    updateFields.refillIntervalValue = config.refillIntervalValue;
-  }
-
-  if (userRecord?.refillIntervalUnit !== config.refillIntervalUnit) {
-    updateFields.refillIntervalUnit = config.refillIntervalUnit;
-  }
-
-  if (userRecord?.refillAmount !== config.refillAmount) {
-    updateFields.refillAmount = config.refillAmount;
-  }
-
-  return updateFields;
-}
-
-module.exports = setBalanceConfig;

api/server/routes/auth.js

@@ -1,75 +1,75 @@
 const express = require('express');
+const { createSetBalanceConfig } = require('@librechat/api');
 const {
-  refreshController,
-  registrationController,
-  resetPasswordController,
   resetPasswordRequestController,
+  resetPasswordController,
+  registrationController,
   graphTokenController,
+  refreshController,
 } = require('~/server/controllers/AuthController');
-const { loginController } = require('~/server/controllers/auth/LoginController');
-const { logoutController } = require('~/server/controllers/auth/LogoutController');
-const { verify2FAWithTempToken } = require('~/server/controllers/auth/TwoFactorAuthController');
 const {
+  regenerateBackupCodes,
+  disable2FA,
+  confirm2FA,
   enable2FA,
   verify2FA,
-  disable2FA,
-  regenerateBackupCodes,
-  confirm2FA,
 } = require('~/server/controllers/TwoFactorController');
-const {
-  checkBan,
-  logHeaders,
-  loginLimiter,
-  requireJwtAuth,
-  checkInviteUser,
-  registerLimiter,
-  requireLdapAuth,
-  setBalanceConfig,
-  requireLocalAuth,
-  resetPasswordLimiter,
-  validateRegistration,
-  validatePasswordReset,
-} = require('~/server/middleware');
+const { verify2FAWithTempToken } = require('~/server/controllers/auth/TwoFactorAuthController');
+const { logoutController } = require('~/server/controllers/auth/LogoutController');
+const { loginController } = require('~/server/controllers/auth/LoginController');
+const { getBalanceConfig } = require('~/server/services/Config');
+const middleware = require('~/server/middleware');
+const { Balance } = require('~/db/models');
+
+const setBalanceConfig = createSetBalanceConfig({
+  getBalanceConfig,
+  Balance,
+});
 
 const router = express.Router();
 
 const ldapAuth = !!process.env.LDAP_URL && !!process.env.LDAP_USER_SEARCH_BASE;
 //Local
-router.post('/logout', requireJwtAuth, logoutController);
+router.post('/logout', middleware.requireJwtAuth, logoutController);
 router.post(
   '/login',
-  logHeaders,
-  loginLimiter,
-  checkBan,
-  ldapAuth ? requireLdapAuth : requireLocalAuth,
+  middleware.logHeaders,
+  middleware.loginLimiter,
+  middleware.checkBan,
+  ldapAuth ? middleware.requireLdapAuth : middleware.requireLocalAuth,
   setBalanceConfig,
   loginController,
 );
 router.post('/refresh', refreshController);
 router.post(
   '/register',
-  registerLimiter,
-  checkBan,
-  checkInviteUser,
-  validateRegistration,
+  middleware.registerLimiter,
+  middleware.checkBan,
+  middleware.checkInviteUser,
+  middleware.validateRegistration,
   registrationController,
 );
 router.post(
   '/requestPasswordReset',
-  resetPasswordLimiter,
-  checkBan,
-  validatePasswordReset,
+  middleware.resetPasswordLimiter,
+  middleware.checkBan,
+  middleware.validatePasswordReset,
   resetPasswordRequestController,
 );
-router.post('/resetPassword', checkBan, validatePasswordReset, resetPasswordController);
+router.post(
+  '/resetPassword',
+  middleware.checkBan,
+  middleware.validatePasswordReset,
+  resetPasswordController,
+);
 
-router.get('/2fa/enable', requireJwtAuth, enable2FA);
-router.post('/2fa/verify', requireJwtAuth, verify2FA);
-router.post('/2fa/verify-temp', checkBan, verify2FAWithTempToken);
-router.post('/2fa/confirm', requireJwtAuth, confirm2FA);
-router.post('/2fa/disable', requireJwtAuth, disable2FA);
-router.post('/2fa/backup/regenerate', requireJwtAuth, regenerateBackupCodes);
+router.get('/2fa/enable', middleware.requireJwtAuth, enable2FA);
+router.post('/2fa/verify', middleware.requireJwtAuth, verify2FA);
+router.post('/2fa/verify-temp', middleware.checkBan, verify2FAWithTempToken);
+router.post('/2fa/confirm', middleware.requireJwtAuth, confirm2FA);
+router.post('/2fa/disable', middleware.requireJwtAuth, disable2FA);
+router.post('/2fa/backup/regenerate', middleware.requireJwtAuth, regenerateBackupCodes);
 
-router.get('/graph-token', requireJwtAuth, graphTokenController);
+router.get('/graph-token', middleware.requireJwtAuth, graphTokenController);
 
 module.exports = router;

api/server/routes/oauth.js

@@ -1,19 +1,20 @@
 // file deepcode ignore NoRateLimitingForLogin: Rate limiting is handled by the `loginLimiter` middleware
 const express = require('express');
 const passport = require('passport');
-const { isEnabled } = require('@librechat/api');
 const { randomState } = require('openid-client');
 const { logger } = require('@librechat/data-schemas');
 const { ErrorTypes } = require('librechat-data-provider');
-const {
-  checkBan,
-  logHeaders,
-  loginLimiter,
-  setBalanceConfig,
-  checkDomainAllowed,
-} = require('~/server/middleware');
+const { isEnabled, createSetBalanceConfig } = require('@librechat/api');
+const { checkDomainAllowed, loginLimiter, logHeaders, checkBan } = require('~/server/middleware');
 const { syncUserEntraGroupMemberships } = require('~/server/services/PermissionService');
 const { setAuthTokens, setOpenIDAuthTokens } = require('~/server/services/AuthService');
+const { getBalanceConfig } = require('~/server/services/Config');
+const { Balance } = require('~/db/models');
+
+const setBalanceConfig = createSetBalanceConfig({
+  getBalanceConfig,
+  Balance,
+});
 
 const router = express.Router();