diff --git a/api/server/controllers/agents/__tests__/usageEvents.live.spec.js b/api/server/controllers/agents/__tests__/usageEvents.live.spec.js new file mode 100644 index 0000000000..ff02c9325f --- /dev/null +++ b/api/server/controllers/agents/__tests__/usageEvents.live.spec.js @@ -0,0 +1,161 @@ +/** + * Live host-layer verification: real Anthropic run through the actual + * getDefaultHandlers pipeline, asserting the SSE usage/context events the + * client consumes and their resume persistence. + * + * Run with: + * RUN_USAGE_LIVE_TESTS=1 ANTHROPIC_API_KEY=... npx jest usageEvents.live --runInBand + */ +const { HumanMessage } = require('@langchain/core/messages'); +const { Run, Providers, GraphEvents } = require('@librechat/agents'); +const { GenerationJobManager } = require('@librechat/api'); +const { getDefaultHandlers } = require('~/server/controllers/agents/callbacks'); + +jest.mock('nanoid', () => ({ + nanoid: jest.fn(() => 'mock-nanoid'), +})); + +jest.mock('~/server/services/Files/Citations', () => ({ + processFileCitations: jest.fn(), +})); + +jest.mock('~/server/services/Files/Code/process', () => ({ + processCodeOutput: jest.fn(), + runPreviewFinalize: jest.fn(), +})); + +jest.mock('~/server/services/Files/process', () => ({ + saveBase64Image: jest.fn(), +})); + +const shouldRunLive = + process.env.RUN_USAGE_LIVE_TESTS === '1' && + process.env.ANTHROPIC_API_KEY != null && + process.env.ANTHROPIC_API_KEY !== ''; + +const describeIfLive = shouldRunLive ? describe : describe.skip; +const modelName = process.env.ANTHROPIC_USAGE_LIVE_MODEL ?? 'claude-haiku-4-5'; +const hasContextUsageEvent = GraphEvents.ON_CONTEXT_USAGE != null; + +const charCounter = (msg) => { + const content = msg.content; + if (typeof content === 'string') { + return Math.ceil(content.length / 4) + 3; + } + if (Array.isArray(content)) { + let length = 3; + for (const part of content) { + if (typeof part === 'string') { + length += Math.ceil(part.length / 4); + } else if (typeof part?.text === 'string') { + length += Math.ceil(part.text.length / 4); + } + } + return length; + } + return 3; +}; + +function createMockRes() { + const events = []; + return { + events, + headersSent: true, + writableEnded: false, + write(payload) { + for (const line of String(payload).split('\n')) { + if (line.startsWith('data: ')) { + events.push(JSON.parse(line.slice(6))); + } + } + return true; + }, + }; +} + +describeIfLive('live usage events through the host pipeline', () => { + jest.setTimeout(120000); + + afterAll(async () => { + await GenerationJobManager.destroy(); + }); + + test('streams real provider usage and persists it for resume', async () => { + const streamId = `usage-live-${Date.now()}`; + await GenerationJobManager.createJob(streamId, 'user-live', 'convo-live'); + + /** streamId mode routes events through the job emitter — capture them + * as a subscribed resumable client would, not via res.write */ + const res = createMockRes(); + await GenerationJobManager.subscribe(streamId, (event) => { + res.events.push(event); + }); + const collectedUsage = []; + const handlers = getDefaultHandlers({ + res, + aggregateContent: () => {}, + toolEndCallback: () => {}, + collectedUsage, + streamId, + }); + + const run = await Run.create({ + runId: 'usage-live-response', + graphConfig: { + type: 'standard', + llmConfig: { + provider: Providers.ANTHROPIC, + model: modelName, + apiKey: process.env.ANTHROPIC_API_KEY, + temperature: 0, + maxTokens: 64, + streaming: true, + streamUsage: true, + }, + instructions: 'You are concise. Reply with one short sentence.', + maxContextTokens: 8000, + }, + returnContent: true, + customHandlers: handlers, + tokenCounter: charCounter, + indexTokenCountMap: {}, + }); + + await run.processStream( + { messages: [new HumanMessage('Say hello in five words or fewer.')] }, + { + configurable: { thread_id: 'usage-live-thread', user_id: 'user-live' }, + streamMode: 'values', + version: 'v2', + }, + ); + + const usageEvents = res.events.filter((e) => e.event === 'on_token_usage'); + expect(usageEvents).toHaveLength(1); + const usage = usageEvents[0].data; + expect(usage.input_tokens).toBeGreaterThan(0); + expect(usage.output_tokens).toBeGreaterThan(0); + expect(usage.provider).toBe(Providers.ANTHROPIC); + expect(usage.model).toBe(modelName); + expect(collectedUsage).toHaveLength(1); + expect(usage.input_tokens).toBe(collectedUsage[0].input_tokens); + + if (hasContextUsageEvent) { + const contextEvents = res.events.filter((e) => e.event === 'on_context_usage'); + expect(contextEvents).toHaveLength(1); + const snapshot = contextEvents[0].data; + expect(snapshot.breakdown.maxContextTokens).toBe(8000); + const estimatedUsed = snapshot.contextBudget - snapshot.remainingContextTokens; + expect(estimatedUsed).toBeGreaterThan(0); + expect(estimatedUsed / usage.input_tokens).toBeGreaterThan(0.2); + expect(estimatedUsed / usage.input_tokens).toBeLessThan(5); + } + + const resumeState = await GenerationJobManager.getResumeState(streamId); + expect(resumeState.collectedUsage).toHaveLength(1); + expect(resumeState.collectedUsage[0].input_tokens).toBe(usage.input_tokens); + if (hasContextUsageEvent) { + expect(resumeState.contextUsage.breakdown.maxContextTokens).toBe(8000); + } + }); +});