feat: Enhance SSE and Redis event handling with heartbeat and cleanup mechanisms

- Added heartbeat functionality to SSEStreamer for improved client connection management.
- Implemented periodic cleanup in RedisEventSubscriber to remove stale subscriptions.
- Updated internal prediction and prediction controllers to handle client disconnects and manage subscriptions more effectively.
- Refactored RedisEventPublisher to streamline event publishing with error handling.

These changes improve the reliability and performance of server-side event streaming and Redis event handling.

Author: HenryHengZJ

packages/server/src/controllers/internal-predictions/index.ts

@@ -42,10 +42,19 @@ const createAndStreamInternalPrediction = async (req: Request, res: Response, ne
         res.setHeader('X-Accel-Buffering', 'no') //nginx config: https://serverfault.com/a/801629
         res.flushHeaders()
 
-        if (process.env.MODE === MODE.QUEUE) {
+        const isQueueMode = process.env.MODE === MODE.QUEUE
+        if (isQueueMode) {
             getRunningExpressApp().redisSubscriber.subscribe(chatId)
         }
 
+        // Detect client disconnect (browser close, network timeout, ALB idle timeout)
+        res.on('close', () => {
+            sseStreamer.removeClient(chatId)
+            if (isQueueMode) {
+                getRunningExpressApp().redisSubscriber.unsubscribe(chatId)
+            }
+        })
+
         const apiResponse = await utilBuildChatflow(req, true)
         sseStreamer.streamMetadataEvent(apiResponse.chatId, apiResponse)
     } catch (error) {
@@ -55,6 +64,9 @@ const createAndStreamInternalPrediction = async (req: Request, res: Response, ne
         next(error)
     } finally {
         sseStreamer.removeClient(chatId)
+        if (process.env.MODE === MODE.QUEUE) {
+            getRunningExpressApp().redisSubscriber.unsubscribe(chatId)
+        }
     }
 }
 export default {

packages/server/src/controllers/predictions/index.ts

@@ -72,10 +72,19 @@ const createPrediction = async (req: Request, res: Response, next: NextFunction)
                     res.setHeader('X-Accel-Buffering', 'no') //nginx config: https://serverfault.com/a/801629
                     res.flushHeaders()
 
-                    if (process.env.MODE === MODE.QUEUE) {
+                    const isQueueMode = process.env.MODE === MODE.QUEUE
+                    if (isQueueMode) {
                         getRunningExpressApp().redisSubscriber.subscribe(chatId)
                     }
 
+                    // Detect client disconnect (browser close, network timeout, ALB idle timeout)
+                    res.on('close', () => {
+                        sseStreamer.removeClient(chatId)
+                        if (isQueueMode) {
+                            getRunningExpressApp().redisSubscriber.unsubscribe(chatId)
+                        }
+                    })
+
                     const apiResponse = await predictionsServices.buildChatflow(req)
                     sseStreamer.streamMetadataEvent(apiResponse.chatId, apiResponse)
                 } catch (error) {
@@ -85,6 +94,9 @@ const createPrediction = async (req: Request, res: Response, next: NextFunction)
                     next(error)
                 } finally {
                     sseStreamer.removeClient(chatId)
+                    if (process.env.MODE === MODE.QUEUE) {
+                        getRunningExpressApp().redisSubscriber.unsubscribe(chatId)
+                    }
                 }
             } else {
                 const apiResponse = await predictionsServices.buildChatflow(req)

packages/server/src/index.ts

@@ -128,6 +128,7 @@ export class App {
 
             // Initialize SSE Streamer
             this.sseStreamer = new SSEStreamer()
+            this.sseStreamer.startHeartbeat()
             logger.info('🌊 [server]: SSE Streamer initialized successfully')
 
             // Init Queues
@@ -148,6 +149,7 @@ export class App {
 
                 this.redisSubscriber = new RedisEventSubscriber(this.sseStreamer)
                 await this.redisSubscriber.connect()
+                this.redisSubscriber.startPeriodicCleanup()
                 logger.info('🔗 [server]: Redis event subscriber connected successfully')
             }
 
@@ -361,6 +363,7 @@ export class App {
 
     async stopApp() {
         try {
+            this.sseStreamer.stopHeartbeat()
             const removePromises: any[] = []
             removePromises.push(this.telemetry.flush())
             if (this.queueManager) {