Recover from broken websocket connection

lib/connection.ts

@@ -3,7 +3,11 @@
  * the Home Assistant websocket API.
  */
 import * as messages from "./messages.js";
-import { ERR_INVALID_AUTH, ERR_CONNECTION_LOST } from "./errors.js";
+import {
+  ERR_INVALID_AUTH,
+  ERR_CONNECTION_LOST,
+  ERR_CONNECTION_TIMEOUT,
+} from "./errors.js";
 import { HassEvent, MessageBase } from "./types.js";
 import { HaWebSocket } from "./socket.js";
 import type { Auth } from "./auth.js";
@@ -76,6 +80,10 @@ type CommandInFlight =
   | SubscribeEventCommmandInFlight<any>
   | CommandWithAnswerInFlight;
 
+type PingTimerData = {
+  timerRef?: number;
+  executing: boolean;
+};
 export class Connection {
   options: ConnectionOptions;
   commandId: number;
@@ -86,6 +94,11 @@ export class Connection {
 
   oldSubscriptions?: Map<number, CommandInFlight>;
 
+  pingTimeout: number;
+  pingInterval: number;
+
+  pingTimer: PingTimerData;
+
   // We use this to queue messages in flight for the first reconnect
   // after the connection has been suspended.
   _queuedMessages?: Array<{
@@ -112,8 +125,16 @@ export class Connection {
     this.eventListeners = new Map();
     // true if a close is requested by the user
     this.closeRequested = false;
+    // Ping timeout in ms
+    this.pingTimeout = 5 * 1000;
+    // Ping interval in ms
+    this.pingInterval = 60 * 1000;
+    // Object holding state of current ping timer
+    this.pingTimer = { executing: false };
 
     this._setSocket(socket);
+
+    this._scheduledPing();
   }
 
   get connected() {
@@ -240,7 +261,14 @@ export class Connection {
   }
 
   ping() {
-    return this.sendMessagePromise(messages.ping());
+    // create a promise that rejects in milliseconds
+    const pingRequest = this.sendMessagePromise(messages.ping());
+    const timeout = new Promise<never>((_, reject) => {
+      setTimeout(() => {
+        reject(ERR_CONNECTION_TIMEOUT);
+      }, this.pingTimeout);
+    });
+    return Promise.race([pingRequest, timeout]);
   }
 
   sendMessage(message: MessageBase, commandId?: number): void {
@@ -353,6 +381,8 @@ export class Connection {
       event.data,
     );
 
+    this._scheduledPing();
+
     if (!Array.isArray(messageGroup)) {
       messageGroup = [messageGroup];
     }
@@ -483,6 +513,36 @@ export class Connection {
     reconnect(0);
   };
 
+  private _scheduledPing() {
+    if (this.pingTimer.executing) {
+      return;
+    }
+    // Reset timer before before scheduling a new one
+    if (this.pingTimer.timerRef) {
+      clearInterval(this.pingTimer.timerRef);
+    }
+
+    const pingLoop = () => {
+      this.pingTimer.timerRef = setTimeout(async () => {
+        this.pingTimer.executing = true;
+        if (this.connected && !this.closeRequested) {
+          try {
+            await this.ping();
+          } catch (error) {
+            if (error === ERR_CONNECTION_TIMEOUT) {
+              // Reconnect needs to be forced since no events are triggered when websocket is broken
+              this.reconnect(true);
+            }
+          }
+        }
+        pingLoop();
+        this.pingTimer.executing = false;
+      }, this.pingInterval);
+    };
+
+    pingLoop();
+  }
+
   private _genCmdId() {
     return ++this.commandId;
   }

lib/errors.ts

@@ -4,3 +4,4 @@ export const ERR_CONNECTION_LOST = 3;
 export const ERR_HASS_HOST_REQUIRED = 4;
 export const ERR_INVALID_HTTPS_TO_HTTP = 5;
 export const ERR_INVALID_AUTH_CALLBACK = 6;
+export const ERR_CONNECTION_TIMEOUT = 7;