🧮 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/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,
+      }),
+    });
+  });
+});