From 03b477a84c73b3bc9af9370cf2c12b6f4160c8b5 Mon Sep 17 00:00:00 2001 From: Ravi Kumar L Date: Sat, 23 May 2026 03:19:45 +0200 Subject: [PATCH] =?UTF-8?q?=F0=9F=96=B2=EF=B8=8F=20feat:=20Trace=20SSE=20S?= =?UTF-8?q?tream=20Lifecycle=20with=20OTel=20(#13266)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- api/server/routes/agents/index.js | 31 ++-- packages/api/src/telemetry/index.ts | 2 + packages/api/src/telemetry/stream.spec.ts | 208 ++++++++++++++++++++++ packages/api/src/telemetry/stream.ts | 145 +++++++++++++++ 4 files changed, 373 insertions(+), 13 deletions(-) create mode 100644 packages/api/src/telemetry/stream.spec.ts create mode 100644 packages/api/src/telemetry/stream.ts diff --git a/api/server/routes/agents/index.js b/api/server/routes/agents/index.js index bbb39f5d2c..f7076ce153 100644 --- a/api/server/routes/agents/index.js +++ b/api/server/routes/agents/index.js @@ -1,5 +1,6 @@ const express = require('express'); const { isEnabled, GenerationJobManager } = require('@librechat/api'); +const { createSseStreamTelemetry } = require('@librechat/api/telemetry'); const { logger } = require('@librechat/data-schemas'); const { uaParser, @@ -76,35 +77,43 @@ router.get('/chat/stream/:streamId', async (req, res) => { return res.status(403).json({ error: 'Unauthorized' }); } + const streamTelemetry = createSseStreamTelemetry({ req, res, streamId, isResume }); + res.setHeader('Content-Encoding', 'identity'); res.setHeader('Content-Type', 'text/event-stream'); res.setHeader('Cache-Control', 'no-cache, no-transform'); res.setHeader('Connection', 'keep-alive'); res.setHeader('X-Accel-Buffering', 'no'); res.flushHeaders(); + streamTelemetry.recordHeadersFlushed(); logger.debug(`[AgentStream] Client subscribed to ${streamId}, resume: ${isResume}`); - const writeEvent = (event) => { + const writeEvent = (event, options = {}) => { if (!res.writableEnded) { - res.write(`event: message\ndata: ${JSON.stringify(event)}\n\n`); + const eventName = options.eventName ?? 'message'; + const payload = `event: ${eventName}\ndata: ${JSON.stringify(event)}\n\n`; + res.write(payload); + streamTelemetry.recordWrite(payload, { final: options.final }); if (typeof res.flush === 'function') { res.flush(); } + return true; } + + return false; }; const onDone = (event) => { - writeEvent(event); + streamTelemetry.recordFinalEventEmitted(); + writeEvent(event, { final: true }); res.end(); }; const onError = (error) => { if (!res.writableEnded) { - res.write(`event: error\ndata: ${JSON.stringify({ error })}\n\n`); - if (typeof res.flush === 'function') { - res.flush(); - } + streamTelemetry.recordErrorEventEmitted(); + writeEvent({ error }, { eventName: 'error' }); res.end(); } }; @@ -117,12 +126,7 @@ router.get('/chat/stream/:streamId', async (req, res) => { if (!res.writableEnded) { if (resumeState) { - res.write( - `event: message\ndata: ${JSON.stringify({ sync: true, resumeState, pendingEvents })}\n\n`, - ); - if (typeof res.flush === 'function') { - res.flush(); - } + writeEvent({ sync: true, resumeState, pendingEvents }); GenerationJobManager.markSyncSent(streamId); logger.debug( `[AgentStream] Sent sync event for ${streamId} with ${resumeState.runSteps.length} run steps, ${pendingEvents.length} pending events`, @@ -143,6 +147,7 @@ router.get('/chat/stream/:streamId', async (req, res) => { } if (!result) { + streamTelemetry.recordSubscribeFailed(); onError('Failed to subscribe to stream'); return; } diff --git a/packages/api/src/telemetry/index.ts b/packages/api/src/telemetry/index.ts index 2eb188bf84..104d8c9b5e 100644 --- a/packages/api/src/telemetry/index.ts +++ b/packages/api/src/telemetry/index.ts @@ -1,5 +1,7 @@ export { getTelemetryConfig } from './config'; export { initializeTelemetry, shutdownTelemetry } from './sdk'; export { telemetryErrorMiddleware, telemetryMiddleware } from './middleware'; +export { createSseStreamTelemetry } from './stream'; +export type { SseStreamTelemetry } from './stream'; export type { TelemetryConfig, TelemetryStatus } from './config'; export type { TelemetryController } from './sdk'; diff --git a/packages/api/src/telemetry/stream.spec.ts b/packages/api/src/telemetry/stream.spec.ts new file mode 100644 index 0000000000..460ef40078 --- /dev/null +++ b/packages/api/src/telemetry/stream.spec.ts @@ -0,0 +1,208 @@ +import { EventEmitter } from 'node:events'; +import { context, SpanKind, SpanStatusCode, trace } from '@opentelemetry/api'; +import type { Response } from 'express'; +import type { Span, Tracer } from '@opentelemetry/api'; +import type { ServerRequest } from '~/types'; +import { createSseStreamTelemetry } from './stream'; + +interface MockResponse extends EventEmitter { + writableEnded: boolean; +} + +function createSpan(): jest.Mocked { + const span = {} as jest.Mocked; + span.addEvent = jest.fn, Parameters>(() => span); + span.addLink = jest.fn, Parameters>(() => span); + span.addLinks = jest.fn, Parameters>(() => span); + span.end = jest.fn>(); + span.isRecording = jest.fn>(() => true); + span.recordException = jest.fn>(); + span.setAttribute = jest.fn, Parameters>(() => span); + span.setAttributes = jest.fn, Parameters>(() => span); + span.setStatus = jest.fn, Parameters>(() => span); + span.spanContext = jest.fn, Parameters>( + () => ({ + spanId: '0000000000000000', + traceFlags: 0, + traceId: '00000000000000000000000000000000', + }), + ); + span.updateName = jest.fn, Parameters>(() => span); + return span; +} + +function createRequest(): ServerRequest { + return { + method: 'GET', + } as ServerRequest; +} + +function createResponse(): MockResponse { + const res = new EventEmitter() as MockResponse; + res.writableEnded = false; + return res; +} + +function mockTracer(span: jest.Mocked): jest.Mock { + const startSpan = jest.fn(() => span); + jest.spyOn(trace, 'getTracer').mockReturnValue({ startSpan } as unknown as Tracer); + return startSpan; +} + +afterEach(() => { + jest.restoreAllMocks(); +}); + +describe('createSseStreamTelemetry', () => { + it('records normal stream completion attributes', () => { + const span = createSpan(); + const startSpan = mockTracer(span); + const res = createResponse(); + const telemetry = createSseStreamTelemetry({ + isResume: false, + req: createRequest(), + res: res as Response, + streamId: 'stream-1', + }); + const payload = 'event: message\ndata: {"final":true}\n\n'; + + telemetry.recordHeadersFlushed(); + telemetry.recordFinalEventEmitted(); + telemetry.recordWrite(payload, { final: true }); + res.writableEnded = true; + res.emit('finish'); + + expect(startSpan).toHaveBeenCalledWith('librechat.sse.stream', { + kind: SpanKind.INTERNAL, + attributes: expect.objectContaining({ + 'http.request.method': 'GET', + 'http.route': '/api/agents/chat/stream/:streamId', + 'librechat.stream.id': 'stream-1', + 'librechat.stream.resume': false, + 'librechat.stream.route': '/api/agents/chat/stream/:streamId', + }), + }, context.active()); + expect(span.addEvent).toHaveBeenCalledWith('headers_flushed'); + expect(span.addEvent).toHaveBeenCalledWith('first_chunk'); + expect(span.addEvent).toHaveBeenCalledWith('final_event_emitted'); + expect(span.addEvent).toHaveBeenCalledWith('final_event_written'); + expect(span.setAttributes).toHaveBeenCalledWith( + expect.objectContaining({ + 'http.response.body.size': Buffer.byteLength(payload), + 'librechat.stream.bytes.sent': Buffer.byteLength(payload), + 'librechat.stream.chunks.count': 1, + 'librechat.stream.completed': true, + 'librechat.stream.end_reason': 'done', + 'librechat.stream.error_event_emitted': false, + 'librechat.stream.final_event_emitted': true, + 'librechat.stream.final_event_written': true, + 'librechat.stream.time_to_first_chunk_ms': expect.any(Number), + }), + ); + expect(span.setStatus).not.toHaveBeenCalled(); + expect(span.end).toHaveBeenCalledTimes(1); + }); + + it('records client aborts on close before writableEnded', () => { + const span = createSpan(); + mockTracer(span); + const res = createResponse(); + const telemetry = createSseStreamTelemetry({ + isResume: true, + req: createRequest(), + res: res as Response, + streamId: 'stream-2', + }); + + telemetry.recordWrite('event: message\ndata: {}\n\n'); + res.emit('close'); + + expect(span.addEvent).toHaveBeenCalledWith('client_aborted'); + expect(span.setAttributes).toHaveBeenCalledWith( + expect.objectContaining({ + 'librechat.stream.completed': false, + 'librechat.stream.end_reason': 'client_aborted', + }), + ); + expect(span.setStatus).toHaveBeenCalledWith({ code: SpanStatusCode.ERROR }); + expect(span.setAttribute).toHaveBeenCalledWith('error.type', 'client_aborted'); + expect(span.end).toHaveBeenCalledTimes(1); + }); + + it('preserves subscribe_failed as the terminal reason when an error event is sent', () => { + const span = createSpan(); + mockTracer(span); + const res = createResponse(); + const telemetry = createSseStreamTelemetry({ + isResume: false, + req: createRequest(), + res: res as Response, + streamId: 'stream-3', + }); + + telemetry.recordSubscribeFailed(); + telemetry.recordErrorEventEmitted(); + telemetry.recordWrite('event: error\ndata: {"error":"Failed to subscribe to stream"}\n\n'); + res.writableEnded = true; + res.emit('finish'); + + expect(span.addEvent).toHaveBeenCalledWith('subscribe_failed'); + expect(span.addEvent).toHaveBeenCalledWith('error_event_emitted'); + expect(span.setAttributes).toHaveBeenCalledWith( + expect.objectContaining({ + 'librechat.stream.completed': false, + 'librechat.stream.end_reason': 'subscribe_failed', + 'librechat.stream.error_event_emitted': true, + }), + ); + expect(span.setStatus).toHaveBeenCalledWith({ code: SpanStatusCode.ERROR }); + expect(span.setAttribute).toHaveBeenCalledWith('error.type', 'subscribe_failed'); + expect(span.end).toHaveBeenCalledTimes(1); + }); + + it('ends only once when finish and close both fire', () => { + const span = createSpan(); + mockTracer(span); + const res = createResponse(); + const telemetry = createSseStreamTelemetry({ + isResume: false, + req: createRequest(), + res: res as Response, + streamId: 'stream-4', + }); + + telemetry.recordFinalEventEmitted(); + telemetry.recordWrite('event: message\ndata: {"final":true}\n\n', { final: true }); + res.writableEnded = true; + res.emit('finish'); + res.emit('close'); + + expect(span.end).toHaveBeenCalledTimes(1); + expect(span.setAttributes).toHaveBeenCalledTimes(1); + }); + + it('does not count writes after the stream span has ended', () => { + const span = createSpan(); + mockTracer(span); + const res = createResponse(); + const telemetry = createSseStreamTelemetry({ + isResume: false, + req: createRequest(), + res: res as Response, + streamId: 'stream-5', + }); + const firstPayload = 'event: message\ndata: {"index":1}\n\n'; + + telemetry.recordWrite(firstPayload); + res.emit('close'); + telemetry.recordWrite('event: message\ndata: {"index":2}\n\n'); + + expect(span.setAttributes).toHaveBeenCalledWith( + expect.objectContaining({ + 'librechat.stream.bytes.sent': Buffer.byteLength(firstPayload), + 'librechat.stream.chunks.count': 1, + }), + ); + expect(span.end).toHaveBeenCalledTimes(1); + }); +}); diff --git a/packages/api/src/telemetry/stream.ts b/packages/api/src/telemetry/stream.ts new file mode 100644 index 0000000000..d696f557c3 --- /dev/null +++ b/packages/api/src/telemetry/stream.ts @@ -0,0 +1,145 @@ +import { performance } from 'node:perf_hooks'; +import { context, SpanKind, SpanStatusCode, trace } from '@opentelemetry/api'; +import type { Attributes, Span } from '@opentelemetry/api'; +import type { Response } from 'express'; +import type { ServerRequest } from '~/types'; + +const STREAM_SPAN_NAME = 'librechat.sse.stream'; +const STREAM_ROUTE = '/api/agents/chat/stream/:streamId'; + +type StreamEndReason = 'done' | 'client_aborted' | 'server_error' | 'subscribe_failed'; + +export interface SseStreamTelemetry { + recordHeadersFlushed: () => void; + recordWrite: (payload: string, options?: { final?: boolean }) => void; + recordFinalEventEmitted: () => void; + recordErrorEventEmitted: () => void; + recordSubscribeFailed: () => void; +} + +interface SseStreamTelemetryOptions { + isResume: boolean; + req: ServerRequest; + res: Response; + streamId: string; +} + +class SseStreamSpanTelemetry implements SseStreamTelemetry { + private readonly span: Span; + private readonly startTimeMs = performance.now(); + private bytesSent = 0; + private chunksCount = 0; + private ended = false; + private errorEventEmitted = false; + private finalEventEmitted = false; + private finalEventWritten = false; + private firstChunkMs: number | undefined; + private plannedEndReason: StreamEndReason | undefined; + + constructor({ isResume, req, res, streamId }: SseStreamTelemetryOptions) { + this.span = trace.getTracer('librechat.telemetry').startSpan(STREAM_SPAN_NAME, { + kind: SpanKind.INTERNAL, + attributes: { + 'http.request.method': req.method, + 'http.route': STREAM_ROUTE, + 'librechat.stream.id': streamId, + 'librechat.stream.resume': isResume, + 'librechat.stream.route': STREAM_ROUTE, + }, + }, context.active()); + + res.once('finish', () => { + this.end(this.plannedEndReason ?? (this.errorEventEmitted ? 'server_error' : 'done')); + }); + + res.once('close', () => { + if (res.writableEnded) { + this.end(this.plannedEndReason ?? (this.errorEventEmitted ? 'server_error' : 'done')); + return; + } + + this.span.addEvent('client_aborted'); + this.end('client_aborted'); + }); + } + + recordHeadersFlushed(): void { + this.span.addEvent('headers_flushed'); + this.span.setAttribute('librechat.stream.headers_flushed', true); + } + + recordWrite(payload: string, options?: { final?: boolean }): void { + if (this.ended) { + return; + } + + this.chunksCount += 1; + this.bytesSent += Buffer.byteLength(payload); + + if (this.firstChunkMs === undefined) { + this.firstChunkMs = performance.now() - this.startTimeMs; + this.span.addEvent('first_chunk'); + this.span.setAttribute('librechat.stream.time_to_first_chunk_ms', this.firstChunkMs); + } + + if (options?.final) { + this.finalEventWritten = true; + this.span.addEvent('final_event_written'); + } + } + + recordFinalEventEmitted(): void { + this.finalEventEmitted = true; + this.plannedEndReason = 'done'; + this.span.addEvent('final_event_emitted'); + } + + recordErrorEventEmitted(): void { + this.errorEventEmitted = true; + this.plannedEndReason ??= 'server_error'; + this.span.addEvent('error_event_emitted'); + } + + recordSubscribeFailed(): void { + this.plannedEndReason = 'subscribe_failed'; + this.span.addEvent('subscribe_failed'); + } + + private end(reason: StreamEndReason): void { + if (this.ended) { + return; + } + + this.ended = true; + const attributes: Attributes = { + 'http.response.body.size': this.bytesSent, + 'librechat.stream.bytes.sent': this.bytesSent, + 'librechat.stream.chunks.count': this.chunksCount, + 'librechat.stream.completed': reason === 'done', + 'librechat.stream.duration_ms': performance.now() - this.startTimeMs, + 'librechat.stream.end_reason': reason, + 'librechat.stream.error_event_emitted': this.errorEventEmitted, + 'librechat.stream.final_event_emitted': this.finalEventEmitted, + 'librechat.stream.final_event_written': this.finalEventWritten, + }; + + if (this.firstChunkMs !== undefined) { + attributes['librechat.stream.time_to_first_chunk_ms'] = this.firstChunkMs; + } + + this.span.setAttributes(attributes); + + if (reason !== 'done') { + this.span.setStatus({ code: SpanStatusCode.ERROR }); + this.span.setAttribute('error.type', reason); + } + + this.span.end(); + } +} + +export function createSseStreamTelemetry( + options: SseStreamTelemetryOptions, +): SseStreamTelemetry { + return new SseStreamSpanTelemetry(options); +}