diff --git a/code/espurna/mqtt.ino b/code/espurna/mqtt.ino
index c44c36b0..4c8f3e56 100644
--- a/code/espurna/mqtt.ino
+++ b/code/espurna/mqtt.ino
@@ -316,9 +316,7 @@ unsigned long _mqttNextMessageId() {
 #if WEB_SUPPORT
 
 bool _mqttWebSocketOnReceive(const char * key, JsonVariant& value) {
-    if (strncmp(key, "mqtt", 3) == 0) return true;
-    if (strncmp(key, "hb", 2) == 0) return true;
-    return false;
+    return (strncmp(key, "mqtt", 3) == 0);
 }
 
 void _mqttWebSocketOnSend(JsonObject& root) {
diff --git a/code/espurna/system.ino b/code/espurna/system.ino
index 27b2d4dd..791ba99c 100644
--- a/code/espurna/system.ino
+++ b/code/espurna/system.ino
@@ -86,6 +86,12 @@ void _systemSetupHeartbeat() {
     _heartbeat_interval = getSetting("hbInterval", HEARTBEAT_INTERVAL).toInt();
 }
 
+#if WEB_SUPPORT
+    bool _systemWebSocketOnReceive(const char * key, JsonVariant& value) {
+        return (strncmp(key, "hb", 2) == 0);
+    }
+#endif
+
 void systemLoop() {
 
     // -------------------------------------------------------------------------
@@ -182,6 +188,10 @@ void systemSetup() {
         systemCheck(false);
     #endif
 
+    #if WEB_SUPPORT
+        wsOnReceiveRegister(_systemWebSocketOnReceive);
+    #endif
+
     // Init device-specific hardware
     _systemSetupSpecificHardware();