diff --git a/esphome/components/water_heater/water_heater.cpp b/esphome/components/water_heater/water_heater.cpp
index e6d15623526..9d7ae0cbc0d 100644
--- a/esphome/components/water_heater/water_heater.cpp
+++ b/esphome/components/water_heater/water_heater.cpp
@@ -65,6 +65,7 @@ WaterHeaterCall &WaterHeaterCall::set_away(bool away) {
   } else {
     this->state_ &= ~WATER_HEATER_STATE_AWAY;
   }
+  this->state_mask_ |= WATER_HEATER_STATE_AWAY;
   return *this;
 }
 
@@ -74,6 +75,7 @@ WaterHeaterCall &WaterHeaterCall::set_on(bool on) {
   } else {
     this->state_ &= ~WATER_HEATER_STATE_ON;
   }
+  this->state_mask_ |= WATER_HEATER_STATE_ON;
   return *this;
 }
 
@@ -92,11 +94,11 @@ void WaterHeaterCall::perform() {
   if (!std::isnan(this->target_temperature_high_)) {
     ESP_LOGD(TAG, "  Target Temperature High: %.2f", this->target_temperature_high_);
   }
-  if (this->state_ & WATER_HEATER_STATE_AWAY) {
-    ESP_LOGD(TAG, "  Away: YES");
+  if (this->state_mask_ & WATER_HEATER_STATE_AWAY) {
+    ESP_LOGD(TAG, "  Away: %s", (this->state_ & WATER_HEATER_STATE_AWAY) ? "YES" : "NO");
   }
-  if (this->state_ & WATER_HEATER_STATE_ON) {
-    ESP_LOGD(TAG, "  On: YES");
+  if (this->state_mask_ & WATER_HEATER_STATE_ON) {
+    ESP_LOGD(TAG, "  On: %s", (this->state_ & WATER_HEATER_STATE_ON) ? "YES" : "NO");
   }
   this->parent_->control(*this);
 }
@@ -137,13 +139,17 @@ void WaterHeaterCall::validate_() {
       this->target_temperature_high_ = NAN;
     }
   }
-  if ((this->state_ & WATER_HEATER_STATE_AWAY) && !traits.get_supports_away_mode()) {
-    ESP_LOGW(TAG, "'%s' - Away mode not supported", this->parent_->get_name().c_str());
+  if (!traits.get_supports_away_mode()) {
+    if (this->state_ & WATER_HEATER_STATE_AWAY) {
+      ESP_LOGW(TAG, "'%s' - Away mode not supported", this->parent_->get_name().c_str());
+    }
     this->state_ &= ~WATER_HEATER_STATE_AWAY;
+    this->state_mask_ &= ~WATER_HEATER_STATE_AWAY;
   }
   // If ON/OFF not supported, device is always on - clear the flag silently
   if (!traits.has_feature_flags(WATER_HEATER_SUPPORTS_ON_OFF)) {
     this->state_ &= ~WATER_HEATER_STATE_ON;
+    this->state_mask_ &= ~WATER_HEATER_STATE_ON;
   }
 }
 
diff --git a/esphome/components/water_heater/water_heater.h b/esphome/components/water_heater/water_heater.h
index 7bd05ba7f5d..93fcf5f401f 100644
--- a/esphome/components/water_heater/water_heater.h
+++ b/esphome/components/water_heater/water_heater.h
@@ -91,6 +91,8 @@ class WaterHeaterCall {
   float get_target_temperature_high() const { return this->target_temperature_high_; }
   /// Get state flags value
   uint32_t get_state() const { return this->state_; }
+  /// Get mask of state flags that are being changed
+  uint32_t get_state_mask() const { return this->state_mask_; }
 
  protected:
   void validate_();
@@ -100,6 +102,7 @@ class WaterHeaterCall {
   float target_temperature_low_{NAN};
   float target_temperature_high_{NAN};
   uint32_t state_{0};
+  uint32_t state_mask_{0};
 };
 
 struct WaterHeaterCallInternal : public WaterHeaterCall {
@@ -111,6 +114,7 @@ struct WaterHeaterCallInternal : public WaterHeaterCall {
     this->target_temperature_low_ = restore.target_temperature_low_;
     this->target_temperature_high_ = restore.target_temperature_high_;
     this->state_ = restore.state_;
+    this->state_mask_ = restore.state_mask_;
     return *this;
   }
 };