📉 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/responses.js

@@ -13,6 +13,9 @@ const {
   buildToolSet,
   createSafeUser,
   initializeAgent,
+  getBalanceConfig,
+  recordCollectedUsage,
+  getTransactionsConfig,
   createToolExecuteHandler,
   // Responses API
   writeDone,
@@ -39,6 +42,7 @@ const {
 const { loadAgentTools, loadToolsForExecution } = require('~/server/services/ToolService');
 const { findAccessibleResources } = require('~/server/services/PermissionService');
 const { getConvoFiles, saveConvo, getConvo } = require('~/models/Conversation');
+const { spendTokens, spendStructuredTokens } = require('~/models/spendTokens');
 const { getAgent, getAgents } = require('~/models/Agent');
 const db = require('~/models');
 
@@ -403,6 +407,9 @@ const createResponse = async (req, res) => {
       const { handlers: responsesHandlers, finalizeStream } =
         createResponsesEventHandlers(handlerConfig);
 
+      // Collect usage for balance tracking
+      const collectedUsage = [];
+
       // Built-in handler for processing raw model stream chunks
       const chatModelStreamHandler = new ChatModelStreamHandler();
 
@@ -445,7 +452,15 @@ const createResponse = async (req, res) => {
         on_reasoning_delta: responsesHandlers.on_reasoning_delta,
         on_run_step: responsesHandlers.on_run_step,
         on_run_step_delta: responsesHandlers.on_run_step_delta,
-        on_chat_model_end: responsesHandlers.on_chat_model_end,
+        on_chat_model_end: {
+          handle: (event, data) => {
+            responsesHandlers.on_chat_model_end.handle(event, data);
+            const usage = data?.output?.usage_metadata;
+            if (usage) {
+              collectedUsage.push(usage);
+            }
+          },
+        },
         on_tool_end: new ToolEndHandler(toolEndCallback, logger),
         on_run_step_completed: { handle: () => {} },
         on_chain_stream: { handle: () => {} },
@@ -499,6 +514,24 @@ const createResponse = async (req, res) => {
         },
       });
 
+      // Record token usage against balance
+      const balanceConfig = getBalanceConfig(req.config);
+      const transactionsConfig = getTransactionsConfig(req.config);
+      recordCollectedUsage(
+        { spendTokens, spendStructuredTokens },
+        {
+          user: userId,
+          conversationId,
+          collectedUsage,
+          context: 'message',
+          balance: balanceConfig,
+          transactions: transactionsConfig,
+          model: primaryConfig.model || agent.model_parameters?.model,
+        },
+      ).catch((err) => {
+        logger.error('[Responses API] Error recording usage:', err);
+      });
+
       // Finalize the stream
       finalizeStream();
       res.end();
@@ -539,6 +572,9 @@ const createResponse = async (req, res) => {
 
       const chatModelStreamHandler = new ChatModelStreamHandler();
 
+      // Collect usage for balance tracking
+      const collectedUsage = [];
+
       /** @type {Promise<import('librechat-data-provider').TAttachment | null>[]} */
       const artifactPromises = [];
       const toolEndCallback = createToolEndCallback({ req, res, artifactPromises, streamId: null });
@@ -569,7 +605,15 @@ const createResponse = async (req, res) => {
         on_reasoning_delta: aggregatorHandlers.on_reasoning_delta,
         on_run_step: aggregatorHandlers.on_run_step,
         on_run_step_delta: aggregatorHandlers.on_run_step_delta,
-        on_chat_model_end: aggregatorHandlers.on_chat_model_end,
+        on_chat_model_end: {
+          handle: (event, data) => {
+            aggregatorHandlers.on_chat_model_end.handle(event, data);
+            const usage = data?.output?.usage_metadata;
+            if (usage) {
+              collectedUsage.push(usage);
+            }
+          },
+        },
         on_tool_end: new ToolEndHandler(toolEndCallback, logger),
         on_run_step_completed: { handle: () => {} },
         on_chain_stream: { handle: () => {} },
@@ -621,6 +665,24 @@ const createResponse = async (req, res) => {
         },
       });
 
+      // Record token usage against balance
+      const balanceConfig = getBalanceConfig(req.config);
+      const transactionsConfig = getTransactionsConfig(req.config);
+      recordCollectedUsage(
+        { spendTokens, spendStructuredTokens },
+        {
+          user: userId,
+          conversationId,
+          collectedUsage,
+          context: 'message',
+          balance: balanceConfig,
+          transactions: transactionsConfig,
+          model: primaryConfig.model || agent.model_parameters?.model,
+        },
+      ).catch((err) => {
+        logger.error('[Responses API] Error recording usage:', err);
+      });
+
       if (artifactPromises.length > 0) {
         try {
           await Promise.all(artifactPromises);