💵 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/server/controllers/assistants/chatV1.js

@@ -19,7 +19,7 @@ const {
   addThreadMetadata,
   saveAssistantMessage,
 } = require('~/server/services/Threads');
-const { sendResponse, sendMessage, sleep, isEnabled, countTokens } = require('~/server/utils');
+const { sendResponse, sendMessage, sleep, countTokens } = require('~/server/utils');
 const { runAssistant, createOnTextProgress } = require('~/server/services/AssistantService');
 const validateAuthor = require('~/server/middleware/assistants/validateAuthor');
 const { formatMessage, createVisionPrompt } = require('~/app/clients/prompts');
@@ -248,7 +248,8 @@ const chatV1 = async (req, res) => {
     }
 
     const checkBalanceBeforeRun = async () => {
-      if (!isEnabled(process.env.CHECK_BALANCE)) {
+      const balance = req.app?.locals?.balance;
+      if (!balance?.enabled) {
         return;
       }
       const transactions =

api/server/controllers/assistants/chatV2.js

@@ -18,11 +18,11 @@ const {
   saveAssistantMessage,
 } = require('~/server/services/Threads');
 const { runAssistant, createOnTextProgress } = require('~/server/services/AssistantService');
-const { sendMessage, sleep, isEnabled, countTokens } = require('~/server/utils');
 const { createErrorHandler } = require('~/server/controllers/assistants/errors');
 const validateAuthor = require('~/server/middleware/assistants/validateAuthor');
 const { createRun, StreamRunManager } = require('~/server/services/Runs');
 const { addTitle } = require('~/server/services/Endpoints/assistants');
+const { sendMessage, sleep, countTokens } = require('~/server/utils');
 const { createRunBody } = require('~/server/services/createRunBody');
 const { getTransactions } = require('~/models/Transaction');
 const checkBalance = require('~/models/checkBalance');
@@ -124,7 +124,8 @@ const chatV2 = async (req, res) => {
     }
 
     const checkBalanceBeforeRun = async () => {
-      if (!isEnabled(process.env.CHECK_BALANCE)) {
+      const balance = req.app?.locals?.balance;
+      if (!balance?.enabled) {
         return;
       }
       const transactions =

api/server/routes/config.js

@@ -69,7 +69,6 @@ router.get('/', async function (req, res) {
         !!process.env.EMAIL_PASSWORD &&
         !!process.env.EMAIL_FROM,
       passwordResetEnabled,
-      checkBalance: isEnabled(process.env.CHECK_BALANCE),
       showBirthdayIcon:
         isBirthday() ||
         isEnabled(process.env.SHOW_BIRTHDAY_ICON) ||
@@ -77,6 +76,7 @@ router.get('/', async function (req, res) {
       helpAndFaqURL: process.env.HELP_AND_FAQ_URL || 'https://librechat.ai',
       interface: req.app.locals.interfaceConfig,
       modelSpecs: req.app.locals.modelSpecs,
+      balance: req.app.locals.balance,
       sharedLinksEnabled,
       publicSharedLinksEnabled,
       analyticsGtmId: process.env.ANALYTICS_GTM_ID,

api/server/services/AppService.js

@@ -9,15 +9,16 @@ const { checkVariables, checkHealth, checkConfig, checkAzureVariables } = requir
 const { azureAssistantsDefaults, assistantsConfigSetup } = require('./start/assistants');
 const { initializeAzureBlobService } = require('./Files/Azure/initialize');
 const { initializeFirebase } = require('./Files/Firebase/initialize');
-const { initializeS3 } = require('./Files/S3/initialize');
 const loadCustomConfig = require('./Config/loadCustomConfig');
 const handleRateLimits = require('./Config/handleRateLimits');
 const { loadDefaultInterface } = require('./start/interface');
 const { azureConfigSetup } = require('./start/azureOpenAI');
 const { processModelSpecs } = require('./start/modelSpecs');
+const { initializeS3 } = require('./Files/S3/initialize');
 const { loadAndFormatTools } = require('./ToolService');
 const { agentsConfigSetup } = require('./start/agents');
 const { initializeRoles } = require('~/models/Role');
+const { isEnabled } = require('~/server/utils');
 const { getMCPManager } = require('~/config');
 const paths = require('~/config/paths');
 
@@ -29,7 +30,7 @@ const paths = require('~/config/paths');
  */
 const AppService = async (app) => {
   await initializeRoles();
-  /** @type {TCustomConfig}*/
+  /** @type {TCustomConfig} */
   const config = (await loadCustomConfig()) ?? {};
   const configDefaults = getConfigDefaults();
 
@@ -37,6 +38,11 @@ const AppService = async (app) => {
   const filteredTools = config.filteredTools;
   const includedTools = config.includedTools;
   const fileStrategy = config.fileStrategy ?? configDefaults.fileStrategy;
+  const startBalance = process.env.START_BALANCE;
+  const balance = config.balance ?? {
+    enabled: isEnabled(process.env.CHECK_BALANCE),
+    startBalance: startBalance ? parseInt(startBalance, 10) : undefined,
+  };
   const imageOutputType = config?.imageOutputType ?? configDefaults.imageOutputType;
 
   process.env.CDN_PROVIDER = fileStrategy;
@@ -52,7 +58,7 @@ const AppService = async (app) => {
     initializeS3();
   }
 
-  /** @type {Record<string, FunctionTool} */
+  /** @type {Record<string, FunctionTool>} */
   const availableTools = loadAndFormatTools({
     adminFilter: filteredTools,
     adminIncluded: includedTools,
@@ -79,6 +85,7 @@ const AppService = async (app) => {
     availableTools,
     imageOutputType,
     interfaceConfig,
+    balance,
   };
 
   if (!Object.keys(config).length) {

api/server/services/AppService.spec.js

@@ -15,6 +15,9 @@ jest.mock('./Config/loadCustomConfig', () => {
     Promise.resolve({
       registration: { socialLogins: ['testLogin'] },
       fileStrategy: 'testStrategy',
+      balance: {
+        enabled: true,
+      },
     }),
   );
 });
@@ -124,6 +127,9 @@ describe('AppService', () => {
       imageOutputType: expect.any(String),
       fileConfig: undefined,
       secureImageLinks: undefined,
+      balance: { enabled: true },
+      filteredTools: undefined,
+      includedTools: undefined,
     });
   });
 
@@ -341,9 +347,6 @@ describe('AppService', () => {
     process.env.FILE_UPLOAD_USER_MAX = 'initialUserMax';
     process.env.FILE_UPLOAD_USER_WINDOW = 'initialUserWindow';
 
-    // Mock a custom configuration without specific rate limits
-    require('./Config/loadCustomConfig').mockImplementationOnce(() => Promise.resolve({}));
-
     await AppService(app);
 
     // Verify that process.env falls back to the initial values
@@ -404,9 +407,6 @@ describe('AppService', () => {
     process.env.IMPORT_USER_MAX = 'initialUserMax';
     process.env.IMPORT_USER_WINDOW = 'initialUserWindow';
 
-    // Mock a custom configuration without specific rate limits
-    require('./Config/loadCustomConfig').mockImplementationOnce(() => Promise.resolve({}));
-
     await AppService(app);
 
     // Verify that process.env falls back to the initial values
@@ -445,13 +445,27 @@ describe('AppService updating app.locals and issuing warnings', () => {
     expect(app.locals.availableTools).toBeDefined();
     expect(app.locals.fileStrategy).toEqual(FileSources.local);
     expect(app.locals.socialLogins).toEqual(defaultSocialLogins);
+    expect(app.locals.balance).toEqual(
+      expect.objectContaining({
+        enabled: false,
+        startBalance: undefined,
+      }),
+    );
   });
 
   it('should update app.locals with values from loadCustomConfig', async () => {
-    // Mock loadCustomConfig to return a specific config object
+    // Mock loadCustomConfig to return a specific config object with a complete balance config
     const customConfig = {
       fileStrategy: 'firebase',
       registration: { socialLogins: ['testLogin'] },
+      balance: {
+        enabled: false,
+        startBalance: 5000,
+        autoRefillEnabled: true,
+        refillIntervalValue: 15,
+        refillIntervalUnit: 'hours',
+        refillAmount: 5000,
+      },
     };
     require('./Config/loadCustomConfig').mockImplementationOnce(() =>
       Promise.resolve(customConfig),
@@ -464,6 +478,7 @@ describe('AppService updating app.locals and issuing warnings', () => {
     expect(app.locals.availableTools).toBeDefined();
     expect(app.locals.fileStrategy).toEqual(customConfig.fileStrategy);
     expect(app.locals.socialLogins).toEqual(customConfig.registration.socialLogins);
+    expect(app.locals.balance).toEqual(customConfig.balance);
   });
 
   it('should apply the assistants endpoint configuration correctly to app.locals', async () => {

api/server/services/Config/getCustomConfig.js

@@ -1,5 +1,5 @@
 const { CacheKeys, EModelEndpoint } = require('librechat-data-provider');
-const { normalizeEndpointName } = require('~/server/utils');
+const { normalizeEndpointName, isEnabled } = require('~/server/utils');
 const loadCustomConfig = require('./loadCustomConfig');
 const getLogStores = require('~/cache/getLogStores');
 
@@ -23,6 +23,29 @@ async function getCustomConfig() {
   return customConfig;
 }
 
+/**
+ * Retrieves the configuration object
+ * @function getBalanceConfig
+ * @returns {Promise<TCustomConfig['balance'] | null>}
+ * */
+async function getBalanceConfig() {
+  const isLegacyEnabled = isEnabled(process.env.CHECK_BALANCE);
+  const startBalance = process.env.START_BALANCE;
+  if (isLegacyEnabled || (startBalance != null && startBalance)) {
+    /** @type {TCustomConfig['balance']} */
+    const config = {
+      enabled: isLegacyEnabled,
+      startBalance: startBalance ? parseInt(startBalance, 10) : undefined,
+    };
+    return config;
+  }
+  const customConfig = await getCustomConfig();
+  if (!customConfig) {
+    return null;
+  }
+  return customConfig?.['balance'] ?? null;
+}
+
 /**
  *
  * @param {string | EModelEndpoint} endpoint
@@ -40,4 +63,4 @@ const getCustomEndpointConfig = async (endpoint) => {
   );
 };
 
-module.exports = { getCustomConfig, getCustomEndpointConfig };
+module.exports = { getCustomConfig, getBalanceConfig, getCustomEndpointConfig };