diff --git a/library.json b/library.json
index d2c4995..b6f29eb 100644
--- a/library.json
+++ b/library.json
@@ -1,7 +1,7 @@
 {
   "$schema": "https://raw.githubusercontent.com/platformio/platformio-core/develop/platformio/assets/schema/library.json",
   "name": "ha-mqtt",
-  "version": "1.5.0",
+  "version": "1.5.1",
   "description": "Home Assistant classes for integration with MQTT auto discovery",
   "repository": {
     "type": "git",
diff --git a/src/ha.h b/src/ha.h
index 3fc75f5..3d627ea 100644
--- a/src/ha.h
+++ b/src/ha.h
@@ -10,8 +10,8 @@ using namespace std;
 #define BASE_TOPIC  "homeassistant/%s/%s/%s"
 
 namespace Ha {
-    uint16_t(*publisher)(const char* topic, const char* message);
-    typedef void (*onMessage)(const char* msg);
+    uint16(*publisher)(const char*, const char*);
+    typedef void (*onMessage)(const char*);
 
     struct Config {
         virtual void buildConfig(JsonDocument& jsonDoc) = 0;
@@ -189,7 +189,8 @@ namespace Ha {
 
         StatefulCommand(Component* cmp, onMessage f) : Command(cmp, f), State(cmp) {}
 
-        void restoreFromState() {
+        void restoreStateFromCommand() {
+            withStateTopic();
             mapRestoreStateTopics.insert({stateTopic, this});
         }
 
@@ -457,8 +458,14 @@ namespace Ha {
             return *this;
         }
 
+        [[deprecated("Use restoreStateFromCommand() instead")]]
         Builder& restoreFromState() {
-            cmp->restoreFromState();
+            cmp->restoreStateFromCommand();
+            return *this;
+        }
+
+        Builder& restoreStateFromCommand() {
+            cmp->restoreStateFromCommand();
             return *this;
         }
     };
diff --git a/src/mqtt.h b/src/mqtt.h
index 796aa78..f4424c7 100644
--- a/src/mqtt.h
+++ b/src/mqtt.h
@@ -22,7 +22,7 @@ namespace Mqtt {
         client.disconnect();
     }
 
-    uint16_t publish(const char* topic, const char* message) {
+    uint16 publish(const char* topic, const char* message) {
         return client.publish(topic, 0, true, message);
     }