/** * 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', }), ); }); }); });