🛟 fix: Keep File Uploads Alive With SSE Heartbeats (#14295)

* fix: Use SSE to upload files in order to avoid idle timeouts.  Idle timeouts can occur for example from gateways and other services like cloudfare when uploading large files.  For example during rag processing the file is uploaded to librechat which then sends it to rag.  While librechat is waiting for the embeddings to come back from rag the file upload is sitting idle.  Gateways tend to want to cancel the upload with an http 408 , 504, or 524.  This change uses SSE to perform the upload so that while librechat is sending the file to rag, it consistently sends back a heartbeat event to the client to keep the connection alive.  This is especially useful when utilizing  EMBEDDING_BATCH_SIZE in librechat rag which will allow rag to process signifigantly larger files without running out of memory.

* added tests to packages\api\src\files\sse.spec.ts in order to test the new sse.ts

* fix: Harden SSE file upload lifecycle

* style: Sort data provider imports

---------

Co-authored-by: Marc Amick <MarcAmick@jhu.edu>
Co-authored-by: Danny Avila <danny@librechat.ai>

packages/api/src/files/index.ts

@@ -10,6 +10,7 @@ export * from './ocr';
 export * from './parse';
 export * from './rag';
 export * from './retention';
+export * from './sse';
 export * from './sweep';
 export * from './usage';
 export * from './validation';

packages/api/src/files/sse.spec.ts

@@ -0,0 +1,218 @@
+import { EventEmitter } from 'events';
+import type { Request, Response } from 'express';
+import { sendUploadSuccess, shouldUseUploadSse, startUploadSseStream } from './sse';
+
+describe('sse', () => {
+  const createMockReq = (accept?: string): Request =>
+    ({ headers: accept ? { accept } : {} }) as Request;
+
+  const createMockRes = (): jest.Mocked<Response> => {
+    const res = new EventEmitter() as unknown as jest.Mocked<Response>;
+    res.writeHead = jest.fn().mockReturnValue(res);
+    res.flushHeaders = jest.fn().mockReturnValue(res);
+    res.write = jest.fn().mockReturnValue(true);
+    res.end = jest.fn().mockReturnValue(res);
+    res.status = jest.fn().mockReturnValue(res);
+    res.json = jest.fn().mockReturnValue(res);
+    Object.defineProperty(res, 'writableEnded', { value: false, writable: true });
+    Object.defineProperty(res, 'destroyed', { value: false, writable: true });
+    return res;
+  };
+
+  const parseEvents = (res: jest.Mocked<Response>): Array<{ event: string; data: unknown }> =>
+    (res.write as jest.Mock).mock.calls.map(([chunk]: [string]) => {
+      const eventMatch = /event:(.*)\n/.exec(chunk);
+      const dataMatch = /data:(.*)\n\n/.exec(chunk);
+      return {
+        event: eventMatch ? eventMatch[1] : '',
+        data: dataMatch ? JSON.parse(dataMatch[1]) : undefined,
+      };
+    });
+
+  describe('shouldUseUploadSse', () => {
+    const originalValue = process.env.FILE_UPLOAD_SSE_ENABLED;
+
+    afterEach(() => {
+      if (originalValue === undefined) {
+        delete process.env.FILE_UPLOAD_SSE_ENABLED;
+        return;
+      }
+      process.env.FILE_UPLOAD_SSE_ENABLED = originalValue;
+    });
+
+    it('requires both the feature flag and an explicit event-stream accept value', () => {
+      process.env.FILE_UPLOAD_SSE_ENABLED = 'true';
+
+      expect(shouldUseUploadSse(createMockReq('application/json, text/event-stream'))).toBe(true);
+      expect(shouldUseUploadSse(createMockReq('text/event-stream; charset=utf-8'))).toBe(true);
+      expect(shouldUseUploadSse(createMockReq('application/json'))).toBe(false);
+      expect(shouldUseUploadSse(createMockReq('*/*'))).toBe(false);
+      expect(shouldUseUploadSse(createMockReq())).toBe(false);
+    });
+
+    it('keeps JSON responses when the server feature flag is disabled', () => {
+      process.env.FILE_UPLOAD_SSE_ENABLED = 'false';
+
+      expect(shouldUseUploadSse(createMockReq('text/event-stream'))).toBe(false);
+    });
+  });
+
+  describe('startUploadSseStream', () => {
+    beforeEach(() => {
+      jest.useFakeTimers();
+    });
+
+    afterEach(() => {
+      jest.useRealTimers();
+    });
+
+    it('writes the SSE headers and flushes them immediately', () => {
+      const res = createMockRes();
+      const stream = startUploadSseStream(res);
+
+      expect(res.writeHead).toHaveBeenCalledWith(
+        200,
+        expect.objectContaining({
+          'Content-Type': 'text/event-stream',
+          Connection: 'keep-alive',
+        }),
+      );
+      expect(res.flushHeaders).toHaveBeenCalledTimes(1);
+      stream.close();
+    });
+
+    it('emits a heartbeat event on every interval tick', () => {
+      const res = createMockRes();
+      const stream = startUploadSseStream(res);
+
+      jest.advanceTimersByTime(3000);
+
+      const heartbeats = parseEvents(res).filter((e) => e.event === 'heartbeat');
+      expect(heartbeats).toHaveLength(3);
+      expect(heartbeats.map((e) => e.data)).toEqual([
+        { keepAlive: 1 },
+        { keepAlive: 2 },
+        { keepAlive: 3 },
+      ]);
+      stream.close();
+    });
+
+    it('stops the heartbeat once the response has already ended', () => {
+      const res = createMockRes();
+      startUploadSseStream(res);
+
+      jest.advanceTimersByTime(1000);
+      Object.defineProperty(res, 'writableEnded', { value: true });
+      jest.advanceTimersByTime(2000);
+
+      const heartbeats = parseEvents(res).filter((e) => e.event === 'heartbeat');
+      expect(heartbeats).toHaveLength(1);
+      expect(jest.getTimerCount()).toBe(0);
+    });
+
+    it('clears the heartbeat interval on response disconnect and stops writes', () => {
+      const res = createMockRes();
+      const stream = startUploadSseStream(res);
+      expect(jest.getTimerCount()).toBe(1);
+
+      Object.defineProperty(res, 'destroyed', { value: true });
+      res.emit('close');
+      stream.sendData({ file_id: 'ignored' });
+
+      expect(jest.getTimerCount()).toBe(0);
+      expect(res.write).not.toHaveBeenCalled();
+    });
+
+    it('emits data and error events', () => {
+      const res = createMockRes();
+      const stream = startUploadSseStream(res);
+
+      stream.sendData({ file_id: 'abc' });
+      stream.sendError({ message: 'failed' });
+
+      const events = parseEvents(res);
+      expect(events.find((event) => event.event === 'data')?.data).toEqual({ file_id: 'abc' });
+      expect(events.find((event) => event.event === 'error')?.data).toEqual({
+        message: 'failed',
+      });
+      stream.close();
+    });
+
+    describe('close', () => {
+      it('emits a close event, ends the response, and clears the heartbeat', () => {
+        const res = createMockRes();
+        const stream = startUploadSseStream(res);
+
+        stream.close();
+
+        const closeEvents = parseEvents(res).filter((e) => e.event === 'close');
+        expect(closeEvents).toHaveLength(1);
+        expect(closeEvents[0].data).toEqual(
+          expect.objectContaining({ closedAt: expect.any(String) }),
+        );
+        expect(res.end).toHaveBeenCalledTimes(1);
+        expect(jest.getTimerCount()).toBe(0);
+      });
+
+      it('does not write after the response already ended', () => {
+        const res = createMockRes();
+        const stream = startUploadSseStream(res);
+        Object.defineProperty(res, 'writableEnded', { value: true });
+
+        stream.close();
+
+        expect(parseEvents(res).filter((e) => e.event === 'close')).toHaveLength(0);
+        expect(res.end).not.toHaveBeenCalled();
+        expect(jest.getTimerCount()).toBe(0);
+      });
+
+      it('does not emit further heartbeats after close', () => {
+        const res = createMockRes();
+        const stream = startUploadSseStream(res);
+
+        stream.close();
+        jest.advanceTimersByTime(5000);
+
+        expect(parseEvents(res).filter((e) => e.event === 'heartbeat')).toHaveLength(0);
+      });
+    });
+  });
+
+  describe('sendUploadSuccess', () => {
+    it('sends the payload over the SSE stream when a stream is active', () => {
+      const res = createMockRes();
+      const sseStream = {
+        sendData: jest.fn(),
+        sendError: jest.fn(),
+        close: jest.fn(),
+      };
+
+      sendUploadSuccess(res, sseStream, 'Upload complete', { file_id: 'abc' });
+
+      expect(sseStream.sendData).toHaveBeenCalledWith({
+        message: 'Upload complete',
+        file_id: 'abc',
+      });
+      expect(res.status).not.toHaveBeenCalled();
+      expect(res.json).not.toHaveBeenCalled();
+    });
+
+    it('falls back to a plain JSON response when no stream is provided', () => {
+      const res = createMockRes();
+
+      sendUploadSuccess(res, null, 'Upload complete', { file_id: 'abc' });
+
+      expect(res.status).toHaveBeenCalledWith(200);
+      expect(res.json).toHaveBeenCalledWith({ message: 'Upload complete', file_id: 'abc' });
+    });
+
+    it('falls back to a plain JSON response when the stream is undefined', () => {
+      const res = createMockRes();
+
+      sendUploadSuccess(res, undefined, 'Upload complete', { file_id: 'abc' });
+
+      expect(res.status).toHaveBeenCalledWith(200);
+      expect(res.json).toHaveBeenCalledWith({ message: 'Upload complete', file_id: 'abc' });
+    });
+  });
+});

packages/api/src/files/sse.ts

@@ -0,0 +1,100 @@
+import type { Request, Response } from 'express';
+import { isEnabled } from '~/utils';
+
+const HEARTBEAT_INTERVAL_MS = 1000;
+const EVENT_STREAM_MEDIA_TYPE = 'text/event-stream';
+
+const SSE_HEADERS = {
+  'Content-Type': EVENT_STREAM_MEDIA_TYPE,
+  'Cache-Control': 'no-cache, no-transform',
+  Connection: 'keep-alive',
+  /** Required so Nginx/other proxies don't buffer the stream. */
+  'X-Accel-Buffering': 'no',
+};
+
+function writeSseEvent<T>(res: Response, event: string, data: T): void {
+  if (res.writableEnded || res.destroyed) {
+    return;
+  }
+  res.write(`event:${event}\nid:${Date.now()}\ndata:${JSON.stringify(data)}\n\n`);
+}
+
+export function shouldUseUploadSse(req: Request): boolean {
+  if (!isEnabled(process.env.FILE_UPLOAD_SSE_ENABLED)) {
+    return false;
+  }
+
+  return (
+    req.headers.accept?.split(',').some((value) => {
+      const [mediaType] = value.split(';', 1);
+      return mediaType.trim().toLowerCase() === EVENT_STREAM_MEDIA_TYPE;
+    }) ?? false
+  );
+}
+
+export interface UploadSseStream {
+  /** Emits the successful upload payload as an `event:data` message. */
+  sendData: <T>(data: T) => void;
+  /** Emits a failure payload as an `event:error` message. */
+  sendError: <T>(data: T) => void;
+  /** Emits the terminal `event:close` message, stops the heartbeat, and ends the response. */
+  close: () => void;
+}
+
+/**
+ * Opens a keep-alive SSE stream for a file-upload response: sends the SSE headers, starts a
+ * heartbeat interval so proxies/clients don't time out the connection during long-running
+ * uploads, and stops the heartbeat when the client disconnects early.
+ *
+ * Callers must only invoke this once all synchronous validation and permission checks that
+ * might still need to send a normal (non-SSE) response have already passed — once the headers
+ * are flushed here, the HTTP status is committed to 200 and can no longer be changed.
+ */
+export function startUploadSseStream(res: Response): UploadSseStream {
+  res.writeHead(200, SSE_HEADERS);
+  res.flushHeaders();
+
+  let counter = 1;
+  const intervalId = setInterval(() => {
+    if (res.writableEnded || res.destroyed) {
+      clearInterval(intervalId);
+      return;
+    }
+    writeSseEvent(res, 'heartbeat', { keepAlive: counter++ });
+  }, HEARTBEAT_INTERVAL_MS);
+
+  const stopHeartbeat = () => {
+    clearInterval(intervalId);
+  };
+  res.once('close', stopHeartbeat);
+
+  return {
+    sendData: (data) => writeSseEvent(res, 'data', data),
+    sendError: (data) => writeSseEvent(res, 'error', data),
+    close: () => {
+      clearInterval(intervalId);
+      res.off('close', stopHeartbeat);
+      if (!res.writableEnded && !res.destroyed) {
+        writeSseEvent(res, 'close', { closedAt: new Date().toISOString() });
+        res.end();
+      }
+    },
+  };
+}
+
+/**
+ * Sends a successful upload response, either as an SSE `event:data` message (when an upload
+ * stream is active) or as a plain JSON response.
+ */
+export function sendUploadSuccess<T extends object>(
+  res: Response,
+  sseStream: UploadSseStream | null | undefined,
+  message: string,
+  result: T,
+): void {
+  if (sseStream) {
+    sseStream.sendData({ message, ...result });
+    return;
+  }
+  res.status(200).json({ message, ...result });
+}