Add Safe Mode Restart Switch (#2437)

CODEOWNERS

@@ -126,6 +126,7 @@ esphome/components/restart/* @esphome/core
 esphome/components/rf_bridge/* @jesserockz
 esphome/components/rgbct/* @jesserockz
 esphome/components/rtttl/* @glmnet
+esphome/components/safe_mode/* @paulmonigatti
 esphome/components/scd4x/* @sjtrny
 esphome/components/script/* @esphome/core
 esphome/components/sdm_meter/* @jesserockz @polyfaces

esphome/components/ota/ota_component.cpp

@@ -89,7 +89,8 @@ void OTAComponent::dump_config() {
     ESP_LOGCONFIG(TAG, "  Using Password.");
   }
 #endif
-  if (this->has_safe_mode_ && this->safe_mode_rtc_value_ > 1) {
+  if (this->has_safe_mode_ && this->safe_mode_rtc_value_ > 1 &&
+      this->safe_mode_rtc_value_ != esphome::ota::OTAComponent::ENTER_SAFE_MODE_MAGIC) {
     ESP_LOGW(TAG, "Last Boot was an unhandled reset, will proceed to safe mode in %d restarts",
              this->safe_mode_num_attempts_ - this->safe_mode_rtc_value_);
   }
@@ -401,6 +402,27 @@ bool OTAComponent::writeall_(const uint8_t *buf, size_t len) {
 float OTAComponent::get_setup_priority() const { return setup_priority::AFTER_WIFI; }
 uint16_t OTAComponent::get_port() const { return this->port_; }
 void OTAComponent::set_port(uint16_t port) { this->port_ = port; }
+
+void OTAComponent::set_safe_mode_pending(const bool &pending) {
+  if (!this->has_safe_mode_)
+    return;
+
+  uint32_t current_rtc = this->read_rtc_();
+
+  if (pending && current_rtc != esphome::ota::OTAComponent::ENTER_SAFE_MODE_MAGIC) {
+    ESP_LOGI(TAG, "Device will enter safe mode on next boot.");
+    this->write_rtc_(esphome::ota::OTAComponent::ENTER_SAFE_MODE_MAGIC);
+  }
+
+  if (!pending && current_rtc == esphome::ota::OTAComponent::ENTER_SAFE_MODE_MAGIC) {
+    ESP_LOGI(TAG, "Safe mode pending has been cleared");
+    this->clean_rtc();
+  }
+}
+bool OTAComponent::get_safe_mode_pending() {
+  return this->has_safe_mode_ && this->read_rtc_() == esphome::ota::OTAComponent::ENTER_SAFE_MODE_MAGIC;
+}
+
 bool OTAComponent::should_enter_safe_mode(uint8_t num_attempts, uint32_t enable_time) {
   this->has_safe_mode_ = true;
   this->safe_mode_start_time_ = millis();
@@ -409,11 +431,17 @@ bool OTAComponent::should_enter_safe_mode(uint8_t num_attempts, uint32_t enable_
   this->rtc_ = global_preferences->make_preference<uint32_t>(233825507UL, false);
   this->safe_mode_rtc_value_ = this->read_rtc_();
 
+  bool is_manual_safe_mode = this->safe_mode_rtc_value_ == esphome::ota::OTAComponent::ENTER_SAFE_MODE_MAGIC;
+
+  if (is_manual_safe_mode)
+    ESP_LOGI(TAG, "Safe mode has been entered manually");
+  else
     ESP_LOGCONFIG(TAG, "There have been %u suspected unsuccessful boot attempts.", this->safe_mode_rtc_value_);
 
-  if (this->safe_mode_rtc_value_ >= num_attempts) {
+  if (this->safe_mode_rtc_value_ >= num_attempts || is_manual_safe_mode) {
     this->clean_rtc();
 
+    if (!is_manual_safe_mode)
       ESP_LOGE(TAG, "Boot loop detected. Proceeding to safe mode.");
 
     this->status_set_error();
@@ -445,7 +473,7 @@ uint32_t OTAComponent::read_rtc_() {
 }
 void OTAComponent::clean_rtc() { this->write_rtc_(0); }
 void OTAComponent::on_safe_shutdown() {
-  if (this->has_safe_mode_)
+  if (this->has_safe_mode_ && this->read_rtc_() != esphome::ota::OTAComponent::ENTER_SAFE_MODE_MAGIC)
     this->clean_rtc();
 }
 

esphome/components/ota/ota_component.h

@@ -48,6 +48,10 @@ class OTAComponent : public Component {
 
   bool should_enter_safe_mode(uint8_t num_attempts, uint32_t enable_time);
 
+  /// Set to true if the next startup will enter safe mode
+  void set_safe_mode_pending(const bool &pending);
+  bool get_safe_mode_pending();
+
 #ifdef USE_OTA_STATE_CALLBACK
   void add_on_state_callback(std::function<void(OTAState, float, uint8_t)> &&callback);
 #endif
@@ -89,6 +93,9 @@ class OTAComponent : public Component {
   uint8_t safe_mode_num_attempts_;
   ESPPreferenceObject rtc_;
 
+  static const uint32_t ENTER_SAFE_MODE_MAGIC =
+      0x5afe5afe;  ///< a magic number to indicate that safe mode should be entered on next boot
+
 #ifdef USE_OTA_STATE_CALLBACK
   CallbackManager<void(OTAState, float, uint8_t)> state_callback_{};
 #endif

esphome/components/safe_mode/__init__.py

@@ -0,0 +1,5 @@
+import esphome.codegen as cg
+
+CODEOWNERS = ["@paulmonigatti"]
+
+safe_mode_ns = cg.esphome_ns.namespace("safe_mode")

esphome/components/safe_mode/switch/__init__.py

@@ -0,0 +1,36 @@
+import esphome.codegen as cg
+import esphome.config_validation as cv
+from esphome.components import switch
+from esphome.components.ota import OTAComponent
+from esphome.const import (
+    CONF_ID,
+    CONF_INVERTED,
+    CONF_ICON,
+    CONF_OTA,
+    ICON_RESTART_ALERT,
+)
+from .. import safe_mode_ns
+
+DEPENDENCIES = ["ota"]
+
+SafeModeSwitch = safe_mode_ns.class_("SafeModeSwitch", switch.Switch, cg.Component)
+
+CONFIG_SCHEMA = switch.SWITCH_SCHEMA.extend(
+    {
+        cv.GenerateID(): cv.declare_id(SafeModeSwitch),
+        cv.GenerateID(CONF_OTA): cv.use_id(OTAComponent),
+        cv.Optional(CONF_INVERTED): cv.invalid(
+            "Safe Mode Restart switches do not support inverted mode!"
+        ),
+        cv.Optional(CONF_ICON, default=ICON_RESTART_ALERT): switch.icon,
+    }
+).extend(cv.COMPONENT_SCHEMA)
+
+
+async def to_code(config):
+    var = cg.new_Pvariable(config[CONF_ID])
+    await cg.register_component(var, config)
+    await switch.register_switch(var, config)
+
+    ota = await cg.get_variable(config[CONF_OTA])
+    cg.add(var.set_ota(ota))

esphome/components/safe_mode/switch/safe_mode_switch.cpp

@@ -0,0 +1,29 @@
+#include "safe_mode_switch.h"
+#include "esphome/core/hal.h"
+#include "esphome/core/log.h"
+#include "esphome/core/application.h"
+
+namespace esphome {
+namespace safe_mode {
+
+static const char *const TAG = "safe_mode_switch";
+
+void SafeModeSwitch::set_ota(ota::OTAComponent *ota) { this->ota_ = ota; }
+
+void SafeModeSwitch::write_state(bool state) {
+  // Acknowledge
+  this->publish_state(false);
+
+  if (state) {
+    ESP_LOGI(TAG, "Restarting device in safe mode...");
+    this->ota_->set_safe_mode_pending(true);
+
+    // Let MQTT settle a bit
+    delay(100);  // NOLINT
+    App.safe_reboot();
+  }
+}
+void SafeModeSwitch::dump_config() { LOG_SWITCH("", "Safe Mode Switch", this); }
+
+}  // namespace safe_mode
+}  // namespace esphome

esphome/components/safe_mode/switch/safe_mode_switch.h

@@ -0,0 +1,21 @@
+#pragma once
+
+#include "esphome/core/component.h"
+#include "esphome/components/ota/ota_component.h"
+#include "esphome/components/switch/switch.h"
+
+namespace esphome {
+namespace safe_mode {
+
+class SafeModeSwitch : public switch_::Switch, public Component {
+ public:
+  void dump_config() override;
+  void set_ota(ota::OTAComponent *ota);
+
+ protected:
+  ota::OTAComponent *ota_;
+  void write_state(bool state) override;
+};
+
+}  // namespace safe_mode
+}  // namespace esphome

esphome/const.py

@@ -761,6 +761,7 @@ ICON_PULSE = "mdi:pulse"
 ICON_RADIATOR = "mdi:radiator"
 ICON_RADIOACTIVE = "mdi:radioactive"
 ICON_RESTART = "mdi:restart"
+ICON_RESTART_ALERT = "mdi:restart-alert"
 ICON_ROTATE_RIGHT = "mdi:rotate-right"
 ICON_RULER = "mdi:ruler"
 ICON_SCALE = "mdi:scale"

tests/test1.yaml

@@ -1905,6 +1905,8 @@ switch:
           state: yes
   - platform: restart
     name: 'Living Room Restart'
+  - platform: safe_mode
+    name: 'Living Room Restart (Safe Mode)'
   - platform: shutdown
     name: 'Living Room Shutdown'
   - platform: output