🧮 feat: Improve structured token spending and testing; fix: Anthropic Cache Spend (#3766)

* chore: update jest and mongodb-memory-server dependencies

* fix: Anthropic edge case of increasing balance

* refactor: Update token calculations in Transaction model/spec

* refactor: `spendTokens` always record transactions, add Tx related tests

* feat: Add error handling for CHECK_BALANCE environment variable

* feat: Update import path for Balance model in Balance controller

* chore: remove logging

* refactor: Improve structured token spend logging in spendTokens.js

* ci: add unit tests for spend token

* ci: Improve structured token spend unit testing

* chore: improve logging phrase for no tx spent from balance

api/models/Transaction.js

@@ -1,5 +1,5 @@
 const mongoose = require('mongoose');
-const { isEnabled } = require('../server/utils/handleText');
+const { isEnabled } = require('~/server/utils/handleText');
 const transactionSchema = require('./schema/transaction');
 const { getMultiplier, getCacheMultiplier } = require('./tx');
 const { logger } = require('~/config');
@@ -76,7 +76,7 @@ transactionSchema.statics.createStructured = async function (txData) {
   await transaction.save();
 
   if (!isEnabled(process.env.CHECK_BALANCE)) {
-    return transaction;
+    return;
   }
 
   let balance = await Balance.findOne({ user: transaction.user }).lean();
@@ -122,28 +122,33 @@ transactionSchema.methods.calculateStructuredTokenValue = function () {
       read: readMultiplier,
     };
 
-    const totalTokens = (this.inputTokens || 0) + (this.writeTokens || 0) + (this.readTokens || 0);
+    const totalPromptTokens =
+      Math.abs(this.inputTokens || 0) +
+      Math.abs(this.writeTokens || 0) +
+      Math.abs(this.readTokens || 0);
 
-    if (totalTokens > 0) {
+    if (totalPromptTokens > 0) {
       this.rate =
-        (inputMultiplier * (this.inputTokens || 0) +
+        (Math.abs(inputMultiplier * (this.inputTokens || 0)) +
-          writeMultiplier * (this.writeTokens || 0) +
+          Math.abs(writeMultiplier * (this.writeTokens || 0)) +
-          readMultiplier * (this.readTokens || 0)) /
+          Math.abs(readMultiplier * (this.readTokens || 0))) /
-        totalTokens;
+        totalPromptTokens;
     } else {
-      this.rate = inputMultiplier; // Default to input rate if no tokens
+      this.rate = Math.abs(inputMultiplier); // Default to input rate if no tokens
     }
 
-    this.tokenValue =
+    this.tokenValue = -(
-      this.inputTokens * inputMultiplier +
+      Math.abs(this.inputTokens || 0) * inputMultiplier +
-      (this.writeTokens || 0) * writeMultiplier +
+      Math.abs(this.writeTokens || 0) * writeMultiplier +
-      (this.readTokens || 0) * readMultiplier;
+      Math.abs(this.readTokens || 0) * readMultiplier
-  } else {
-    const multiplier = Math.abs(
-      getMultiplier({ tokenType: this.tokenType, model, endpointTokenConfig }),
     );
-    this.rate = multiplier;
+
-    this.tokenValue = this.rawAmount * multiplier;
+    this.rawAmount = -totalPromptTokens;
+  } else if (this.tokenType === 'completion') {
+    const multiplier = getMultiplier({ tokenType: this.tokenType, model, endpointTokenConfig });
+    this.rate = Math.abs(multiplier);
+    this.tokenValue = -Math.abs(this.rawAmount) * multiplier;
+    this.rawAmount = -Math.abs(this.rawAmount);
   }
 
   if (this.context && this.tokenType === 'completion' && this.context === 'incomplete') {

api/models/Transaction.spec.js

@@ -0,0 +1,348 @@
+const mongoose = require('mongoose');
+const { MongoMemoryServer } = require('mongodb-memory-server');
+const Balance = require('./Balance');
+const { spendTokens, spendStructuredTokens } = require('./spendTokens');
+const { getMultiplier, getCacheMultiplier } = require('./tx');
+
+let mongoServer;
+
+beforeAll(async () => {
+  mongoServer = await MongoMemoryServer.create();
+  const mongoUri = mongoServer.getUri();
+  await mongoose.connect(mongoUri);
+});
+
+afterAll(async () => {
+  await mongoose.disconnect();
+  await mongoServer.stop();
+});
+
+beforeEach(async () => {
+  await mongoose.connection.dropDatabase();
+});
+
+describe('Regular Token Spending Tests', () => {
+  test('Balance should decrease when spending tokens with spendTokens', async () => {
+    // Arrange
+    const userId = new mongoose.Types.ObjectId();
+    const initialBalance = 10000000; // $10.00
+    await Balance.create({ user: userId, tokenCredits: initialBalance });
+
+    const model = 'gpt-3.5-turbo';
+    const txData = {
+      user: userId,
+      conversationId: 'test-conversation-id',
+      model,
+      context: 'test',
+      endpointTokenConfig: null,
+    };
+
+    const tokenUsage = {
+      promptTokens: 100,
+      completionTokens: 50,
+    };
+
+    // Act
+    process.env.CHECK_BALANCE = 'true';
+    await spendTokens(txData, tokenUsage);
+
+    // Assert
+    console.log('Initial Balance:', initialBalance);
+
+    const updatedBalance = await Balance.findOne({ user: userId });
+    console.log('Updated Balance:', updatedBalance.tokenCredits);
+
+    const promptMultiplier = getMultiplier({ model, tokenType: 'prompt' });
+    const completionMultiplier = getMultiplier({ model, tokenType: 'completion' });
+
+    const expectedPromptCost = tokenUsage.promptTokens * promptMultiplier;
+    const expectedCompletionCost = tokenUsage.completionTokens * completionMultiplier;
+    const expectedTotalCost = expectedPromptCost + expectedCompletionCost;
+    const expectedBalance = initialBalance - expectedTotalCost;
+
+    expect(updatedBalance.tokenCredits).toBeLessThan(initialBalance);
+    expect(updatedBalance.tokenCredits).toBeCloseTo(expectedBalance, 0);
+
+    console.log('Expected Total Cost:', expectedTotalCost);
+    console.log('Actual Balance Decrease:', initialBalance - updatedBalance.tokenCredits);
+  });
+
+  test('spendTokens should handle zero completion tokens', async () => {
+    // Arrange
+    const userId = new mongoose.Types.ObjectId();
+    const initialBalance = 10000000; // $10.00
+    await Balance.create({ user: userId, tokenCredits: initialBalance });
+
+    const model = 'gpt-3.5-turbo';
+    const txData = {
+      user: userId,
+      conversationId: 'test-conversation-id',
+      model,
+      context: 'test',
+      endpointTokenConfig: null,
+    };
+
+    const tokenUsage = {
+      promptTokens: 100,
+      completionTokens: 0,
+    };
+
+    // Act
+    process.env.CHECK_BALANCE = 'true';
+    await spendTokens(txData, tokenUsage);
+
+    // Assert
+    const updatedBalance = await Balance.findOne({ user: userId });
+
+    const promptMultiplier = getMultiplier({ model, tokenType: 'prompt' });
+    const expectedCost = tokenUsage.promptTokens * promptMultiplier;
+    expect(updatedBalance.tokenCredits).toBeCloseTo(initialBalance - expectedCost, 0);
+
+    console.log('Initial Balance:', initialBalance);
+    console.log('Updated Balance:', updatedBalance.tokenCredits);
+    console.log('Expected Cost:', expectedCost);
+  });
+
+  test('spendTokens should handle undefined token counts', async () => {
+    const userId = new mongoose.Types.ObjectId();
+    const initialBalance = 10000000; // $10.00
+    await Balance.create({ user: userId, tokenCredits: initialBalance });
+
+    const model = 'gpt-3.5-turbo';
+    const txData = {
+      user: userId,
+      conversationId: 'test-conversation-id',
+      model,
+      context: 'test',
+      endpointTokenConfig: null,
+    };
+
+    const tokenUsage = {};
+
+    const result = await spendTokens(txData, tokenUsage);
+
+    expect(result).toBeUndefined();
+  });
+
+  test('spendTokens should handle only prompt tokens', async () => {
+    const userId = new mongoose.Types.ObjectId();
+    const initialBalance = 10000000; // $10.00
+    await Balance.create({ user: userId, tokenCredits: initialBalance });
+
+    const model = 'gpt-3.5-turbo';
+    const txData = {
+      user: userId,
+      conversationId: 'test-conversation-id',
+      model,
+      context: 'test',
+      endpointTokenConfig: null,
+    };
+
+    const tokenUsage = { promptTokens: 100 };
+
+    await spendTokens(txData, tokenUsage);
+
+    const updatedBalance = await Balance.findOne({ user: userId });
+
+    const promptMultiplier = getMultiplier({ model, tokenType: 'prompt' });
+    const expectedCost = 100 * promptMultiplier;
+    expect(updatedBalance.tokenCredits).toBeCloseTo(initialBalance - expectedCost, 0);
+  });
+});
+
+describe('Structured Token Spending Tests', () => {
+  test('Balance should decrease and rawAmount should be set when spending a large number of structured tokens', async () => {
+    // Arrange
+    const userId = new mongoose.Types.ObjectId();
+    const initialBalance = 17613154.55; // $17.61
+    await Balance.create({ user: userId, tokenCredits: initialBalance });
+
+    const model = 'claude-3-5-sonnet';
+    const txData = {
+      user: userId,
+      conversationId: 'c23a18da-706c-470a-ac28-ec87ed065199',
+      model,
+      context: 'message',
+      endpointTokenConfig: null, // We'll use the default rates
+    };
+
+    const tokenUsage = {
+      promptTokens: {
+        input: 11,
+        write: 140522,
+        read: 0,
+      },
+      completionTokens: 5,
+    };
+
+    // Get the actual multipliers
+    const promptMultiplier = getMultiplier({ model, tokenType: 'prompt' });
+    const completionMultiplier = getMultiplier({ model, tokenType: 'completion' });
+    const writeMultiplier = getCacheMultiplier({ model, cacheType: 'write' });
+    const readMultiplier = getCacheMultiplier({ model, cacheType: 'read' });
+
+    console.log('Multipliers:', {
+      promptMultiplier,
+      completionMultiplier,
+      writeMultiplier,
+      readMultiplier,
+    });
+
+    // Act
+    process.env.CHECK_BALANCE = 'true';
+    const result = await spendStructuredTokens(txData, tokenUsage);
+
+    // Assert
+    console.log('Initial Balance:', initialBalance);
+    console.log('Updated Balance:', result.completion.balance);
+    console.log('Transaction Result:', result);
+
+    const expectedPromptCost =
+      tokenUsage.promptTokens.input * promptMultiplier +
+      tokenUsage.promptTokens.write * writeMultiplier +
+      tokenUsage.promptTokens.read * readMultiplier;
+    const expectedCompletionCost = tokenUsage.completionTokens * completionMultiplier;
+    const expectedTotalCost = expectedPromptCost + expectedCompletionCost;
+    const expectedBalance = initialBalance - expectedTotalCost;
+
+    console.log('Expected Cost:', expectedTotalCost);
+    console.log('Expected Balance:', expectedBalance);
+
+    expect(result.completion.balance).toBeLessThan(initialBalance);
+
+    // Allow for a small difference (e.g., 100 token credits, which is $0.0001)
+    const allowedDifference = 100;
+    expect(Math.abs(result.completion.balance - expectedBalance)).toBeLessThan(allowedDifference);
+
+    // Check if the decrease is approximately as expected
+    const balanceDecrease = initialBalance - result.completion.balance;
+    expect(balanceDecrease).toBeCloseTo(expectedTotalCost, 0);
+
+    // Check token values
+    const expectedPromptTokenValue = -(
+      tokenUsage.promptTokens.input * promptMultiplier +
+      tokenUsage.promptTokens.write * writeMultiplier +
+      tokenUsage.promptTokens.read * readMultiplier
+    );
+    const expectedCompletionTokenValue = -tokenUsage.completionTokens * completionMultiplier;
+
+    expect(result.prompt.prompt).toBeCloseTo(expectedPromptTokenValue, 1);
+    expect(result.completion.completion).toBe(expectedCompletionTokenValue);
+
+    console.log('Expected prompt tokenValue:', expectedPromptTokenValue);
+    console.log('Actual prompt tokenValue:', result.prompt.prompt);
+    console.log('Expected completion tokenValue:', expectedCompletionTokenValue);
+    console.log('Actual completion tokenValue:', result.completion.completion);
+  });
+
+  test('should handle zero completion tokens in structured spending', async () => {
+    const userId = new mongoose.Types.ObjectId();
+    const initialBalance = 17613154.55;
+    await Balance.create({ user: userId, tokenCredits: initialBalance });
+
+    const model = 'claude-3-5-sonnet';
+    const txData = {
+      user: userId,
+      conversationId: 'test-convo',
+      model,
+      context: 'message',
+    };
+
+    const tokenUsage = {
+      promptTokens: {
+        input: 10,
+        write: 100,
+        read: 5,
+      },
+      completionTokens: 0,
+    };
+
+    process.env.CHECK_BALANCE = 'true';
+    const result = await spendStructuredTokens(txData, tokenUsage);
+
+    expect(result.prompt).toBeDefined();
+    expect(result.completion).toBeUndefined();
+    expect(result.prompt.prompt).toBeLessThan(0);
+  });
+
+  test('should handle only prompt tokens in structured spending', async () => {
+    const userId = new mongoose.Types.ObjectId();
+    const initialBalance = 17613154.55;
+    await Balance.create({ user: userId, tokenCredits: initialBalance });
+
+    const model = 'claude-3-5-sonnet';
+    const txData = {
+      user: userId,
+      conversationId: 'test-convo',
+      model,
+      context: 'message',
+    };
+
+    const tokenUsage = {
+      promptTokens: {
+        input: 10,
+        write: 100,
+        read: 5,
+      },
+    };
+
+    process.env.CHECK_BALANCE = 'true';
+    const result = await spendStructuredTokens(txData, tokenUsage);
+
+    expect(result.prompt).toBeDefined();
+    expect(result.completion).toBeUndefined();
+    expect(result.prompt.prompt).toBeLessThan(0);
+  });
+
+  test('should handle undefined token counts in structured spending', async () => {
+    const userId = new mongoose.Types.ObjectId();
+    const initialBalance = 17613154.55;
+    await Balance.create({ user: userId, tokenCredits: initialBalance });
+
+    const model = 'claude-3-5-sonnet';
+    const txData = {
+      user: userId,
+      conversationId: 'test-convo',
+      model,
+      context: 'message',
+    };
+
+    const tokenUsage = {};
+
+    process.env.CHECK_BALANCE = 'true';
+    const result = await spendStructuredTokens(txData, tokenUsage);
+
+    expect(result).toEqual({
+      prompt: undefined,
+      completion: undefined,
+    });
+  });
+
+  test('should handle incomplete context for completion tokens', async () => {
+    const userId = new mongoose.Types.ObjectId();
+    const initialBalance = 17613154.55;
+    await Balance.create({ user: userId, tokenCredits: initialBalance });
+
+    const model = 'claude-3-5-sonnet';
+    const txData = {
+      user: userId,
+      conversationId: 'test-convo',
+      model,
+      context: 'incomplete',
+    };
+
+    const tokenUsage = {
+      promptTokens: {
+        input: 10,
+        write: 100,
+        read: 5,
+      },
+      completionTokens: 50,
+    };
+
+    process.env.CHECK_BALANCE = 'true';
+    const result = await spendStructuredTokens(txData, tokenUsage);
+
+    expect(result.completion.completion).toBeCloseTo(-50 * 15 * 1.15, 0); // Assuming multiplier is 15 and cancelRate is 1.15
+  });
+});

api/models/spendTokens.js

@@ -32,35 +32,34 @@ const spendTokens = async (txData, tokenUsage) => {
   );
   let prompt, completion;
   try {
-    if (promptTokens >= 0) {
+    if (promptTokens !== undefined) {
       prompt = await Transaction.create({
         ...txData,
         tokenType: 'prompt',
-        rawAmount: -promptTokens,
+        rawAmount: -Math.max(promptTokens, 0),
       });
     }
 
-    if (!completionTokens && isNaN(completionTokens)) {
+    if (completionTokens !== undefined) {
-      logger.debug('[spendTokens] !completionTokens', { prompt, completion });
+      completion = await Transaction.create({
-      return;
+        ...txData,
+        tokenType: 'completion',
+        rawAmount: -Math.max(completionTokens, 0),
+      });
     }
 
-    completion = await Transaction.create({
+    if (prompt || completion) {
-      ...txData,
-      tokenType: 'completion',
-      rawAmount: -completionTokens,
-    });
-
-    prompt &&
-      completion &&
       logger.debug('[spendTokens] Transaction data record against balance:', {
         user: txData.user,
-        prompt: prompt.prompt,
+        prompt: prompt?.prompt,
-        promptRate: prompt.rate,
+        promptRate: prompt?.rate,
-        completion: completion.completion,
+        completion: completion?.completion,
-        completionRate: completion.rate,
+        completionRate: completion?.rate,
-        balance: completion.balance,
+        balance: completion?.balance ?? prompt?.balance,
       });
+    } else {
+      logger.debug('[spendTokens] No transactions incurred against balance');
+    }
   } catch (err) {
     logger.error('[spendTokens]', err);
   }
@@ -102,14 +101,12 @@ const spendStructuredTokens = async (txData, tokenUsage) => {
   try {
     if (promptTokens) {
       const { input = 0, write = 0, read = 0 } = promptTokens;
-      const promptAmount = input + write + read;
       prompt = await Transaction.createStructured({
         ...txData,
         tokenType: 'prompt',
-        rawAmount: -promptAmount,
+        inputTokens: -input,
-        inputTokens: input,
+        writeTokens: -write,
-        writeTokens: write,
+        readTokens: -read,
-        readTokens: read,
       });
     }
 
@@ -121,19 +118,23 @@ const spendStructuredTokens = async (txData, tokenUsage) => {
       });
     }
 
-    prompt &&
+    if (prompt || completion) {
-      completion &&
       logger.debug('[spendStructuredTokens] Transaction data record against balance:', {
         user: txData.user,
-        prompt: prompt.tokenValue,
+        prompt: prompt?.prompt,
-        promptRate: prompt.rate,
+        promptRate: prompt?.rate,
-        completion: completion.tokenValue,
+        completion: completion?.completion,
-        completionRate: completion.rate,
+        completionRate: completion?.rate,
-        balance: completion.balance,
+        balance: completion?.balance ?? prompt?.balance,
       });
+    } else {
+      logger.debug('[spendStructuredTokens] No transactions incurred against balance');
+    }
   } catch (err) {
     logger.error('[spendStructuredTokens]', err);
   }
+
+  return { prompt, completion };
 };
 
 module.exports = { spendTokens, spendStructuredTokens };

api/models/spendTokens.spec.js

@@ -0,0 +1,197 @@
+const mongoose = require('mongoose');
+
+jest.mock('./Transaction', () => ({
+  Transaction: {
+    create: jest.fn(),
+    createStructured: jest.fn(),
+  },
+}));
+
+jest.mock('./Balance', () => ({
+  findOne: jest.fn(),
+  findOneAndUpdate: jest.fn(),
+}));
+
+jest.mock('~/config', () => ({
+  logger: {
+    debug: jest.fn(),
+    error: jest.fn(),
+  },
+}));
+
+// Import after mocking
+const { spendTokens, spendStructuredTokens } = require('./spendTokens');
+const { Transaction } = require('./Transaction');
+const Balance = require('./Balance');
+describe('spendTokens', () => {
+  beforeEach(() => {
+    jest.clearAllMocks();
+    process.env.CHECK_BALANCE = 'true';
+  });
+
+  it('should create transactions for both prompt and completion tokens', async () => {
+    const txData = {
+      user: new mongoose.Types.ObjectId(),
+      conversationId: 'test-convo',
+      model: 'gpt-3.5-turbo',
+      context: 'test',
+    };
+    const tokenUsage = {
+      promptTokens: 100,
+      completionTokens: 50,
+    };
+
+    Transaction.create.mockResolvedValueOnce({ tokenType: 'prompt', rawAmount: -100 });
+    Transaction.create.mockResolvedValueOnce({ tokenType: 'completion', rawAmount: -50 });
+    Balance.findOne.mockResolvedValue({ tokenCredits: 10000 });
+    Balance.findOneAndUpdate.mockResolvedValue({ tokenCredits: 9850 });
+
+    await spendTokens(txData, tokenUsage);
+
+    expect(Transaction.create).toHaveBeenCalledTimes(2);
+    expect(Transaction.create).toHaveBeenCalledWith(
+      expect.objectContaining({
+        tokenType: 'prompt',
+        rawAmount: -100,
+      }),
+    );
+    expect(Transaction.create).toHaveBeenCalledWith(
+      expect.objectContaining({
+        tokenType: 'completion',
+        rawAmount: -50,
+      }),
+    );
+  });
+
+  it('should handle zero completion tokens', async () => {
+    const txData = {
+      user: new mongoose.Types.ObjectId(),
+      conversationId: 'test-convo',
+      model: 'gpt-3.5-turbo',
+      context: 'test',
+    };
+    const tokenUsage = {
+      promptTokens: 100,
+      completionTokens: 0,
+    };
+
+    Transaction.create.mockResolvedValueOnce({ tokenType: 'prompt', rawAmount: -100 });
+    Transaction.create.mockResolvedValueOnce({ tokenType: 'completion', rawAmount: -0 });
+    Balance.findOne.mockResolvedValue({ tokenCredits: 10000 });
+    Balance.findOneAndUpdate.mockResolvedValue({ tokenCredits: 9850 });
+
+    await spendTokens(txData, tokenUsage);
+
+    expect(Transaction.create).toHaveBeenCalledTimes(2);
+    expect(Transaction.create).toHaveBeenCalledWith(
+      expect.objectContaining({
+        tokenType: 'prompt',
+        rawAmount: -100,
+      }),
+    );
+    expect(Transaction.create).toHaveBeenCalledWith(
+      expect.objectContaining({
+        tokenType: 'completion',
+        rawAmount: -0, // Changed from 0 to -0
+      }),
+    );
+  });
+
+  it('should handle undefined token counts', async () => {
+    const txData = {
+      user: new mongoose.Types.ObjectId(),
+      conversationId: 'test-convo',
+      model: 'gpt-3.5-turbo',
+      context: 'test',
+    };
+    const tokenUsage = {};
+
+    await spendTokens(txData, tokenUsage);
+
+    expect(Transaction.create).not.toHaveBeenCalled();
+  });
+
+  it('should not update balance when CHECK_BALANCE is false', async () => {
+    process.env.CHECK_BALANCE = 'false';
+    const txData = {
+      user: new mongoose.Types.ObjectId(),
+      conversationId: 'test-convo',
+      model: 'gpt-3.5-turbo',
+      context: 'test',
+    };
+    const tokenUsage = {
+      promptTokens: 100,
+      completionTokens: 50,
+    };
+
+    Transaction.create.mockResolvedValueOnce({ tokenType: 'prompt', rawAmount: -100 });
+    Transaction.create.mockResolvedValueOnce({ tokenType: 'completion', rawAmount: -50 });
+
+    await spendTokens(txData, tokenUsage);
+
+    expect(Transaction.create).toHaveBeenCalledTimes(2);
+    expect(Balance.findOne).not.toHaveBeenCalled();
+    expect(Balance.findOneAndUpdate).not.toHaveBeenCalled();
+  });
+
+  it('should create structured transactions for both prompt and completion tokens', async () => {
+    const txData = {
+      user: new mongoose.Types.ObjectId(),
+      conversationId: 'test-convo',
+      model: 'claude-3-5-sonnet',
+      context: 'test',
+    };
+    const tokenUsage = {
+      promptTokens: {
+        input: 10,
+        write: 100,
+        read: 5,
+      },
+      completionTokens: 50,
+    };
+
+    Transaction.createStructured.mockResolvedValueOnce({
+      rate: 3.75,
+      user: txData.user.toString(),
+      balance: 9570,
+      prompt: -430,
+    });
+    Transaction.create.mockResolvedValueOnce({
+      rate: 15,
+      user: txData.user.toString(),
+      balance: 8820,
+      completion: -750,
+    });
+
+    const result = await spendStructuredTokens(txData, tokenUsage);
+
+    expect(Transaction.createStructured).toHaveBeenCalledWith(
+      expect.objectContaining({
+        tokenType: 'prompt',
+        inputTokens: -10,
+        writeTokens: -100,
+        readTokens: -5,
+      }),
+    );
+    expect(Transaction.create).toHaveBeenCalledWith(
+      expect.objectContaining({
+        tokenType: 'completion',
+        rawAmount: -50,
+      }),
+    );
+    expect(result).toEqual({
+      prompt: expect.objectContaining({
+        rate: 3.75,
+        user: txData.user.toString(),
+        balance: 9570,
+        prompt: -430,
+      }),
+      completion: expect.objectContaining({
+        rate: 15,
+        user: txData.user.toString(),
+        balance: 8820,
+        completion: -750,
+      }),
+    });
+  });
+});

api/package.json

@@ -101,7 +101,8 @@
     "zod": "^3.22.4"
   },
   "devDependencies": {
-    "jest": "^29.5.0",
+    "jest": "^29.7.0",
+    "mongodb-memory-server": "^10.0.0",
     "nodemon": "^3.0.1",
     "supertest": "^6.3.3"
   }

api/server/controllers/Balance.js

@@ -1,4 +1,4 @@
-const Balance = require('../../models/Balance');
+const Balance = require('~/models/Balance');
 
 async function balanceController(req, res) {
   const { tokenCredits: balance = '' } =

config/add-balance.js

@@ -1,6 +1,7 @@
 const path = require('path');
 require('module-alias')({ base: path.resolve(__dirname, '..', 'api') });
 const { askQuestion, silentExit } = require('./helpers');
+const { isEnabled } = require('~/server/utils/handleText');
 const { Transaction } = require('~/models/Transaction');
 const User = require('~/models/User');
 const connect = require('./connect');
@@ -36,6 +37,12 @@ const connect = require('./connect');
     );
     silentExit(1);
   }
+  if (isEnabled(process.env.CHECK_BALANCE) === false) {
+    console.red(
+      'Error: CHECK_BALANCE environment variable is set to `false`! Please configure: `CHECK_BALANCE=true`',
+    );
+    silentExit(1);
+  }
 
   /**
    * If we don't have the right number of arguments, lets prompt the user for them

package-lock.json

@@ -110,7 +110,8 @@
         "zod": "^3.22.4"
       },
       "devDependencies": {
-        "jest": "^29.5.0",
+        "jest": "^29.7.0",
+        "mongodb-memory-server": "^10.0.0",
         "nodemon": "^3.0.1",
         "supertest": "^6.3.3"
       }
@@ -7857,10 +7858,10 @@
       }
     },
     "node_modules/@mongodb-js/saslprep": {
-      "version": "1.1.4",
+      "version": "1.1.8",
-      "resolved": "https://registry.npmjs.org/@mongodb-js/saslprep/-/saslprep-1.1.4.tgz",
+      "resolved": "https://registry.npmjs.org/@mongodb-js/saslprep/-/saslprep-1.1.8.tgz",
-      "integrity": "sha512-8zJ8N1x51xo9hwPh6AWnKdLGEC5N3lDa6kms1YHmFBoRhTpJR6HG8wWk0td1MVCu9cD4YBrvjZEtd5Obw0Fbnw==",
+      "integrity": "sha512-qKwC/M/nNNaKUBMQ0nuzm47b7ZYWQHN3pcXq4IIcoSBc2hOIrflAxJduIvvqmhoz3gR2TacTAs8vlsCVPkiEdQ==",
-      "optional": true,
+      "devOptional": true,
       "dependencies": {
         "sparse-bitfield": "^3.0.3"
       }
@@ -13056,6 +13057,15 @@
       "resolved": "https://registry.npmjs.org/async/-/async-3.2.5.tgz",
       "integrity": "sha512-baNZyqaaLhyLVKm/DlvdW051MSgO6b8eVfIezl9E5PqWxFgzLm/wQntEW4zOytVburDEr0JlALEpdOFwvErLsg=="
     },
+    "node_modules/async-mutex": {
+      "version": "0.5.0",
+      "resolved": "https://registry.npmjs.org/async-mutex/-/async-mutex-0.5.0.tgz",
+      "integrity": "sha512-1A94B18jkJ3DYq284ohPxoXbfTA5HsQ7/Mf4DEhcyLx3Bz27Rh59iScbB6EPiP+B+joue6YCxcMXSbFC1tZKwA==",
+      "dev": true,
+      "dependencies": {
+        "tslib": "^2.4.0"
+      }
+    },
     "node_modules/asynckit": {
       "version": "0.4.0",
       "resolved": "https://registry.npmjs.org/asynckit/-/asynckit-0.4.0.tgz",
@@ -13820,6 +13830,15 @@
         "ieee754": "^1.2.1"
       }
     },
+    "node_modules/buffer-crc32": {
+      "version": "0.2.13",
+      "resolved": "https://registry.npmjs.org/buffer-crc32/-/buffer-crc32-0.2.13.tgz",
+      "integrity": "sha512-VO9Ht/+p3SN7SKWqcrgEzjGbRSJYTx+Q1pTQC0wrWqHx0vpJraQ6GtHx8tvcg1rlK1byhU5gccxgOgj7B0TDkQ==",
+      "dev": true,
+      "engines": {
+        "node": "*"
+      }
+    },
     "node_modules/buffer-equal-constant-time": {
       "version": "1.0.1",
       "resolved": "https://registry.npmjs.org/buffer-equal-constant-time/-/buffer-equal-constant-time-1.0.1.tgz",
@@ -22146,7 +22165,7 @@
       "version": "1.5.0",
       "resolved": "https://registry.npmjs.org/memory-pager/-/memory-pager-1.5.0.tgz",
       "integrity": "sha512-ZS4Bp4r/Zoeq6+NLJpP+0Zzm0pR8whtGPf1XExKLJBAczGMnSi3It14OiNCStjQjM6NU1okjQGSxgEZN8eBYKg==",
-      "optional": true
+      "devOptional": true
     },
     "node_modules/merge-descriptors": {
       "version": "1.0.1",
@@ -22985,6 +23004,184 @@
         "whatwg-url": "^11.0.0"
       }
     },
+    "node_modules/mongodb-memory-server": {
+      "version": "10.0.0",
+      "resolved": "https://registry.npmjs.org/mongodb-memory-server/-/mongodb-memory-server-10.0.0.tgz",
+      "integrity": "sha512-7Geo/s4lst/QHw+N8/stdnyb08xn68O0zbSee62jgoPfWOXfSPhX9a8OvyOY/o23oYk9ra2EpA2Oejenb3JKfw==",
+      "dev": true,
+      "hasInstallScript": true,
+      "dependencies": {
+        "mongodb-memory-server-core": "10.0.0",
+        "tslib": "^2.6.3"
+      },
+      "engines": {
+        "node": ">=16.20.1"
+      }
+    },
+    "node_modules/mongodb-memory-server-core": {
+      "version": "10.0.0",
+      "resolved": "https://registry.npmjs.org/mongodb-memory-server-core/-/mongodb-memory-server-core-10.0.0.tgz",
+      "integrity": "sha512-AdYi4nVqe3Pk95fRJ+DegbDdEfAG9wujNsVvJWbwh8+ZJd+d3JJK1PHxRyJ9rMvoczvlli5M30eMig7zBuF5pQ==",
+      "dev": true,
+      "dependencies": {
+        "async-mutex": "^0.5.0",
+        "camelcase": "^6.3.0",
+        "debug": "^4.3.5",
+        "find-cache-dir": "^3.3.2",
+        "follow-redirects": "^1.15.6",
+        "https-proxy-agent": "^7.0.5",
+        "mongodb": "^6.7.0",
+        "new-find-package-json": "^2.0.0",
+        "semver": "^7.6.3",
+        "tar-stream": "^3.1.7",
+        "tslib": "^2.6.3",
+        "yauzl": "^3.1.3"
+      },
+      "engines": {
+        "node": ">=16.20.1"
+      }
+    },
+    "node_modules/mongodb-memory-server-core/node_modules/@types/whatwg-url": {
+      "version": "11.0.5",
+      "resolved": "https://registry.npmjs.org/@types/whatwg-url/-/whatwg-url-11.0.5.tgz",
+      "integrity": "sha512-coYR071JRaHa+xoEvvYqvnIHaVqaYrLPbsufM9BF63HkwI5Lgmy2QR8Q5K/lYDYo5AK82wOvSOS0UsLTpTG7uQ==",
+      "dev": true,
+      "dependencies": {
+        "@types/webidl-conversions": "*"
+      }
+    },
+    "node_modules/mongodb-memory-server-core/node_modules/agent-base": {
+      "version": "7.1.1",
+      "resolved": "https://registry.npmjs.org/agent-base/-/agent-base-7.1.1.tgz",
+      "integrity": "sha512-H0TSyFNDMomMNJQBn8wFV5YC/2eJ+VXECwOadZJT554xP6cODZHPX3H9QMQECxvrgiSOP1pHjy1sMWQVYJOUOA==",
+      "dev": true,
+      "dependencies": {
+        "debug": "^4.3.4"
+      },
+      "engines": {
+        "node": ">= 14"
+      }
+    },
+    "node_modules/mongodb-memory-server-core/node_modules/bson": {
+      "version": "6.8.0",
+      "resolved": "https://registry.npmjs.org/bson/-/bson-6.8.0.tgz",
+      "integrity": "sha512-iOJg8pr7wq2tg/zSlCCHMi3hMm5JTOxLTagf3zxhcenHsFp+c6uOs6K7W5UE7A4QIJGtqh/ZovFNMP4mOPJynQ==",
+      "dev": true,
+      "engines": {
+        "node": ">=16.20.1"
+      }
+    },
+    "node_modules/mongodb-memory-server-core/node_modules/debug": {
+      "version": "4.3.6",
+      "resolved": "https://registry.npmjs.org/debug/-/debug-4.3.6.tgz",
+      "integrity": "sha512-O/09Bd4Z1fBrU4VzkhFqVgpPzaGbw6Sm9FEkBT1A/YBXQFGuuSxa1dN2nxgxS34JmKXqYx8CZAwEVoJFImUXIg==",
+      "dev": true,
+      "dependencies": {
+        "ms": "2.1.2"
+      },
+      "engines": {
+        "node": ">=6.0"
+      },
+      "peerDependenciesMeta": {
+        "supports-color": {
+          "optional": true
+        }
+      }
+    },
+    "node_modules/mongodb-memory-server-core/node_modules/https-proxy-agent": {
+      "version": "7.0.5",
+      "resolved": "https://registry.npmjs.org/https-proxy-agent/-/https-proxy-agent-7.0.5.tgz",
+      "integrity": "sha512-1e4Wqeblerz+tMKPIq2EMGiiWW1dIjZOksyHWSUm1rmuvw/how9hBHZ38lAGj5ID4Ik6EdkOw7NmWPy6LAwalw==",
+      "dev": true,
+      "dependencies": {
+        "agent-base": "^7.0.2",
+        "debug": "4"
+      },
+      "engines": {
+        "node": ">= 14"
+      }
+    },
+    "node_modules/mongodb-memory-server-core/node_modules/mongodb": {
+      "version": "6.8.0",
+      "resolved": "https://registry.npmjs.org/mongodb/-/mongodb-6.8.0.tgz",
+      "integrity": "sha512-HGQ9NWDle5WvwMnrvUxsFYPd3JEbqD3RgABHBQRuoCEND0qzhsd0iH5ypHsf1eJ+sXmvmyKpP+FLOKY8Il7jMw==",
+      "dev": true,
+      "dependencies": {
+        "@mongodb-js/saslprep": "^1.1.5",
+        "bson": "^6.7.0",
+        "mongodb-connection-string-url": "^3.0.0"
+      },
+      "engines": {
+        "node": ">=16.20.1"
+      },
+      "peerDependencies": {
+        "@aws-sdk/credential-providers": "^3.188.0",
+        "@mongodb-js/zstd": "^1.1.0",
+        "gcp-metadata": "^5.2.0",
+        "kerberos": "^2.0.1",
+        "mongodb-client-encryption": ">=6.0.0 <7",
+        "snappy": "^7.2.2",
+        "socks": "^2.7.1"
+      },
+      "peerDependenciesMeta": {
+        "@aws-sdk/credential-providers": {
+          "optional": true
+        },
+        "@mongodb-js/zstd": {
+          "optional": true
+        },
+        "gcp-metadata": {
+          "optional": true
+        },
+        "kerberos": {
+          "optional": true
+        },
+        "mongodb-client-encryption": {
+          "optional": true
+        },
+        "snappy": {
+          "optional": true
+        },
+        "socks": {
+          "optional": true
+        }
+      }
+    },
+    "node_modules/mongodb-memory-server-core/node_modules/mongodb-connection-string-url": {
+      "version": "3.0.1",
+      "resolved": "https://registry.npmjs.org/mongodb-connection-string-url/-/mongodb-connection-string-url-3.0.1.tgz",
+      "integrity": "sha512-XqMGwRX0Lgn05TDB4PyG2h2kKO/FfWJyCzYQbIhXUxz7ETt0I/FqHjUeqj37irJ+Dl1ZtU82uYyj14u2XsZKfg==",
+      "dev": true,
+      "dependencies": {
+        "@types/whatwg-url": "^11.0.2",
+        "whatwg-url": "^13.0.0"
+      }
+    },
+    "node_modules/mongodb-memory-server-core/node_modules/tr46": {
+      "version": "4.1.1",
+      "resolved": "https://registry.npmjs.org/tr46/-/tr46-4.1.1.tgz",
+      "integrity": "sha512-2lv/66T7e5yNyhAAC4NaKe5nVavzuGJQVVtRYLyQ2OI8tsJ61PMLlelehb0wi2Hx6+hT/OJUWZcw8MjlSRnxvw==",
+      "dev": true,
+      "dependencies": {
+        "punycode": "^2.3.0"
+      },
+      "engines": {
+        "node": ">=14"
+      }
+    },
+    "node_modules/mongodb-memory-server-core/node_modules/whatwg-url": {
+      "version": "13.0.0",
+      "resolved": "https://registry.npmjs.org/whatwg-url/-/whatwg-url-13.0.0.tgz",
+      "integrity": "sha512-9WWbymnqj57+XEuqADHrCJ2eSXzn8WXIW/YSGaZtb2WKAInQ6CHfaUUcTyyver0p8BDg5StLQq8h1vtZuwmOig==",
+      "dev": true,
+      "dependencies": {
+        "tr46": "^4.1.1",
+        "webidl-conversions": "^7.0.0"
+      },
+      "engines": {
+        "node": ">=16"
+      }
+    },
     "node_modules/mongoose": {
       "version": "7.6.8",
       "resolved": "https://registry.npmjs.org/mongoose/-/mongoose-7.6.8.tgz",
@@ -23215,6 +23412,18 @@
       "resolved": "https://registry.npmjs.org/neo-async/-/neo-async-2.6.2.tgz",
       "integrity": "sha512-Yd3UES5mWCSqR+qNT93S3UoYUkqAZ9lLg8a7g9rimsWmYGK8cVToA4/sF3RrshdyV3sAGMXVUmpMYOw+dLpOuw=="
     },
+    "node_modules/new-find-package-json": {
+      "version": "2.0.0",
+      "resolved": "https://registry.npmjs.org/new-find-package-json/-/new-find-package-json-2.0.0.tgz",
+      "integrity": "sha512-lDcBsjBSMlj3LXH2v/FW3txlh2pYTjmbOXPYJD93HI5EwuLzI11tdHSIpUMmfq/IOsldj4Ps8M8flhm+pCK4Ew==",
+      "dev": true,
+      "dependencies": {
+        "debug": "^4.3.4"
+      },
+      "engines": {
+        "node": ">=12.22.0"
+      }
+    },
     "node_modules/node-abi": {
       "version": "3.54.0",
       "resolved": "https://registry.npmjs.org/node-abi/-/node-abi-3.54.0.tgz",
@@ -24424,6 +24633,12 @@
         "url": "https://github.com/sponsors/Borewit"
       }
     },
+    "node_modules/pend": {
+      "version": "1.2.0",
+      "resolved": "https://registry.npmjs.org/pend/-/pend-1.2.0.tgz",
+      "integrity": "sha512-F3asv42UuXchdzt+xXqfW1OGlVBe+mxa2mqI0pg5yAHZPvFmY3Y6drSf/GQ1A86WgWEN9Kzh/WrgKa6iGcHXLg==",
+      "dev": true
+    },
     "node_modules/pica": {
       "version": "9.0.1",
       "resolved": "https://registry.npmjs.org/pica/-/pica-9.0.1.tgz",
@@ -27525,12 +27740,9 @@
       "integrity": "sha512-6aU+Rwsezw7VR8/nyvKTx8QpWH9FrcYiXXlqC4z5d5XQBDRqtbfsRjnwGyqbi3gddNtWHuEk9OANUotL26qKUw=="
     },
     "node_modules/semver": {
-      "version": "7.6.0",
+      "version": "7.6.3",
-      "resolved": "https://registry.npmjs.org/semver/-/semver-7.6.0.tgz",
+      "resolved": "https://registry.npmjs.org/semver/-/semver-7.6.3.tgz",
-      "integrity": "sha512-EnwXhrlwXMk9gKu5/flx5sv/an57AkRplG3hTK68W7FRDN+k+OWBj65M7719OkA82XLBxrcX0KSHj+X5COhOVg==",
+      "integrity": "sha512-oVekP1cKtI+CTDvHWYFUcMtsK/00wmAEfyqKfNdARm8u1wNVhSgaX7A8d4UuIlUI5e84iEwOhs7ZPYRmzU9U6A==",
-      "dependencies": {
-        "lru-cache": "^6.0.0"
-      },
       "bin": {
         "semver": "bin/semver.js"
       },
@@ -27538,22 +27750,6 @@
         "node": ">=10"
       }
     },
-    "node_modules/semver/node_modules/lru-cache": {
-      "version": "6.0.0",
-      "resolved": "https://registry.npmjs.org/lru-cache/-/lru-cache-6.0.0.tgz",
-      "integrity": "sha512-Jo6dJ04CmSjuznwJSS3pUeWmd/H0ffTlkXXgwZi+eq1UCmqQwCh+eLsYOYCwY991i2Fah4h1BEMCx4qThGbsiA==",
-      "dependencies": {
-        "yallist": "^4.0.0"
-      },
-      "engines": {
-        "node": ">=10"
-      }
-    },
-    "node_modules/semver/node_modules/yallist": {
-      "version": "4.0.0",
-      "resolved": "https://registry.npmjs.org/yallist/-/yallist-4.0.0.tgz",
-      "integrity": "sha512-3wdGidZyq5PB084XLES5TpOSRA3wjXAlIWMhum2kRcv/41Sn2emQ0dycQW4uZXLejwKvg6EsvbdlVL+FYEct7A=="
-    },
     "node_modules/send": {
       "version": "0.18.0",
       "resolved": "https://registry.npmjs.org/send/-/send-0.18.0.tgz",
@@ -27965,7 +28161,7 @@
       "version": "3.0.3",
       "resolved": "https://registry.npmjs.org/sparse-bitfield/-/sparse-bitfield-3.0.3.tgz",
       "integrity": "sha512-kvzhi7vqKTfkh0PZU+2D2PIllw2ymqJKujUcyPMd9Y75Nv4nPbGJZXNhxsgdQab2BmlDct1YnfQCguEvHr7VsQ==",
-      "optional": true,
+      "devOptional": true,
       "dependencies": {
         "memory-pager": "^1.0.2"
       }
@@ -29193,9 +29389,9 @@
       }
     },
     "node_modules/tslib": {
-      "version": "2.6.2",
+      "version": "2.7.0",
-      "resolved": "https://registry.npmjs.org/tslib/-/tslib-2.6.2.tgz",
+      "resolved": "https://registry.npmjs.org/tslib/-/tslib-2.7.0.tgz",
-      "integrity": "sha512-AEYxH93jGFPn/a2iVAwW87VuUIkR1FVUKB77NwMF7nBTDkDrrT/Hpt/IrCJ0QXhW27jTBDcf5ZY7w6RiqTMw2Q=="
+      "integrity": "sha512-gLXCKdN1/j47AiHiOkJN69hJmcbGTHI0ImLmbYLHykhgeN0jVGola9yVjFgzCUklsZQMW55o+dW7IXv3RCXDzA=="
     },
     "node_modules/tsutils": {
       "version": "3.21.0",
@@ -31454,6 +31650,19 @@
         "node": ">=8"
       }
     },
+    "node_modules/yauzl": {
+      "version": "3.1.3",
+      "resolved": "https://registry.npmjs.org/yauzl/-/yauzl-3.1.3.tgz",
+      "integrity": "sha512-JCCdmlJJWv7L0q/KylOekyRaUrdEoUxWkWVcgorosTROCFWiS9p2NNPE9Yb91ak7b1N5SxAZEliWpspbZccivw==",
+      "dev": true,
+      "dependencies": {
+        "buffer-crc32": "~0.2.3",
+        "pend": "~1.2.0"
+      },
+      "engines": {
+        "node": ">=12"
+      }
+    },
     "node_modules/yocto-queue": {
       "version": "0.1.0",
       "resolved": "https://registry.npmjs.org/yocto-queue/-/yocto-queue-0.1.0.tgz",