💵 feat: Add Automatic Balance Refill (#6452)

* 🚀 feat: Add automatic refill settings to balance schema

* 🚀 feat: Refactor balance feature to use global interface configuration

* 🚀 feat: Implement auto-refill functionality for balance management

* 🚀 feat: Enhance auto-refill logic and configuration for balance management

* 🚀 chore: Bump version to 0.7.74 in package.json and package-lock.json

* 🚀 chore: Bump version to 0.0.5 in package.json and package-lock.json

* 🚀 docs: Update comment for balance settings in librechat.example.yaml

* chore: space in `.env.example`

* 🚀 feat: Implement balance configuration loading and refactor related components

* 🚀 test: Refactor tests to use custom config for balance feature

* 🚀 fix: Update balance response handling in Transaction.js to use Balance model

* 🚀 test: Update AppService tests to include balance configuration in mock setup

* 🚀 test: Enhance AppService tests with complete balance configuration scenarios

* 🚀 refactor: Rename balanceConfig to balance and update related tests for clarity

* 🚀 refactor: Remove loadDefaultBalance and update balance handling in AppService

* 🚀 test: Update AppService tests to reflect new balance structure and defaults

* 🚀 test: Mock getCustomConfig in BaseClient tests to control balance configuration

* 🚀 test: Add get method to mockCache in OpenAIClient tests for improved cache handling

* 🚀 test: Mock getCustomConfig in OpenAIClient tests to control balance configuration

* 🚀 test: Remove mock for getCustomConfig in OpenAIClient tests to streamline configuration handling

* 🚀 fix: Update balance configuration reference in config.js for consistency

* refactor: Add getBalanceConfig function to retrieve balance configuration

* chore: Comment out example balance settings in librechat.example.yaml

* refactor: Replace getCustomConfig with getBalanceConfig for balance handling

* fix: tests

* refactor: Replace getBalanceConfig call with balance from request locals

* refactor: Update balance handling to use environment variables for configuration

* refactor: Replace getBalanceConfig calls with balance from request locals

* refactor: Simplify balance configuration logic in getBalanceConfig

---------

Co-authored-by: Danny Avila <danny@librechat.ai>

api/models/Balance.js

@@ -3,6 +3,40 @@ const { balanceSchema } = require('@librechat/data-schemas');
 const { getMultiplier } = require('./tx');
 const { logger } = require('~/config');
 
+/**
+ * Adds a time interval to a given date.
+ * @param {Date} date - The starting date.
+ * @param {number} value - The numeric value of the interval.
+ * @param {'seconds'|'minutes'|'hours'|'days'|'weeks'|'months'} unit - The unit of time.
+ * @returns {Date} A new Date representing the starting date plus the interval.
+ */
+const addIntervalToDate = (date, value, unit) => {
+  const result = new Date(date);
+  switch (unit) {
+    case 'seconds':
+      result.setSeconds(result.getSeconds() + value);
+      break;
+    case 'minutes':
+      result.setMinutes(result.getMinutes() + value);
+      break;
+    case 'hours':
+      result.setHours(result.getHours() + value);
+      break;
+    case 'days':
+      result.setDate(result.getDate() + value);
+      break;
+    case 'weeks':
+      result.setDate(result.getDate() + value * 7);
+      break;
+    case 'months':
+      result.setMonth(result.getMonth() + value);
+      break;
+    default:
+      break;
+  }
+  return result;
+};
+
 balanceSchema.statics.check = async function ({
   user,
   model,
@@ -14,9 +48,20 @@ balanceSchema.statics.check = async function ({
 }) {
   const multiplier = getMultiplier({ valueKey, tokenType, model, endpoint, endpointTokenConfig });
   const tokenCost = amount * multiplier;
-  const { tokenCredits: balance } = (await this.findOne({ user }, 'tokenCredits').lean()) ?? {};
 
-  logger.debug('[Balance.check]', {
+  // Retrieve the complete balance record
+  let record = await this.findOne({ user }).lean();
+  if (!record) {
+    logger.debug('[Balance.check] No balance record found for user', { user });
+    return {
+      canSpend: false,
+      balance: 0,
+      tokenCost,
+    };
+  }
+  let balance = record.tokenCredits;
+
+  logger.debug('[Balance.check] Initial state', {
     user,
     model,
     endpoint,
@@ -28,15 +73,31 @@ balanceSchema.statics.check = async function ({
     endpointTokenConfig: !!endpointTokenConfig,
   });
 
-  if (!balance) {
-    return {
-      canSpend: false,
-      balance: 0,
-      tokenCost,
-    };
+  // 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) {
+      record = await this.findOneAndUpdate(
+        { user },
+        {
+          $inc: { tokenCredits: record.refillAmount },
+          $set: { lastRefill: new Date() },
+        },
+        { new: true },
+      ).lean();
+      balance = record.tokenCredits;
+      logger.debug('[Balance.check] Auto-refill performed', { balance });
+    }
   }
 
-  logger.debug('[Balance.check]', { tokenCost });
+  logger.debug('[Balance.check] Token cost', { tokenCost });
 
   return { canSpend: balance >= tokenCost, balance, tokenCost };
 };