api/models/spendTokens.spec.js

const mongoose = require('mongoose');
const { MongoMemoryServer } = require('mongodb-memory-server');
const { spendTokens, spendStructuredTokens } = require('./spendTokens');
const { createTransaction, createAutoRefillTransaction } = require('./Transaction');

require('~/db/models');

jest.mock('~/config', () => ({
  logger: {
    debug: jest.fn(),
    error: jest.fn(),
  },
}));

describe('spendTokens', () => {
  let mongoServer;
  let userId;
  let Transaction;
  let Balance;

  beforeAll(async () => {
    mongoServer = await MongoMemoryServer.create();
    await mongoose.connect(mongoServer.getUri());

    Transaction = mongoose.model('Transaction');
    Balance = mongoose.model('Balance');
  });

  afterAll(async () => {
    await mongoose.disconnect();
    await mongoServer.stop();
  });

  beforeEach(async () => {
    // Clear collections before each test
    await Transaction.deleteMany({});
    await Balance.deleteMany({});

    // Create a new user ID for each test
    userId = new mongoose.Types.ObjectId();

    // Balance config is now passed directly in txData
  });

  it('should create transactions for both prompt and completion tokens', async () => {
    // Create a balance for the user
    await Balance.create({
      user: userId,
      tokenCredits: 10000,
    });

    const txData = {
      user: userId,
      conversationId: 'test-convo',
      model: 'gpt-3.5-turbo',
      context: 'test',
      balance: { enabled: true },
    };
    const tokenUsage = {
      promptTokens: 100,
      completionTokens: 50,
    };

    await spendTokens(txData, tokenUsage);

    // Verify transactions were created
    const transactions = await Transaction.find({ user: userId }).sort({ tokenType: 1 });
    expect(transactions).toHaveLength(2);

    // Check completion transaction
    expect(transactions[0].tokenType).toBe('completion');
    expect(transactions[0].rawAmount).toBe(-50);

    // Check prompt transaction
    expect(transactions[1].tokenType).toBe('prompt');
    expect(transactions[1].rawAmount).toBe(-100);

    // Verify balance was updated
    const balance = await Balance.findOne({ user: userId });
    expect(balance).toBeDefined();
    expect(balance.tokenCredits).toBeLessThan(10000); // Balance should be reduced
  });

  it('should handle zero completion tokens', async () => {
    // Create a balance for the user
    await Balance.create({
      user: userId,
      tokenCredits: 10000,
    });

    const txData = {
      user: userId,
      conversationId: 'test-convo',
      model: 'gpt-3.5-turbo',
      context: 'test',
      balance: { enabled: true },
    };
    const tokenUsage = {
      promptTokens: 100,
      completionTokens: 0,
    };

    await spendTokens(txData, tokenUsage);

    // Verify transactions were created
    const transactions = await Transaction.find({ user: userId }).sort({ tokenType: 1 });
    expect(transactions).toHaveLength(2);

    // Check completion transaction
    expect(transactions[0].tokenType).toBe('completion');
    // In JavaScript -0 and 0 are different but functionally equivalent
    // Use Math.abs to handle both 0 and -0
    expect(Math.abs(transactions[0].rawAmount)).toBe(0);

    // Check prompt transaction
    expect(transactions[1].tokenType).toBe('prompt');
    expect(transactions[1].rawAmount).toBe(-100);
  });

  it('should handle undefined token counts', async () => {
    const txData = {
      user: userId,
      conversationId: 'test-convo',
      model: 'gpt-3.5-turbo',
      context: 'test',
      balance: { enabled: true },
    };
    const tokenUsage = {};

    await spendTokens(txData, tokenUsage);

    // Verify no transactions were created
    const transactions = await Transaction.find({ user: userId });
    expect(transactions).toHaveLength(0);
  });

  it('should not update balance when the balance feature is disabled', async () => {
    // Balance is now passed directly in txData
    // Create a balance for the user
    await Balance.create({
      user: userId,
      tokenCredits: 10000,
    });

    const txData = {
      user: userId,
      conversationId: 'test-convo',
      model: 'gpt-3.5-turbo',
      context: 'test',
      balance: { enabled: false },
    };
    const tokenUsage = {
      promptTokens: 100,
      completionTokens: 50,
    };

    await spendTokens(txData, tokenUsage);

    // Verify transactions were created
    const transactions = await Transaction.find({ user: userId });
    expect(transactions).toHaveLength(2);

    // Verify balance was not updated (should still be 10000)
    const balance = await Balance.findOne({ user: userId });
    expect(balance.tokenCredits).toBe(10000);
  });

  it('should not allow balance to go below zero when spending tokens', async () => {
    // Create a balance with a low amount
    await Balance.create({
      user: userId,
      tokenCredits: 5000,
    });

    const txData = {
      user: userId,
      conversationId: 'test-convo',
      model: 'gpt-4', // Using a more expensive model
      context: 'test',
      balance: { enabled: true },
    };

    // Spending more tokens than the user has balance for
    const tokenUsage = {
      promptTokens: 1000,
      completionTokens: 500,
    };

    await spendTokens(txData, tokenUsage);

    // Verify transactions were created
    const transactions = await Transaction.find({ user: userId }).sort({ tokenType: 1 });
    expect(transactions).toHaveLength(2);

    // Verify balance was reduced to exactly 0, not negative
    const balance = await Balance.findOne({ user: userId });
    expect(balance).toBeDefined();
    expect(balance.tokenCredits).toBe(0);

    // Check that the transaction records show the adjusted values
    const transactionResults = await Promise.all(
      transactions.map((t) =>
        createTransaction({
          ...txData,
          tokenType: t.tokenType,
          rawAmount: t.rawAmount,
        }),
      ),
    );

    // The second transaction should have an adjusted value since balance is already 0
    expect(transactionResults[1]).toEqual(
      expect.objectContaining({
        balance: 0,
      }),
    );
  });

  it('should handle multiple transactions in sequence with low balance and not increase balance', async () => {
    // This test is specifically checking for the issue reported in production
    // where the balance increases after a transaction when it should remain at 0
    // Create a balance with a very low amount
    await Balance.create({
      user: userId,
      tokenCredits: 100,
    });

    // First transaction - should reduce balance to 0
    const txData1 = {
      user: userId,
      conversationId: 'test-convo-1',
      model: 'gpt-4',
      context: 'test',
      balance: { enabled: true },
    };

    const tokenUsage1 = {
      promptTokens: 100,
      completionTokens: 50,
    };

    await spendTokens(txData1, tokenUsage1);

    // Check balance after first transaction
    let balance = await Balance.findOne({ user: userId });
    expect(balance.tokenCredits).toBe(0);

    // Second transaction - should keep balance at 0, not make it negative or increase it
    const txData2 = {
      user: userId,
      conversationId: 'test-convo-2',
      model: 'gpt-4',
      context: 'test',
      balance: { enabled: true },
    };

    const tokenUsage2 = {
      promptTokens: 200,
      completionTokens: 100,
    };

    await spendTokens(txData2, tokenUsage2);

    // Check balance after second transaction - should still be 0
    balance = await Balance.findOne({ user: userId });
    expect(balance.tokenCredits).toBe(0);

    // Verify all transactions were created
    const transactions = await Transaction.find({ user: userId });
    expect(transactions).toHaveLength(4); // 2 transactions (prompt+completion) for each call

    // Let's examine the actual transaction records to see what's happening
    const transactionDetails = await Transaction.find({ user: userId }).sort({ createdAt: 1 });

    // Log the transaction details for debugging
    console.log('Transaction details:');
    transactionDetails.forEach((tx, i) => {
      console.log(`Transaction ${i + 1}:`, {
        tokenType: tx.tokenType,
        rawAmount: tx.rawAmount,
        tokenValue: tx.tokenValue,
        model: tx.model,
      });
    });

    // Check the return values from Transaction.create directly
    // This is to verify that the incrementValue is not becoming positive
    const directResult = await createTransaction({
      user: userId,
      conversationId: 'test-convo-3',
      model: 'gpt-4',
      tokenType: 'completion',
      rawAmount: -100,
      context: 'test',
      balance: { enabled: true },
    });

    console.log('Direct Transaction.create result:', directResult);

    // The completion value should never be positive
    expect(directResult.completion).not.toBeGreaterThan(0);
  });

  it('should ensure tokenValue is always negative for spending tokens', async () => {
    // Create a balance for the user
    await Balance.create({
      user: userId,
      tokenCredits: 10000,
    });

    // Test with various models to check multiplier calculations
    const models = ['gpt-3.5-turbo', 'gpt-4', 'claude-3-5-sonnet'];

    for (const model of models) {
      const txData = {
        user: userId,
        conversationId: `test-convo-${model}`,
        model,
        context: 'test',
        balance: { enabled: true },
      };

      const tokenUsage = {
        promptTokens: 100,
        completionTokens: 50,
      };

      await spendTokens(txData, tokenUsage);

      // Get the transactions for this model
      const transactions = await Transaction.find({
        user: userId,
        model,
      });

      // Verify tokenValue is negative for all transactions
      transactions.forEach((tx) => {
        console.log(`Model ${model}, Type ${tx.tokenType}: tokenValue = ${tx.tokenValue}`);
        expect(tx.tokenValue).toBeLessThan(0);
      });
    }
  });

  it('should handle structured transactions in sequence with low balance', async () => {
    // Create a balance with a very low amount
    await Balance.create({
      user: userId,
      tokenCredits: 100,
    });

    // First transaction - should reduce balance to 0
    const txData1 = {
      user: userId,
      conversationId: 'test-convo-1',
      model: 'claude-3-5-sonnet',
      context: 'test',
      balance: { enabled: true },
    };

    const tokenUsage1 = {
      promptTokens: {
        input: 10,
        write: 100,
        read: 5,
      },
      completionTokens: 50,
    };

    await spendStructuredTokens(txData1, tokenUsage1);

    // Check balance after first transaction
    let balance = await Balance.findOne({ user: userId });
    expect(balance.tokenCredits).toBe(0);

    // Second transaction - should keep balance at 0, not make it negative or increase it
    const txData2 = {
      user: userId,
      conversationId: 'test-convo-2',
      model: 'claude-3-5-sonnet',
      context: 'test',
      balance: { enabled: true },
    };

    const tokenUsage2 = {
      promptTokens: {
        input: 20,
        write: 200,
        read: 10,
      },
      completionTokens: 100,
    };

    await spendStructuredTokens(txData2, tokenUsage2);

    // Check balance after second transaction - should still be 0
    balance = await Balance.findOne({ user: userId });
    expect(balance.tokenCredits).toBe(0);

    // Verify all transactions were created
    const transactions = await Transaction.find({ user: userId });
    expect(transactions).toHaveLength(4); // 2 transactions (prompt+completion) for each call

    // Let's examine the actual transaction records to see what's happening
    const transactionDetails = await Transaction.find({ user: userId }).sort({ createdAt: 1 });

    // Log the transaction details for debugging
    console.log('Structured transaction details:');
    transactionDetails.forEach((tx, i) => {
      console.log(`Transaction ${i + 1}:`, {
        tokenType: tx.tokenType,
        rawAmount: tx.rawAmount,
        tokenValue: tx.tokenValue,
        inputTokens: tx.inputTokens,
        writeTokens: tx.writeTokens,
        readTokens: tx.readTokens,
        model: tx.model,
      });
    });
  });

  it('should not allow balance to go below zero when spending structured tokens', async () => {
    // Create a balance with a low amount
    await Balance.create({
      user: userId,
      tokenCredits: 5000,
    });

    const txData = {
      user: userId,
      conversationId: 'test-convo',
      model: 'claude-3-5-sonnet', // Using a model that supports structured tokens
      context: 'test',
      balance: { enabled: true },
    };

    // Spending more tokens than the user has balance for
    const tokenUsage = {
      promptTokens: {
        input: 100,
        write: 1000,
        read: 50,
      },
      completionTokens: 500,
    };

    const result = await spendStructuredTokens(txData, tokenUsage);

    // Verify transactions were created
    const transactions = await Transaction.find({ user: userId }).sort({ tokenType: 1 });
    expect(transactions).toHaveLength(2);

    // Verify balance was reduced to exactly 0, not negative
    const balance = await Balance.findOne({ user: userId });
    expect(balance).toBeDefined();
    expect(balance.tokenCredits).toBe(0);

    // The result should show the adjusted values
    expect(result).toEqual({
      prompt: expect.objectContaining({
        user: userId.toString(),
        balance: expect.any(Number),
      }),
      completion: expect.objectContaining({
        user: userId.toString(),
        balance: 0, // Final balance should be 0
      }),
    });
  });

  it('should handle multiple concurrent transactions correctly with a high balance', async () => {
    // Create a balance with a high amount
    const initialBalance = 10000000;
    await Balance.create({
      user: userId,
      tokenCredits: initialBalance,
    });

    // Simulate the recordCollectedUsage function from the production code
    const conversationId = 'test-concurrent-convo';
    const context = 'message';
    const model = 'gpt-4';

    const amount = 50;
    // Create `amount` of usage records to simulate multiple transactions
    const collectedUsage = Array.from({ length: amount }, (_, i) => ({
      model,
      input_tokens: 100 + i * 10, // Increasing input tokens
      output_tokens: 50 + i * 5, // Increasing output tokens
      input_token_details: {
        cache_creation: i % 2 === 0 ? 20 : 0, // Some have cache creation
        cache_read: i % 3 === 0 ? 10 : 0, // Some have cache read
      },
    }));

    // Process all transactions concurrently to simulate race conditions
    const promises = [];
    let expectedTotalSpend = 0;

    for (let i = 0; i < collectedUsage.length; i++) {
      const usage = collectedUsage[i];
      if (!usage) {
        continue;
      }

      const cache_creation = Number(usage.input_token_details?.cache_creation) || 0;
      const cache_read = Number(usage.input_token_details?.cache_read) || 0;

      const txMetadata = {
        context,
        conversationId,
        user: userId,
        model: usage.model,
        balance: { enabled: true },
      };

      // Calculate expected spend for this transaction
      const promptTokens = usage.input_tokens;
      const completionTokens = usage.output_tokens;

      // For regular transactions
      if (cache_creation === 0 && cache_read === 0) {
        // Add to expected spend using the correct multipliers from tx.js
        // For gpt-4, the multipliers are: prompt=30, completion=60
        expectedTotalSpend += promptTokens * 30; // gpt-4 prompt rate is 30
        expectedTotalSpend += completionTokens * 60; // gpt-4 completion rate is 60

        promises.push(
          spendTokens(txMetadata, {
            promptTokens,
            completionTokens,
          }),
        );
      } else {
        // For structured transactions with cache operations
        // The multipliers for claude models with cache operations are different
        // But since we're using gpt-4 in the test, we need to use appropriate values
        expectedTotalSpend += promptTokens * 30; // Base prompt rate for gpt-4
        // Since gpt-4 doesn't have cache multipliers defined, we'll use the prompt rate
        expectedTotalSpend += cache_creation * 30; // Write rate (using prompt rate as fallback)
        expectedTotalSpend += cache_read * 30; // Read rate (using prompt rate as fallback)
        expectedTotalSpend += completionTokens * 60; // Completion rate for gpt-4

        promises.push(
          spendStructuredTokens(txMetadata, {
            promptTokens: {
              input: promptTokens,
              write: cache_creation,
              read: cache_read,
            },
            completionTokens,
          }),
        );
      }
    }

    // Wait for all transactions to complete
    await Promise.all(promises);

    // Verify final balance
    const finalBalance = await Balance.findOne({ user: userId });
    expect(finalBalance).toBeDefined();

    // The final balance should be the initial balance minus the expected total spend
    const expectedFinalBalance = initialBalance - expectedTotalSpend;

    console.log('Initial balance:', initialBalance);
    console.log('Expected total spend:', expectedTotalSpend);
    console.log('Expected final balance:', expectedFinalBalance);
    console.log('Actual final balance:', finalBalance.tokenCredits);

    // Allow for small rounding differences
    expect(finalBalance.tokenCredits).toBeCloseTo(expectedFinalBalance, 0);

    // Verify all transactions were created
    const transactions = await Transaction.find({
      user: userId,
      conversationId,
    });

    // We should have 2 transactions (prompt + completion) for each usage record
    // Some might be structured, some regular
    expect(transactions.length).toBeGreaterThanOrEqual(collectedUsage.length);

    // Log transaction details for debugging
    console.log('Transaction summary:');
    let totalTokenValue = 0;
    transactions.forEach((tx) => {
      console.log(`${tx.tokenType}: rawAmount=${tx.rawAmount}, tokenValue=${tx.tokenValue}`);
      totalTokenValue += tx.tokenValue;
    });
    console.log('Total token value from transactions:', totalTokenValue);

    // The difference between expected and actual is significant
    // This is likely due to the multipliers being different in the test environment
    // Let's adjust our expectation based on the actual transactions
    const actualSpend = initialBalance - finalBalance.tokenCredits;
    console.log('Actual spend:', actualSpend);

    // Instead of checking the exact balance, let's verify that:
    // 1. The balance was reduced (tokens were spent)
    expect(finalBalance.tokenCredits).toBeLessThan(initialBalance);
    // 2. The total token value from transactions matches the actual spend
    expect(Math.abs(totalTokenValue)).toBeCloseTo(actualSpend, -3); // Allow for larger differences
  });

  // Add this new test case
  it('should handle multiple concurrent balance increases correctly', async () => {
    // Start with zero balance
    const initialBalance = 0;
    await Balance.create({
      user: userId,
      tokenCredits: initialBalance,
    });

    const numberOfRefills = 25;
    const refillAmount = 1000;

    const promises = [];
    for (let i = 0; i < numberOfRefills; i++) {
      promises.push(
        createAutoRefillTransaction({
          user: userId,
          tokenType: 'credits',
          context: 'concurrent-refill-test',
          rawAmount: refillAmount,
          balance: { enabled: true },
        }),
      );
    }

    // Wait for all refill transactions to complete
    const results = await Promise.all(promises);

    // Verify final balance
    const finalBalance = await Balance.findOne({ user: userId });
    expect(finalBalance).toBeDefined();

    // The final balance should be the initial balance plus the sum of all refills
    const expectedFinalBalance = initialBalance + numberOfRefills * refillAmount;

    console.log('Initial balance (Increase Test):', initialBalance);
    console.log(`Performed ${numberOfRefills} refills of ${refillAmount} each.`);
    console.log('Expected final balance (Increase Test):', expectedFinalBalance);
    console.log('Actual final balance (Increase Test):', finalBalance.tokenCredits);

    // Use toBeCloseTo for safety, though toBe should work for integer math
    expect(finalBalance.tokenCredits).toBeCloseTo(expectedFinalBalance, 0);

    // Verify all transactions were created
    const transactions = await Transaction.find({
      user: userId,
      context: 'concurrent-refill-test',
    });

    // We should have one transaction for each refill attempt
    expect(transactions.length).toBe(numberOfRefills);

    // Optional: Verify the sum of increments from the results matches the balance change
    const totalIncrementReported = results.reduce((sum, result) => {
      // Assuming createAutoRefillTransaction returns an object with the increment amount
      // Adjust this based on the actual return structure.
      // Let's assume it returns { balance: newBalance, transaction: { rawAmount: ... } }
      // Or perhaps we check the transaction.rawAmount directly
      return sum + (result?.transaction?.rawAmount || 0);
    }, 0);
    console.log('Total increment reported by results:', totalIncrementReported);
    expect(totalIncrementReported).toBe(expectedFinalBalance - initialBalance);

    // Optional: Check the sum of tokenValue from saved transactions
    let totalTokenValueFromDb = 0;
    transactions.forEach((tx) => {
      // For refills, rawAmount is positive, and tokenValue might be calculated based on it
      // Let's assume tokenValue directly reflects the increment for simplicity here
      // If calculation is involved, adjust accordingly
      totalTokenValueFromDb += tx.rawAmount; // Or tx.tokenValue if that holds the increment
    });
    console.log('Total rawAmount from DB transactions:', totalTokenValueFromDb);
    expect(totalTokenValueFromDb).toBeCloseTo(expectedFinalBalance - initialBalance, 0);
  });

  it('should create structured transactions for both prompt and completion tokens', async () => {
    // Create a balance for the user
    await Balance.create({
      user: userId,
      tokenCredits: 10000,
    });

    const txData = {
      user: userId,
      conversationId: 'test-convo',
      model: 'claude-3-5-sonnet',
      context: 'test',
      balance: { enabled: true },
    };
    const tokenUsage = {
      promptTokens: {
        input: 10,
        write: 100,
        read: 5,
      },
      completionTokens: 50,
    };

    const result = await spendStructuredTokens(txData, tokenUsage);

    // Verify transactions were created
    const transactions = await Transaction.find({ user: userId }).sort({ tokenType: 1 });
    expect(transactions).toHaveLength(2);

    // Check completion transaction
    expect(transactions[0].tokenType).toBe('completion');
    expect(transactions[0].rawAmount).toBe(-50);

    // Check prompt transaction
    expect(transactions[1].tokenType).toBe('prompt');
    expect(transactions[1].inputTokens).toBe(-10);
    expect(transactions[1].writeTokens).toBe(-100);
    expect(transactions[1].readTokens).toBe(-5);

    // Verify result contains transaction info
    expect(result).toEqual({
      prompt: expect.objectContaining({
        user: userId.toString(),
        prompt: expect.any(Number),
      }),
      completion: expect.objectContaining({
        user: userId.toString(),
        completion: expect.any(Number),
      }),
    });

    // Verify balance was updated
    const balance = await Balance.findOne({ user: userId });
    expect(balance).toBeDefined();
    expect(balance.tokenCredits).toBeLessThan(10000); // Balance should be reduced
  });
});