📉 feat: Add Token Usage Tracking for Agents API Routes (#11600)

* feat: Implement token usage tracking for OpenAI and Responses controllers

- Added functionality to record token usage against user balances in OpenAIChatCompletionController and createResponse functions.
- Introduced new utility functions for managing token spending and structured token usage.
- Enhanced error handling for token recording to improve logging and debugging capabilities.
- Updated imports to include new usage tracking methods and configurations.

* test: Add unit tests for recordCollectedUsage function in usage.spec.ts

- Introduced comprehensive tests for the recordCollectedUsage function, covering various scenarios including handling empty and null collectedUsage, single and multiple usage entries, and sequential and parallel execution cases.
- Enhanced token handling tests to ensure correct calculations for both OpenAI and Anthropic formats, including cache token management.
- Improved overall test coverage for usage tracking functionality, ensuring robust validation of expected behaviors and outcomes.

* test: Add unit tests for OpenAI and Responses API controllers

- Introduced comprehensive unit tests for the OpenAIChatCompletionController and createResponse functions, focusing on the correct invocation of recordCollectedUsage for token spending.
- Enhanced tests to validate the passing of balance and transactions configuration to the recordCollectedUsage function.
- Ensured proper dependency injection of spendTokens and spendStructuredTokens in the usage recording process.
- Improved overall test coverage for token usage tracking, ensuring robust validation of expected behaviors and outcomes.

api/server/controllers/agents/__tests__/openai.spec.js

@@ -0,0 +1,207 @@
+/**
+ * Unit tests for OpenAI-compatible API controller
+ * Tests that recordCollectedUsage is called correctly for token spending
+ */
+
+const mockSpendTokens = jest.fn().mockResolvedValue({});
+const mockSpendStructuredTokens = jest.fn().mockResolvedValue({});
+const mockRecordCollectedUsage = jest
+  .fn()
+  .mockResolvedValue({ input_tokens: 100, output_tokens: 50 });
+const mockGetBalanceConfig = jest.fn().mockReturnValue({ enabled: true });
+const mockGetTransactionsConfig = jest.fn().mockReturnValue({ enabled: true });
+
+jest.mock('nanoid', () => ({
+  nanoid: jest.fn(() => 'mock-nanoid-123'),
+}));
+
+jest.mock('@librechat/data-schemas', () => ({
+  logger: {
+    debug: jest.fn(),
+    error: jest.fn(),
+    warn: jest.fn(),
+  },
+}));
+
+jest.mock('@librechat/agents', () => ({
+  Callback: { TOOL_ERROR: 'TOOL_ERROR' },
+  ToolEndHandler: jest.fn(),
+  formatAgentMessages: jest.fn().mockReturnValue({
+    messages: [],
+    indexTokenCountMap: {},
+  }),
+  ChatModelStreamHandler: jest.fn().mockImplementation(() => ({
+    handle: jest.fn(),
+  })),
+}));
+
+jest.mock('@librechat/api', () => ({
+  writeSSE: jest.fn(),
+  createRun: jest.fn().mockResolvedValue({
+    processStream: jest.fn().mockResolvedValue(undefined),
+  }),
+  createChunk: jest.fn().mockReturnValue({}),
+  buildToolSet: jest.fn().mockReturnValue(new Set()),
+  sendFinalChunk: jest.fn(),
+  createSafeUser: jest.fn().mockReturnValue({ id: 'user-123' }),
+  validateRequest: jest
+    .fn()
+    .mockReturnValue({ request: { model: 'agent-123', messages: [], stream: false } }),
+  initializeAgent: jest.fn().mockResolvedValue({
+    model: 'gpt-4',
+    model_parameters: {},
+    toolRegistry: {},
+  }),
+  getBalanceConfig: mockGetBalanceConfig,
+  createErrorResponse: jest.fn(),
+  getTransactionsConfig: mockGetTransactionsConfig,
+  recordCollectedUsage: mockRecordCollectedUsage,
+  buildNonStreamingResponse: jest.fn().mockReturnValue({ id: 'resp-123' }),
+  createOpenAIStreamTracker: jest.fn().mockReturnValue({
+    addText: jest.fn(),
+    addReasoning: jest.fn(),
+    toolCalls: new Map(),
+    usage: { promptTokens: 0, completionTokens: 0, reasoningTokens: 0 },
+  }),
+  createOpenAIContentAggregator: jest.fn().mockReturnValue({
+    addText: jest.fn(),
+    addReasoning: jest.fn(),
+    getText: jest.fn().mockReturnValue(''),
+    getReasoning: jest.fn().mockReturnValue(''),
+    toolCalls: new Map(),
+    usage: { promptTokens: 100, completionTokens: 50, reasoningTokens: 0 },
+  }),
+  createToolExecuteHandler: jest.fn().mockReturnValue({ handle: jest.fn() }),
+  isChatCompletionValidationFailure: jest.fn().mockReturnValue(false),
+}));
+
+jest.mock('~/server/services/ToolService', () => ({
+  loadAgentTools: jest.fn().mockResolvedValue([]),
+  loadToolsForExecution: jest.fn().mockResolvedValue([]),
+}));
+
+jest.mock('~/models/spendTokens', () => ({
+  spendTokens: mockSpendTokens,
+  spendStructuredTokens: mockSpendStructuredTokens,
+}));
+
+jest.mock('~/server/controllers/agents/callbacks', () => ({
+  createToolEndCallback: jest.fn().mockReturnValue(jest.fn()),
+}));
+
+jest.mock('~/server/services/PermissionService', () => ({
+  findAccessibleResources: jest.fn().mockResolvedValue([]),
+}));
+
+jest.mock('~/models/Conversation', () => ({
+  getConvoFiles: jest.fn().mockResolvedValue([]),
+}));
+
+jest.mock('~/models/Agent', () => ({
+  getAgent: jest.fn().mockResolvedValue({
+    id: 'agent-123',
+    provider: 'openAI',
+    model_parameters: { model: 'gpt-4' },
+  }),
+  getAgents: jest.fn().mockResolvedValue([]),
+}));
+
+jest.mock('~/models', () => ({
+  getFiles: jest.fn(),
+  getUserKey: jest.fn(),
+  getMessages: jest.fn(),
+  updateFilesUsage: jest.fn(),
+  getUserKeyValues: jest.fn(),
+  getUserCodeFiles: jest.fn(),
+  getToolFilesByIds: jest.fn(),
+  getCodeGeneratedFiles: jest.fn(),
+}));
+
+describe('OpenAIChatCompletionController', () => {
+  let OpenAIChatCompletionController;
+  let req, res;
+
+  beforeEach(() => {
+    jest.clearAllMocks();
+
+    const controller = require('../openai');
+    OpenAIChatCompletionController = controller.OpenAIChatCompletionController;
+
+    req = {
+      body: {
+        model: 'agent-123',
+        messages: [{ role: 'user', content: 'Hello' }],
+        stream: false,
+      },
+      user: { id: 'user-123' },
+      config: {
+        endpoints: {
+          agents: { allowedProviders: ['openAI'] },
+        },
+      },
+      on: jest.fn(),
+    };
+
+    res = {
+      status: jest.fn().mockReturnThis(),
+      json: jest.fn(),
+      setHeader: jest.fn(),
+      flushHeaders: jest.fn(),
+      end: jest.fn(),
+      write: jest.fn(),
+    };
+  });
+
+  describe('token usage recording', () => {
+    it('should call recordCollectedUsage after successful non-streaming completion', async () => {
+      await OpenAIChatCompletionController(req, res);
+
+      expect(mockRecordCollectedUsage).toHaveBeenCalledTimes(1);
+      expect(mockRecordCollectedUsage).toHaveBeenCalledWith(
+        { spendTokens: mockSpendTokens, spendStructuredTokens: mockSpendStructuredTokens },
+        expect.objectContaining({
+          user: 'user-123',
+          conversationId: expect.any(String),
+          collectedUsage: expect.any(Array),
+          context: 'message',
+          balance: { enabled: true },
+          transactions: { enabled: true },
+        }),
+      );
+    });
+
+    it('should pass balance and transactions config to recordCollectedUsage', async () => {
+      mockGetBalanceConfig.mockReturnValue({ enabled: true, startBalance: 1000 });
+      mockGetTransactionsConfig.mockReturnValue({ enabled: true, rateLimit: 100 });
+
+      await OpenAIChatCompletionController(req, res);
+
+      expect(mockRecordCollectedUsage).toHaveBeenCalledWith(
+        expect.any(Object),
+        expect.objectContaining({
+          balance: { enabled: true, startBalance: 1000 },
+          transactions: { enabled: true, rateLimit: 100 },
+        }),
+      );
+    });
+
+    it('should pass spendTokens and spendStructuredTokens as dependencies', async () => {
+      await OpenAIChatCompletionController(req, res);
+
+      const [deps] = mockRecordCollectedUsage.mock.calls[0];
+      expect(deps).toHaveProperty('spendTokens', mockSpendTokens);
+      expect(deps).toHaveProperty('spendStructuredTokens', mockSpendStructuredTokens);
+    });
+
+    it('should include model from primaryConfig in recordCollectedUsage params', async () => {
+      await OpenAIChatCompletionController(req, res);
+
+      expect(mockRecordCollectedUsage).toHaveBeenCalledWith(
+        expect.any(Object),
+        expect.objectContaining({
+          model: 'gpt-4',
+        }),
+      );
+    });
+  });
+});