/** * Tests for abortMiddleware - spendCollectedUsage function * * This tests the token spending logic for abort scenarios, * particularly for parallel agents (addedConvo) where multiple * models need their tokens spent. * * spendCollectedUsage delegates to recordCollectedUsage from @librechat/api, * passing pricing + bulkWriteOps deps, with context: 'abort'. * After spending, it clears the collectedUsage array to prevent double-spending * from the AgentClient finally block (which shares the same array reference). */ const mockSpendTokens = jest.fn().mockResolvedValue(); const mockSpendStructuredTokens = jest.fn().mockResolvedValue(); const mockRecordCollectedUsage = jest .fn() .mockResolvedValue({ input_tokens: 100, output_tokens: 50 }); const mockGetMultiplier = jest.fn().mockReturnValue(1); const mockGetCacheMultiplier = jest.fn().mockReturnValue(null); jest.mock('@librechat/data-schemas', () => ({ logger: { debug: jest.fn(), error: jest.fn(), warn: jest.fn(), info: jest.fn(), }, })); jest.mock('@librechat/api', () => ({ countTokens: jest.fn().mockResolvedValue(100), isEnabled: jest.fn().mockReturnValue(false), sendEvent: jest.fn(), GenerationJobManager: { abortJob: jest.fn(), }, recordCollectedUsage: mockRecordCollectedUsage, sanitizeMessageForTransmit: jest.fn((msg) => msg), })); jest.mock('librechat-data-provider', () => ({ isAssistantsEndpoint: jest.fn().mockReturnValue(false), ErrorTypes: { INVALID_REQUEST: 'INVALID_REQUEST', NO_SYSTEM_MESSAGES: 'NO_SYSTEM_MESSAGES' }, })); jest.mock('~/app/clients/prompts', () => ({ truncateText: jest.fn((text) => text), smartTruncateText: jest.fn((text) => text), })); jest.mock('~/cache/clearPendingReq', () => jest.fn().mockResolvedValue()); jest.mock('~/server/middleware/error', () => ({ sendError: jest.fn(), })); const mockUpdateBalance = jest.fn().mockResolvedValue({}); const mockBulkInsertTransactions = jest.fn().mockResolvedValue(undefined); jest.mock('~/models', () => ({ saveMessage: jest.fn().mockResolvedValue(), getConvo: jest.fn().mockResolvedValue({ title: 'Test Chat' }), updateBalance: mockUpdateBalance, bulkInsertTransactions: mockBulkInsertTransactions, spendTokens: (...args) => mockSpendTokens(...args), spendStructuredTokens: (...args) => mockSpendStructuredTokens(...args), getMultiplier: mockGetMultiplier, getCacheMultiplier: mockGetCacheMultiplier, })); jest.mock('./abortRun', () => ({ abortRun: jest.fn(), })); const { logger } = require('@librechat/data-schemas'); const { sendError } = require('~/server/middleware/error'); const { handleAbortError, spendCollectedUsage } = require('./abortMiddleware'); const buildAbortRequest = () => ({ body: { model: 'gpt-4', }, user: { id: 'user-123', }, }); describe('abortMiddleware - spendCollectedUsage', () => { beforeEach(() => { jest.clearAllMocks(); }); describe('spendCollectedUsage delegation', () => { it('should return early if collectedUsage is empty', async () => { await spendCollectedUsage({ userId: 'user-123', conversationId: 'convo-123', collectedUsage: [], fallbackModel: 'gpt-4', }); expect(mockRecordCollectedUsage).not.toHaveBeenCalled(); }); it('should return early if collectedUsage is null', async () => { await spendCollectedUsage({ userId: 'user-123', conversationId: 'convo-123', collectedUsage: null, fallbackModel: 'gpt-4', }); expect(mockRecordCollectedUsage).not.toHaveBeenCalled(); }); it('should call recordCollectedUsage with abort context and full deps', async () => { const collectedUsage = [{ input_tokens: 100, output_tokens: 50, model: 'gpt-4' }]; await spendCollectedUsage({ userId: 'user-123', conversationId: 'convo-123', collectedUsage, fallbackModel: 'gpt-4', messageId: 'msg-123', }); expect(mockRecordCollectedUsage).toHaveBeenCalledTimes(1); expect(mockRecordCollectedUsage).toHaveBeenCalledWith( { spendTokens: expect.any(Function), spendStructuredTokens: expect.any(Function), pricing: { getMultiplier: mockGetMultiplier, getCacheMultiplier: mockGetCacheMultiplier, }, bulkWriteOps: { insertMany: mockBulkInsertTransactions, updateBalance: mockUpdateBalance, }, }, { user: 'user-123', conversationId: 'convo-123', collectedUsage, context: 'abort', messageId: 'msg-123', model: 'gpt-4', }, ); }); it('should pass context abort for multiple models (parallel agents)', async () => { const collectedUsage = [ { input_tokens: 100, output_tokens: 50, model: 'gpt-4' }, { input_tokens: 80, output_tokens: 40, model: 'claude-3' }, { input_tokens: 120, output_tokens: 60, model: 'gemini-pro' }, ]; await spendCollectedUsage({ userId: 'user-123', conversationId: 'convo-123', collectedUsage, fallbackModel: 'gpt-4', }); expect(mockRecordCollectedUsage).toHaveBeenCalledTimes(1); expect(mockRecordCollectedUsage).toHaveBeenCalledWith( expect.any(Object), expect.objectContaining({ context: 'abort', collectedUsage, }), ); }); it('should handle real-world parallel agent abort scenario', async () => { const collectedUsage = [ { input_tokens: 31596, output_tokens: 151, model: 'gemini-3-flash-preview' }, { input_tokens: 28000, output_tokens: 120, model: 'gpt-5.2' }, ]; await spendCollectedUsage({ userId: 'user-123', conversationId: 'convo-123', collectedUsage, fallbackModel: 'gemini-3-flash-preview', }); expect(mockRecordCollectedUsage).toHaveBeenCalledTimes(1); expect(mockRecordCollectedUsage).toHaveBeenCalledWith( expect.any(Object), expect.objectContaining({ user: 'user-123', conversationId: 'convo-123', context: 'abort', model: 'gemini-3-flash-preview', }), ); }); /** * Race condition prevention: after abort middleware spends tokens, * the collectedUsage array is cleared so AgentClient.recordCollectedUsage() * (which shares the same array reference) sees an empty array and returns early. */ it('should clear collectedUsage array after spending to prevent double-spending', async () => { const collectedUsage = [ { input_tokens: 100, output_tokens: 50, model: 'gpt-4' }, { input_tokens: 80, output_tokens: 40, model: 'claude-3' }, ]; expect(collectedUsage.length).toBe(2); await spendCollectedUsage({ userId: 'user-123', conversationId: 'convo-123', collectedUsage, fallbackModel: 'gpt-4', }); expect(mockRecordCollectedUsage).toHaveBeenCalledTimes(1); expect(collectedUsage.length).toBe(0); }); it('should await recordCollectedUsage before clearing array', async () => { let resolved = false; mockRecordCollectedUsage.mockImplementation(async () => { await new Promise((resolve) => setTimeout(resolve, 10)); resolved = true; return { input_tokens: 100, output_tokens: 50 }; }); const collectedUsage = [ { input_tokens: 100, output_tokens: 50, model: 'gpt-4' }, { input_tokens: 80, output_tokens: 40, model: 'claude-3' }, ]; await spendCollectedUsage({ userId: 'user-123', conversationId: 'convo-123', collectedUsage, fallbackModel: 'gpt-4', }); expect(resolved).toBe(true); expect(collectedUsage.length).toBe(0); }); }); }); describe('abortMiddleware - handleAbortError', () => { beforeEach(() => { jest.clearAllMocks(); }); it.each([ [ 'native DOMException AbortError', new DOMException('The operation was aborted', 'AbortError'), 'AbortError', ], [ 'wrapped AbortError message', new Error('SSE stream disconnected: AbortError: The operation was aborted'), 'Error', ], [ 'cause-nested AbortError', new Error('Request failed', { cause: new DOMException('The operation was aborted', 'AbortError'), }), 'Error', ], ])('logs a %s as a debug event instead of an error', async (_label, error, name) => { await handleAbortError({}, buildAbortRequest(), error, { sender: 'AI', conversationId: 'convo-123', messageId: 'message-123', parentMessageId: 'parent-123', userMessageId: 'user-message-123', }); expect(logger.error).not.toHaveBeenCalled(); expect(logger.debug).toHaveBeenCalledWith('[handleAbortError] AI response aborted by user', { conversationId: 'convo-123', code: error.code, name, message: error.message, }); expect(sendError).toHaveBeenCalledTimes(1); }); it('keeps unexpected generation errors classified as errors', async () => { const error = new Error('Provider failed'); await handleAbortError({}, buildAbortRequest(), error, { sender: 'AI', conversationId: 'convo-123', messageId: 'message-123', parentMessageId: 'parent-123', userMessageId: 'user-message-123', }); expect(logger.error).toHaveBeenCalledWith( '[handleAbortError] AI response error; aborting request:', error, ); expect(logger.debug).not.toHaveBeenCalled(); expect(sendError).toHaveBeenCalledTimes(1); }); });