Andreas/LibreChat
1
0
Fork 0
mirror of https://github.com/danny-avila/LibreChat.git synced 2025-12-17 08:50:15 +01:00
Code Issues Projects Releases Packages Wiki Activity Actions 7

🏗️ fix: Agents Token Spend Race Conditions, Add Auto-refill Tx, Add Relevant Tests (#6480)

Browse source 
* 🏗️ refactor: Improve spendTokens logic to handle zero completion tokens and enhance test coverage

* 🏗️ test: Add tests to ensure balance does not go below zero when spending tokens

* 🏗️ fix: Ensure proper continuation in AgentClient when handling errors

* fix: spend token race conditions

* 🏗️ test: Add test for handling multiple concurrent transactions with high balance

* fix: Handle Omni models prompt prefix handling for user messages with array content in OpenAIClient

* refactor: Update checkBalance import paths to use new balanceMethods module

* refactor: Update checkBalance imports and implement updateBalance function for atomic balance updates

* fix: import from replace method

* feat: Add createAutoRefillTransaction method to handle non-balance updating transactions

* refactor: Move auto-refill logic to balanceMethods and enhance checkBalance functionality

* feat: Implement logging for auto-refill transactions in balance checks

* refactor: Remove logRefill calls from multiple client and handler files

* refactor: Move balance checking and auto-refill logic to balanceMethods for improved structure

* refactor: Simplify balance check calls by removing unnecessary balanceRecord assignments

* fix: Prevent negative rawAmount in spendTokens when promptTokens is zero

* fix: Update balanceMethods to use Balance model for findOneAndUpdate

* chore: import order

* refactor: remove unused txMethods file to streamline codebase

* feat: enhance updateBalance and createAutoRefillTransaction methods to support additional parameters for improved balance management
This commit is contained in:
Danny Avila 2025-03-22 17:54:25 -04:00 • committed by GitHub
parent 5e6a3ec219
commit 842b68fc32
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
13 changed files with 807 additions and 279 deletions
Download patch file Download diff file Expand all files Collapse all files

98
api/models/Transaction.js
View file

@ -4,8 +4,45 @@ const { getBalanceConfig } = require('~/server/services/Config');
const { getMultiplier, getCacheMultiplier } = require('./tx');
const { logger } = require('~/config');
const Balance = require('./Balance');
const cancelRate = 1.15;
/**
* Updates a user's token balance based on a transaction.
*
* @async
* @function
* @param {Object} params - The function parameters.
* @param {string} params.user - The user ID.
* @param {number} params.incrementValue - The value to increment the balance by (can be negative).
* @param {import('mongoose').UpdateQuery<import('@librechat/data-schemas').IBalance>['$set']} params.setValues
* @returns {Promise<Object>} Returns the updated balance response.
*/
const updateBalance = async ({ user, incrementValue, setValues }) => {
// Use findOneAndUpdate with a conditional update to make the balance update atomic
// This prevents race conditions when multiple transactions are processed concurrently
const balanceResponse = await Balance.findOneAndUpdate(
{ user },
[
{
$set: {
tokenCredits: {
$cond: {
if: { $lt: [{ $add: ['$tokenCredits', incrementValue] }, 0] },
then: 0,
else: { $add: ['$tokenCredits', incrementValue] },
},
},
...setValues,
},
},
],
{ upsert: true, new: true },
).lean();
return balanceResponse;
};
/** Method to calculate and set the tokenValue for a transaction */
transactionSchema.methods.calculateTokenValue = function () {
if (!this.valueKey || !this.tokenType) {
@ -21,6 +58,39 @@ transactionSchema.methods.calculateTokenValue = function () {
}
};
/**
* New static method to create an auto-refill transaction that does NOT trigger a balance update.
* @param {object} txData - Transaction data.
* @param {string} txData.user - The user ID.
* @param {string} txData.tokenType - The type of token.
* @param {string} txData.context - The context of the transaction.
* @param {number} txData.rawAmount - The raw amount of tokens.
* @returns {Promise<object>} - The created transaction.
*/
transactionSchema.statics.createAutoRefillTransaction = async function (txData) {
if (txData.rawAmount != null && isNaN(txData.rawAmount)) {
return;
}
const transaction = new this(txData);
transaction.endpointTokenConfig = txData.endpointTokenConfig;
transaction.calculateTokenValue();
await transaction.save();
const balanceResponse = await updateBalance({
user: transaction.user,
incrementValue: txData.rawAmount,
setValues: { lastRefill: new Date() },
});
const result = {
rate: transaction.rate,
user: transaction.user.toString(),
balance: balanceResponse.tokenCredits,
};
logger.debug('[Balance.check] Auto-refill performed', result);
result.transaction = transaction;
return result;
};
/**
* Static method to create a transaction and update the balance
* @param {txData} txData - Transaction data.
@ -42,18 +112,12 @@ transactionSchema.statics.create = async function (txData) {
return;
}
let balanceResponse = await Balance.findOne({ user: transaction.user }).lean();
let incrementValue = transaction.tokenValue;
if (balanceResponse && balanceResponse.tokenCredits + incrementValue < 0) {
incrementValue = -balanceResponse.tokenCredits;
}
balanceResponse = await Balance.findOneAndUpdate(
{ user: transaction.user },
{ $inc: { tokenCredits: incrementValue } },
{ upsert: true, new: true },
).lean();
const balanceResponse = await updateBalance({
user: transaction.user,
incrementValue,
});
return {
rate: transaction.rate,
@ -84,18 +148,12 @@ transactionSchema.statics.createStructured = async function (txData) {
return;
}
let balanceResponse = await Balance.findOne({ user: transaction.user }).lean();
let incrementValue = transaction.tokenValue;
if (balanceResponse && balanceResponse.tokenCredits + incrementValue < 0) {
incrementValue = -balanceResponse.tokenCredits;
}
balanceResponse = await Balance.findOneAndUpdate(
{ user: transaction.user },
{ $inc: { tokenCredits: incrementValue } },
{ upsert: true, new: true },
).lean();
const balanceResponse = await updateBalance({
user: transaction.user,
incrementValue,
});
return {
rate: transaction.rate,