🖲️ feat: Trace SSE Stream Lifecycle with OTel (#13266)

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;
   }

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';

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<Span> {
+  const span = {} as jest.Mocked<Span>;
+  span.addEvent = jest.fn<jest.Mocked<Span>, Parameters<Span['addEvent']>>(() => span);
+  span.addLink = jest.fn<jest.Mocked<Span>, Parameters<Span['addLink']>>(() => span);
+  span.addLinks = jest.fn<jest.Mocked<Span>, Parameters<Span['addLinks']>>(() => span);
+  span.end = jest.fn<void, Parameters<Span['end']>>();
+  span.isRecording = jest.fn<boolean, Parameters<Span['isRecording']>>(() => true);
+  span.recordException = jest.fn<void, Parameters<Span['recordException']>>();
+  span.setAttribute = jest.fn<jest.Mocked<Span>, Parameters<Span['setAttribute']>>(() => span);
+  span.setAttributes = jest.fn<jest.Mocked<Span>, Parameters<Span['setAttributes']>>(() => span);
+  span.setStatus = jest.fn<jest.Mocked<Span>, Parameters<Span['setStatus']>>(() => span);
+  span.spanContext = jest.fn<ReturnType<Span['spanContext']>, Parameters<Span['spanContext']>>(
+    () => ({
+      spanId: '0000000000000000',
+      traceFlags: 0,
+      traceId: '00000000000000000000000000000000',
+    }),
+  );
+  span.updateName = jest.fn<jest.Mocked<Span>, Parameters<Span['updateName']>>(() => 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<Span>): 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);
+  });
+});

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);
+}