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

* 🏗️ 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

View file

@ -36,7 +36,7 @@ const spendTokens = async (txData, tokenUsage) => {
prompt = await Transaction.create({
...txData,
tokenType: 'prompt',
rawAmount: -Math.max(promptTokens, 0),
rawAmount: promptTokens === 0 ? 0 : -Math.max(promptTokens, 0),
});
}
@ -44,7 +44,7 @@ const spendTokens = async (txData, tokenUsage) => {
completion = await Transaction.create({
...txData,
tokenType: 'completion',
rawAmount: -Math.max(completionTokens, 0),
rawAmount: completionTokens === 0 ? 0 : -Math.max(completionTokens, 0),
});
}