🪙 feat: Sync Balance Config on Login (danny-avila#6671)

* chore: Add deprecation warnings for environment variables in checks

* chore: Change deprecatedVariables to a const declaration in checks.js

* fix: Add date validation in checkBalanceRecord to prevent invalid date errors

* feat: Add setBalanceConfig middleware to synchronize user balance settings

* chore: Reorder middleware imports in oauth.js for better readability

Author: danny-avila

api/models/balanceMethods.js

@@ -5,6 +5,10 @@ const { getMultiplier } = require('./tx');
 const { logger } = require('~/config');
 const Balance = require('./Balance');
 
+function isInvalidDate(date) {
+  return isNaN(date);
+}
+
 /**
  * Simple check method that calculates token cost and returns balance info.
  * The auto-refill logic has been moved to balanceMethods.js to prevent circular dependencies.
@@ -48,13 +52,12 @@ const checkBalanceRecord = async function ({
   // Only perform auto-refill if spending would bring the balance to 0 or below
   if (balance - tokenCost <= 0 && record.autoRefillEnabled && record.refillAmount > 0) {
     const lastRefillDate = new Date(record.lastRefill);
-    const nextRefillDate = addIntervalToDate(
-      lastRefillDate,
-      record.refillIntervalValue,
-      record.refillIntervalUnit,
-    );
     const now = new Date();
-    if (now >= nextRefillDate) {
+    if (
+      isInvalidDate(lastRefillDate) ||
+      now >=
+        addIntervalToDate(lastRefillDate, record.refillIntervalValue, record.refillIntervalUnit)
+    ) {
       try {
         /** @type {{ rate: number, user: string, balance: number, transaction: import('@librechat/data-schemas').ITransaction}} */
         const result = await Transaction.createAutoRefillTransaction({

api/server/middleware/index.js

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

api/server/middleware/setBalanceConfig.js

@@ -0,0 +1,91 @@
+const { getBalanceConfig } = require('~/server/services/Config');
+const Balance = require('~/models/Balance');
+const { logger } = require('~/config');
+
+/**
+ * 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

@@ -23,6 +23,7 @@ const {
   checkInviteUser,
   registerLimiter,
   requireLdapAuth,
+  setBalanceConfig,
   requireLocalAuth,
   resetPasswordLimiter,
   validateRegistration,
@@ -40,6 +41,7 @@ router.post(
   loginLimiter,
   checkBan,
   ldapAuth ? requireLdapAuth : requireLocalAuth,
+  setBalanceConfig,
   loginController,
 );
 router.post('/refresh', refreshController);

api/server/routes/oauth.js

@@ -1,7 +1,13 @@
 // file deepcode ignore NoRateLimitingForLogin: Rate limiting is handled by the `loginLimiter` middleware
 const express = require('express');
 const passport = require('passport');
-const { loginLimiter, logHeaders, checkBan, checkDomainAllowed } = require('~/server/middleware');
+const {
+  checkBan,
+  logHeaders,
+  loginLimiter,
+  setBalanceConfig,
+  checkDomainAllowed,
+} = require('~/server/middleware');
 const { setAuthTokens } = require('~/server/services/AuthService');
 const { logger } = require('~/config');
 
@@ -56,6 +62,7 @@ router.get(
     session: false,
     scope: ['openid', 'profile', 'email'],
   }),
+  setBalanceConfig,
   oauthHandler,
 );
 
@@ -80,6 +87,7 @@ router.get(
     scope: ['public_profile'],
     profileFields: ['id', 'email', 'name'],
   }),
+  setBalanceConfig,
   oauthHandler,
 );
 
@@ -100,6 +108,7 @@ router.get(
     failureMessage: true,
     session: false,
   }),
+  setBalanceConfig,
   oauthHandler,
 );
 
@@ -122,6 +131,7 @@ router.get(
     session: false,
     scope: ['user:email', 'read:user'],
   }),
+  setBalanceConfig,
   oauthHandler,
 );
 
@@ -144,6 +154,7 @@ router.get(
     session: false,
     scope: ['identify', 'email'],
   }),
+  setBalanceConfig,
   oauthHandler,
 );
 
@@ -164,6 +175,7 @@ router.post(
     failureMessage: true,
     session: false,
   }),
+  setBalanceConfig,
   oauthHandler,
 );
 

api/server/services/start/checks.js

@@ -13,6 +13,24 @@ const secretDefaults = {
   JWT_REFRESH_SECRET: 'eaa5191f2914e30b9387fd84e254e4ba6fc51b4654968a9b0803b456a54b8418',
 };
 
+const deprecatedVariables = [
+  {
+    key: 'CHECK_BALANCE',
+    description:
+      'Please use the `balance` field in the `librechat.yaml` config file instead.\nMore info: https://librechat.ai/docs/configuration/librechat_yaml/object_structure/balance#overview',
+  },
+  {
+    key: 'START_BALANCE',
+    description:
+      'Please use the `balance` field in the `librechat.yaml` config file instead.\nMore info: https://librechat.ai/docs/configuration/librechat_yaml/object_structure/balance#overview',
+  },
+  {
+    key: 'GOOGLE_API_KEY',
+    description:
+      'Please use the `GOOGLE_SEARCH_API_KEY` environment variable for the Google Search Tool instead.',
+  },
+];
+
 /**
  * Checks environment variables for default secrets and deprecated variables.
  * Logs warnings for any default secret values being used and for usage of deprecated `GOOGLE_API_KEY`.
@@ -37,19 +55,11 @@ function checkVariables() {
     \u200B`);
   }
 
-  if (process.env.GOOGLE_API_KEY) {
-    logger.warn(
-      'The `GOOGLE_API_KEY` environment variable is deprecated.\nPlease use the `GOOGLE_SEARCH_API_KEY` environment variable instead.',
-    );
-  }
-
-  if (process.env.OPENROUTER_API_KEY) {
-    logger.warn(
-      `The \`OPENROUTER_API_KEY\` environment variable is deprecated and its functionality will be removed soon.
-      Use of this environment variable is highly discouraged as it can lead to unexpected errors when using custom endpoints.
-      Please use the config (\`librechat.yaml\`) file for setting up OpenRouter, and use \`OPENROUTER_KEY\` or another environment variable instead.`,
-    );
-  }
+  deprecatedVariables.forEach(({ key, description }) => {
+    if (process.env[key]) {
+      logger.warn(`The \`${key}\` environment variable is deprecated. ${description}`);
+    }
+  });
 
   checkPasswordReset();
 }

api/typedefs.js

@@ -771,6 +771,11 @@
  * @typedef {import('@librechat/data-schemas').IMongoFile} MongoFile
  * @memberof typedefs
  */
+/**
+ * @exports IBalance
+ * @typedef {import('@librechat/data-schemas').IBalance} IBalance
+ * @memberof typedefs
+ */
 
 /**
  * @exports MongoUser