diff --git a/api/server/middleware/abortMiddleware.js b/api/server/middleware/abortMiddleware.js index feed002b0e..96db0b60d4 100644 --- a/api/server/middleware/abortMiddleware.js +++ b/api/server/middleware/abortMiddleware.js @@ -6,6 +6,7 @@ const { countTokens, GenerationJobManager, recordCollectedUsage, + getTransactionsConfig, sanitizeMessageForTransmit, buildAbortedResponseMetadata, } = require('@librechat/api'); @@ -59,6 +60,7 @@ const isAbortError = (error) => { * @param {Array} params.collectedUsage - Usage metadata from all models * @param {string} [params.fallbackModel] - Fallback model name if not in usage * @param {string} [params.messageId] - The response message ID for transaction correlation + * @param {AppConfig['transactions']} [params.transactions] - Resolved transactions config */ async function spendCollectedUsage({ userId, @@ -66,6 +68,7 @@ async function spendCollectedUsage({ collectedUsage, fallbackModel, messageId, + transactions, }) { if (!collectedUsage || collectedUsage.length === 0) { return; @@ -85,6 +88,7 @@ async function spendCollectedUsage({ context: 'abort', messageId, model: fallbackModel, + transactions, }, ); @@ -149,6 +153,8 @@ async function abortMessage(req, res) { responseMessage.metadata = abortMetadata; } + const transactions = getTransactionsConfig(req.config); + // Spend tokens for ALL models from collectedUsage (handles parallel agents/addedConvo) if (collectedUsage && collectedUsage.length > 0) { await spendCollectedUsage({ @@ -157,11 +163,12 @@ async function abortMessage(req, res) { collectedUsage, fallbackModel: jobData?.model, messageId: jobData?.responseMessageId, + transactions, }); } else { // Fallback: no collected usage, use text-based token counting for primary model only await db.spendTokens( - { ...responseMessage, context: 'incomplete', user: userId }, + { ...responseMessage, context: 'incomplete', user: userId, transactions }, { promptTokens, completionTokens }, ); } diff --git a/api/server/middleware/abortMiddleware.spec.js b/api/server/middleware/abortMiddleware.spec.js index 06e434065a..20c4700421 100644 --- a/api/server/middleware/abortMiddleware.spec.js +++ b/api/server/middleware/abortMiddleware.spec.js @@ -19,6 +19,7 @@ const mockRecordCollectedUsage = jest const mockGetMultiplier = jest.fn().mockReturnValue(1); const mockGetCacheMultiplier = jest.fn().mockReturnValue(null); +const mockGetTransactionsConfig = jest.fn().mockReturnValue({ enabled: false }); jest.mock('@librechat/data-schemas', () => ({ logger: { @@ -37,7 +38,9 @@ jest.mock('@librechat/api', () => ({ abortJob: jest.fn(), }, recordCollectedUsage: mockRecordCollectedUsage, + getTransactionsConfig: (...args) => mockGetTransactionsConfig(...args), sanitizeMessageForTransmit: jest.fn((msg) => msg), + buildAbortedResponseMetadata: jest.fn().mockReturnValue(null), })); jest.mock('librechat-data-provider', () => ({ @@ -75,7 +78,9 @@ jest.mock('./abortRun', () => ({ const { logger } = require('@librechat/data-schemas'); const { sendError } = require('~/server/middleware/error'); -const { handleAbortError, spendCollectedUsage } = require('./abortMiddleware'); +const { GenerationJobManager } = require('@librechat/api'); +const db = require('~/models'); +const { handleAbort, handleAbortError, spendCollectedUsage } = require('./abortMiddleware'); const buildAbortRequest = () => ({ body: { @@ -310,3 +315,107 @@ describe('abortMiddleware - handleAbortError', () => { expect(sendError).toHaveBeenCalledTimes(1); }); }); + +/** + * The transactions config is resolved from the request's app config and must reach + * every write path in this file. `createTransaction` reads `transactions` from the + * caller-supplied data, so an omitted value is indistinguishable from enabled and + * the write proceeds even when `transactions.enabled` is false. + */ +describe('abortMiddleware - transactions config', () => { + const buildJobData = () => ({ + model: 'gpt-4', + responseMessageId: 'msg-123', + conversationId: 'convo-123', + endpoint: 'agents', + sender: 'AI', + promptTokens: 25, + userMessage: { + messageId: 'user-msg-123', + parentMessageId: 'parent-123', + conversationId: 'convo-123', + text: 'hello', + }, + }); + + const buildReq = () => ({ + body: { abortKey: 'convo-123:1', endpoint: 'agents' }, + user: { id: 'user-123', email: 'user@example.com' }, + config: { transactions: { enabled: false } }, + }); + + const buildRes = () => ({ + headersSent: false, + setHeader: jest.fn(), + send: jest.fn(), + }); + + beforeEach(() => { + jest.clearAllMocks(); + mockGetTransactionsConfig.mockReturnValue({ enabled: false }); + mockRecordCollectedUsage.mockResolvedValue({ input_tokens: 100, output_tokens: 50 }); + db.getConvo.mockResolvedValue({ title: 'Test Chat' }); + }); + + it('forwards transactions through spendCollectedUsage to recordCollectedUsage', async () => { + const collectedUsage = [{ input_tokens: 100, output_tokens: 50, model: 'gpt-4' }]; + + await spendCollectedUsage({ + userId: 'user-123', + conversationId: 'convo-123', + collectedUsage, + fallbackModel: 'gpt-4', + transactions: { enabled: false }, + }); + + expect(mockRecordCollectedUsage).toHaveBeenCalledTimes(1); + expect(mockRecordCollectedUsage).toHaveBeenCalledWith( + expect.any(Object), + expect.objectContaining({ context: 'abort', transactions: { enabled: false } }), + ); + }); + + it('resolves the config from req and forwards it on the collected-usage path', async () => { + const collectedUsage = [{ input_tokens: 100, output_tokens: 50, model: 'gpt-4' }]; + GenerationJobManager.abortJob.mockResolvedValue({ + success: true, + jobData: buildJobData(), + content: [], + text: 'partial', + collectedUsage, + }); + + const req = buildReq(); + await handleAbort()(req, buildRes()); + + expect(logger.error).not.toHaveBeenCalled(); + expect(mockGetTransactionsConfig).toHaveBeenCalledWith(req.config); + expect(mockRecordCollectedUsage).toHaveBeenCalledTimes(1); + expect(mockRecordCollectedUsage).toHaveBeenCalledWith( + expect.any(Object), + expect.objectContaining({ context: 'abort', transactions: { enabled: false } }), + ); + }); + + it('resolves the config from req and forwards it on the token-count fallback path', async () => { + GenerationJobManager.abortJob.mockResolvedValue({ + success: true, + jobData: buildJobData(), + content: [], + text: 'partial', + collectedUsage: [], + }); + + const req = buildReq(); + await handleAbort()(req, buildRes()); + + expect(logger.error).not.toHaveBeenCalled(); + expect(mockGetTransactionsConfig).toHaveBeenCalledWith(req.config); + expect(mockRecordCollectedUsage).not.toHaveBeenCalled(); + expect(mockSpendTokens).toHaveBeenCalledTimes(1); + expect(mockSpendTokens).toHaveBeenCalledWith( + expect.objectContaining({ context: 'incomplete', transactions: { enabled: false } }), + expect.any(Object), + ); + }); +});