fix(core): Setup websocket keep-live messages (#6866)

fixes #6757

packages/cli/src/push/index.ts

@@ -26,7 +26,7 @@ export class Push extends EventEmitter {
 		} else if (!useWebSockets) {
 			(this.backend as SSEPush).add(req.query.sessionId, { req, res });
 		} else {
-			res.status(1008).send('Unauthorized');
+			res.status(401).send('Unauthorized');
 		}
 		this.emit('editorUiConnected', req.query.sessionId);
 	}
@@ -88,7 +88,7 @@ export const setupPushHandler = (restEndpoint: string, app: Application) => {
 				ws.send(`Unauthorized: ${(error as Error).message}`);
 				ws.close(1008);
 			} else {
-				res.status(1008).send('Unauthorized');
+				res.status(401).send('Unauthorized');
 			}
 			return;
 		}

packages/cli/src/push/websocket.push.ts

@@ -1,12 +1,29 @@
 import type WebSocket from 'ws';
 import { AbstractPush } from './abstract.push';
 
+function heartbeat(this: WebSocket) {
+	this.isAlive = true;
+}
+
 export class WebSocketPush extends AbstractPush<WebSocket> {
+	constructor() {
+		super();
+
+		// Ping all connected clients every 60 seconds
+		setInterval(() => this.pingAll(), 60 * 1000);
+	}
+
 	add(sessionId: string, connection: WebSocket) {
+		connection.isAlive = true;
+		connection.on('pong', heartbeat);
+
 		super.add(sessionId, connection);
 
 		// Makes sure to remove the session if the connection is closed
-		connection.once('close', () => this.remove(sessionId));
+		connection.once('close', () => {
+			connection.off('pong', heartbeat);
+			this.remove(sessionId);
+		});
 	}
 
 	protected close(connection: WebSocket): void {
@@ -16,4 +33,18 @@ export class WebSocketPush extends AbstractPush<WebSocket> {
 	protected sendToOne(connection: WebSocket, data: string): void {
 		connection.send(data);
 	}
+
+	private pingAll() {
+		for (const sessionId in this.connections) {
+			const connection = this.connections[sessionId];
+			// If a connection did not respond with a `PONG` in the last 60 seconds, disconnect
+			if (!connection.isAlive) {
+				delete this.connections[sessionId];
+				return connection.terminate();
+			}
+
+			connection.isAlive = false;
+			connection.ping();
+		}
+	}
 }