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

fixes #6757

package.json

@@ -89,7 +89,8 @@
       "typedi@0.10.0": "patches/typedi@0.10.0.patch",
       "@sentry/cli@2.17.0": "patches/@sentry__cli@2.17.0.patch",
       "pkce-challenge@3.0.0": "patches/pkce-challenge@3.0.0.patch",
-      "pyodide@0.23.4": "patches/pyodide@0.23.4.patch"
+      "pyodide@0.23.4": "patches/pyodide@0.23.4.patch",
+      "@types/ws@8.5.4": "patches/@types__ws@8.5.4.patch"
     }
   }
 }

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

patches/@types__ws@8.5.4.patch

@@ -0,0 +1,14 @@
+diff --git a/index.d.ts b/index.d.ts
+index 7a8182a94289524851cb08a3b24897f2b6bce747..f5bfb61bdacbae81ca274cc4b5a61e6e7322b7cd 100755
+--- a/index.d.ts
++++ b/index.d.ts
+@@ -72,6 +72,9 @@ declare class WebSocket extends EventEmitter {
+         | typeof WebSocket.CLOSED;
+     readonly url: string;
+ 
++    /** Indicates if the connection has replied to the last PING */
++    isAlive: boolean;
++
+     /** The connection is not yet open. */
+     readonly CONNECTING: 0;
+     /** The connection is open and ready to communicate. */

pnpm-lock.yaml

@@ -26,6 +26,9 @@ patchedDependencies:
   '@sentry/cli@2.17.0':
     hash: nchnoezkq6p37qaiku3vrpwraq
     path: patches/@sentry__cli@2.17.0.patch
+  '@types/ws@8.5.4':
+    hash: nbzuqaoyqbrfwipijj5qriqqju
+    path: patches/@types__ws@8.5.4.patch
   pkce-challenge@3.0.0:
     hash: dypouzb3lve7vncq25i5fuanki
     path: patches/pkce-challenge@3.0.0.patch
@@ -554,7 +557,7 @@ importers:
         version: 13.7.7
       '@types/ws':
         specifier: ^8.5.4
-        version: 8.5.4
+        version: 8.5.4(patch_hash=nbzuqaoyqbrfwipijj5qriqqju)
       '@types/xml2js':
         specifier: ^0.4.11
         version: 0.4.11
@@ -6719,11 +6722,12 @@ packages:
       '@types/webidl-conversions': 7.0.0
     dev: false
 
-  /@types/ws@8.5.4:
+  /@types/ws@8.5.4(patch_hash=nbzuqaoyqbrfwipijj5qriqqju):
     resolution: {integrity: sha512-zdQDHKUgcX/zBc4GrwsE/7dVdAD8JR4EuiAXiiUhhfyIJXXb2+PrGshFyeXWQPMmmZ2XxgaqclgpIC7eTXc1mg==}
     dependencies:
       '@types/node': 18.16.16
     dev: true
+    patched: true
 
   /@types/xml2js@0.4.11:
     resolution: {integrity: sha512-JdigeAKmCyoJUiQljjr7tQG3if9NkqGUgwEUqBvV0N7LM4HyQk7UXCnusRa1lnvXAEYJ8mw8GtZWioagNztOwA==}