diff --git a/api/models/balanceMethods.js b/api/models/balanceMethods.js index e700cc96e7..4b788160aa 100644 --- a/api/models/balanceMethods.js +++ b/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({ diff --git a/api/server/middleware/index.js b/api/server/middleware/index.js index 789ec6a82d..6a41d6f157 100644 --- a/api/server/middleware/index.js +++ b/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, diff --git a/api/server/middleware/setBalanceConfig.js b/api/server/middleware/setBalanceConfig.js new file mode 100644 index 0000000000..98d3cf1145 --- /dev/null +++ b/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; diff --git a/api/server/routes/auth.js b/api/server/routes/auth.js index 2d9fae7ae7..187d908abd 100644 --- a/api/server/routes/auth.js +++ b/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); diff --git a/api/server/routes/oauth.js b/api/server/routes/oauth.js index 9ea896e30e..b2037683d2 100644 --- a/api/server/routes/oauth.js +++ b/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, ); diff --git a/api/server/services/start/checks.js b/api/server/services/start/checks.js index 100424d35a..fe9cd79edf 100644 --- a/api/server/services/start/checks.js +++ b/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(); } diff --git a/api/typedefs.js b/api/typedefs.js index 9d35c64580..a89866c07e 100644 --- a/api/typedefs.js +++ b/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