Merge branch 'partition-table-ota' into integration

esphome/__main__.py

@@ -1092,6 +1092,15 @@ def upload_program(
 
     port_type = get_port_type(host)
 
+    # MQTT and MQTTIP are also OTA paths; MQTTIP gets resolved to a real IP later by
+    # _resolve_network_devices(). Only SERIAL and BOOTSEL are non-OTA upload paths.
+    if port_type in (PortType.SERIAL, PortType.BOOTSEL) and getattr(
+        args, "partition_table", False
+    ):
+        raise EsphomeError(
+            "The option --partition-table can only be used for Over The Air updates."
+        )
+
     if port_type == PortType.BOOTSEL:
         exit_code = upload_using_picotool(config)
         # Return None for device - BOOTSEL can't be used for logging,
@@ -1132,12 +1141,80 @@ def upload_program(
 
     binary = CORE.firmware_bin
     ota_type = espota2.OTA_TYPE_UPDATE_APP
+    if getattr(args, "partition_table", False):
+        # Fail fast if the resolved ESPHome OTA config does not enable allow_partition_access.
+        # The device-side handshake also rejects this with "Device only supports app updates",
+        # but checking here surfaces the misconfiguration before opening a network connection.
+        if not ota_conf.get("allow_partition_access"):
+            raise EsphomeError(
+                "The option --partition-table requires 'allow_partition_access: true' on the "
+                "esphome OTA platform in the device's YAML configuration. Add it, recompile, "
+                "flash a build with the option enabled, and then retry --partition-table."
+            )
+        binary = CORE.partition_table_bin
+        ota_type = espota2.OTA_TYPE_UPDATE_PARTITION_TABLE
     if getattr(args, "file", None) is not None:
         binary = Path(args.file)
 
+    if ota_type == espota2.OTA_TYPE_UPDATE_PARTITION_TABLE:
+        _validate_partition_table_binary(binary)
+
     return espota2.run_ota(network_devices, remote_port, password, binary, ota_type)
 
 
+# Layout of esp_partition_info_t on flash. Each entry is 32 bytes, leading with a
+# 16-bit little-endian magic. ESP-IDF defines ESP_PARTITION_MAGIC = 0x50AA (stored as
+# bytes 0xAA, 0x50) for partition entries and ESP_PARTITION_MAGIC_MD5 = 0xEBEB for the
+# trailing checksum entry. Padding past the last entry is 0xFF. The full table is
+# exactly ESP_PARTITION_TABLE_MAX_LEN bytes.
+_PARTITION_TABLE_MAX_LEN = 0xC00
+_ESP_PARTITION_MAGIC = 0x50AA
+_ESP_PARTITION_MAGIC_MD5 = 0xEBEB
+
+
+def _validate_partition_table_binary(binary: Path) -> None:
+    """Validate that ``binary`` looks like an ESP32 partition table image.
+
+    Catches common mistakes (wrong file, truncated build output, swapped --file path)
+    before opening a network connection so the failure mode is a clear local error
+    instead of a post-handshake device rejection.
+    """
+    try:
+        data = binary.read_bytes()
+    except OSError as err:
+        raise EsphomeError(
+            f"Cannot read partition table file '{binary}': {err}"
+        ) from err
+
+    if len(data) != _PARTITION_TABLE_MAX_LEN:
+        raise EsphomeError(
+            f"Partition table file '{binary}' has wrong size: expected "
+            f"{_PARTITION_TABLE_MAX_LEN} bytes, got {len(data)}. "
+            "Pass the partition table image (e.g. partitions.bin / partition-table.bin), "
+            "not the firmware image."
+        )
+
+    first_magic = data[0] | (data[1] << 8)
+    if first_magic != _ESP_PARTITION_MAGIC:
+        raise EsphomeError(
+            f"Partition table file '{binary}' does not start with the expected "
+            f"partition magic 0x{_ESP_PARTITION_MAGIC:04X} (got 0x{first_magic:04X}). "
+            "This file does not look like an ESP32 partition table."
+        )
+
+    # The MD5 checksum entry is required: without it the device-side
+    # esp_partition_table_verify will accept the table but the bootloader will
+    # refuse to boot from it. Scan the 32-byte entries for the MD5 magic.
+    if not any(
+        (data[off] | (data[off + 1] << 8)) == _ESP_PARTITION_MAGIC_MD5
+        for off in range(0, _PARTITION_TABLE_MAX_LEN, 32)
+    ):
+        raise EsphomeError(
+            f"Partition table file '{binary}' is missing the MD5 checksum entry. "
+            "Regenerate the partition table with gen_esp32part.py or rebuild the project."
+        )
+
+
 def show_logs(config: ConfigType, args: ArgsProtocol, devices: list[str]) -> int | None:
     try:
         module = importlib.import_module("esphome.components." + CORE.target_platform)
@@ -1805,6 +1882,11 @@ def parse_args(argv):
         "--file",
         help="Manually specify the binary file to upload.",
     )
+    parser_upload.add_argument(
+        "--partition-table",
+        help="Upload as partition table (OTA).",
+        action="store_true",
+    )
 
     parser_logs = subparsers.add_parser(
         "logs",

esphome/components/esphome/ota/__init__.py

@@ -16,11 +16,13 @@ from esphome.const import (
     CONF_SAFE_MODE,
     CONF_VERSION,
 )
-from esphome.core import coroutine_with_priority
+from esphome.core import CORE, coroutine_with_priority
 from esphome.coroutine import CoroPriority
 import esphome.final_validate as fv
 from esphome.types import ConfigType
 
+CONF_ALLOW_PARTITION_ACCESS = "allow_partition_access"
+
 _LOGGER = logging.getLogger(__name__)
 
 
@@ -75,6 +77,10 @@ def ota_esphome_final_validate(config):
                 merged_ota_esphome_configs_by_port[conf_port] = merge_config(
                     merged_ota_esphome_configs_by_port[conf_port], ota_conf
                 )
+            if ota_conf.get(CONF_ALLOW_PARTITION_ACCESS) and not CORE.is_esp32:
+                raise cv.Invalid(
+                    f"{CONF_ALLOW_PARTITION_ACCESS} is only supported on the esp32"
+                )
         else:
             new_ota_conf.append(ota_conf)
 
@@ -125,6 +131,7 @@ CONFIG_SCHEMA = cv.All(
                 ln882x=8820,
                 rtl87xx=8892,
             ): cv.port,
+            cv.Optional(CONF_ALLOW_PARTITION_ACCESS, default=False): cv.boolean,
             cv.Optional(CONF_PASSWORD): cv.string,
             cv.Optional(CONF_NUM_ATTEMPTS): cv.invalid(
                 f"'{CONF_SAFE_MODE}' (and its related configuration variables) has moved from 'ota' to its own component. See https://esphome.io/components/safe_mode"
@@ -159,6 +166,10 @@ async def to_code(config: ConfigType) -> None:
         if config[CONF_PASSWORD]:
             cg.add(var.set_auth_password(config[CONF_PASSWORD]))
     cg.add_define("USE_OTA_VERSION", config[CONF_VERSION])
+
+    if config.get(CONF_ALLOW_PARTITION_ACCESS):
+        cg.add_define("USE_OTA_PARTITIONS")
+
     # Build flag so lwip_fast_select.c (a .c file that can't include defines.h) sees it.
     cg.add_build_flag("-DUSE_OTA_PLATFORM_ESPHOME")
 

esphome/components/esphome/ota/ota_esphome.cpp

@@ -87,6 +87,10 @@ void ESPHomeOTAComponent::setup() {
   // no wakes fire and loop() falls back to the self-disable safety net.
   esphome_fast_select_set_ota_listener_sock(esphome_lwip_get_sock(this->server_->get_fd()));
 #endif
+
+#ifdef USE_OTA_PARTITIONS
+  ota::get_running_app_position(this->running_app_offset_, this->running_app_size_);
+#endif
 }
 
 void ESPHomeOTAComponent::dump_config() {
@@ -100,6 +104,29 @@ void ESPHomeOTAComponent::dump_config() {
     ESP_LOGCONFIG(TAG, "  Password configured");
   }
 #endif
+#ifdef USE_OTA_PARTITIONS
+  ESP_LOGCONFIG(TAG,
+                "  Partition access allowed\n"
+                "  Running app:\n"
+                "    Partition address: 0x%X\n"
+                "    Used size: %zu bytes (0x%X)",
+                this->running_app_offset_, this->running_app_size_, this->running_app_size_);
+
+#ifdef USE_ESP32
+  ESP_LOGCONFIG(TAG,
+                "  Partition table:\n"
+                "    %-12s %-4s %-8s %-10s %-10s",
+                "Name", "Type", "Subtype", "Address", "Size");
+  esp_partition_iterator_t it = esp_partition_find(ESP_PARTITION_TYPE_ANY, ESP_PARTITION_SUBTYPE_ANY, NULL);
+  while (it != NULL) {
+    const esp_partition_t *partition = esp_partition_get(it);
+    ESP_LOGCONFIG(TAG, "    %-12s 0x%-2X 0x%-6X 0x%-8" PRIX32 " 0x%-8" PRIX32, partition->label, partition->type,
+                  partition->subtype, partition->address, partition->size);
+    it = esp_partition_next(it);
+  }
+  esp_partition_iterator_release(it);
+#endif
+#endif
 }
 
 void ESPHomeOTAComponent::loop() {
@@ -118,6 +145,7 @@ static constexpr uint8_t CLIENT_FEATURE_SUPPORTS_COMPRESSION = 0x01;
 static constexpr uint8_t CLIENT_FEATURE_SUPPORTS_SHA256_AUTH = 0x02;
 static constexpr uint8_t CLIENT_FEATURE_SUPPORTS_EXTENDED_PROTOCOL = 0x04;
 static constexpr uint8_t SERVER_FEATURE_SUPPORTS_COMPRESSION = 0x01;
+static constexpr uint8_t SERVER_FEATURE_SUPPORTS_PARTITION_ACCESS = 0x02;
 
 void ESPHomeOTAComponent::handle_handshake_() {
   /// Handle the OTA handshake and authentication.
@@ -215,6 +243,9 @@ void ESPHomeOTAComponent::handle_handshake_() {
         static_assert(HANDSHAKE_BUF_SIZE >= 2, "handshake_buf_ must hold the 2-byte extended-protocol feature ack");
         this->handshake_buf_[0] = ota::OTA_RESPONSE_FEATURE_FLAGS;
         this->handshake_buf_[1] = (supports_compression ? SERVER_FEATURE_SUPPORTS_COMPRESSION : 0);
+#ifdef USE_OTA_PARTITIONS
+        this->handshake_buf_[1] |= SERVER_FEATURE_SUPPORTS_PARTITION_ACCESS;
+#endif
       } else {
         this->handshake_buf_[0] =
             supports_compression ? ota::OTA_RESPONSE_SUPPORTS_COMPRESSION : ota::OTA_RESPONSE_HEADER_OK;
@@ -347,10 +378,12 @@ void ESPHomeOTAComponent::handle_data_() {
              (static_cast<size_t>(buf[2]) << 8) | buf[3];
   ESP_LOGV(TAG, "Size is %u bytes", ota_size);
 
+#ifndef USE_OTA_PARTITIONS
   if (ota_type != ota::OTA_TYPE_UPDATE_APP) {
     error_code = ota::OTA_RESPONSE_ERROR_UNSUPPORTED_OTA_TYPE;
     goto error;  // NOLINT(cppcoreguidelines-avoid-goto)
   }
+#endif
 
   // Now that we've passed authentication and are actually
   // starting the update, set the warning status and notify
@@ -362,8 +395,8 @@ void ESPHomeOTAComponent::handle_data_() {
   this->notify_state_(ota::OTA_STARTED, 0.0f, 0);
 #endif
 
-  // This will block for a few seconds as it locks flash
-  error_code = this->backend_->begin(ota_size);
+  // begin() may block for a few seconds while it locks flash.
+  error_code = this->backend_->begin(ota_size, ota_type);
   if (error_code != ota::OTA_RESPONSE_OK)
     goto error;  // NOLINT(cppcoreguidelines-avoid-goto)
   update_started = true;
@@ -465,6 +498,13 @@ void ESPHomeOTAComponent::handle_data_() {
   this->notify_state_(ota::OTA_COMPLETED, 100.0f, 0);
 #endif
   delay(100);  // NOLINT
+#ifdef USE_OTA_PARTITIONS
+  if (ota_type == ota::OTA_TYPE_UPDATE_PARTITION_TABLE) {
+    // Skip on_safe_shutdown: nvs_flash_deinit() has already invalidated open NVS handles, so
+    // preferences flush would emit ESP_ERR_NVS_INVALID_HANDLE for every entry. Reboot directly.
+    App.reboot();
+  }
+#endif
   App.safe_reboot();
 
 error:

esphome/components/esphome/ota/ota_esphome.h

@@ -98,6 +98,10 @@ class ESPHomeOTAComponent final : public ota::OTAComponent {
 
   uint32_t client_connect_time_{0};
   static constexpr size_t HANDSHAKE_BUF_SIZE = 5;
+#ifdef USE_OTA_PARTITIONS
+  uint32_t running_app_offset_{0};
+  size_t running_app_size_{0};
+#endif
   uint16_t port_;
   uint8_t handshake_buf_[HANDSHAKE_BUF_SIZE];
   OTAState ota_state_{OTAState::IDLE};

esphome/components/ota/ota_backend.h

@@ -42,6 +42,8 @@ enum OTAResponseTypes {
   OTA_RESPONSE_ERROR_RP2040_NOT_ENOUGH_SPACE = 0x8C,
   OTA_RESPONSE_ERROR_SIGNATURE_INVALID = 0x8D,
   OTA_RESPONSE_ERROR_UNSUPPORTED_OTA_TYPE = 0x8E,
+  OTA_RESPONSE_ERROR_PARTITION_TABLE_VERIFY = 0x8F,
+  OTA_RESPONSE_ERROR_PARTITION_TABLE_UPDATE = 0x90,
   OTA_RESPONSE_ERROR_UNKNOWN = 0xFF,
 };
 
@@ -55,6 +57,7 @@ enum OTAState {
 
 enum OTAType : uint8_t {
   OTA_TYPE_UPDATE_APP = 0x00,
+  OTA_TYPE_UPDATE_PARTITION_TABLE = 0x01,
 };
 
 /** Listener interface for OTA state changes.

esphome/components/ota/ota_backend_arduino_libretiny.cpp

@@ -13,7 +13,10 @@ static const char *const TAG = "ota.arduino_libretiny";
 
 std::unique_ptr<ArduinoLibreTinyOTABackend> make_ota_backend() { return make_unique<ArduinoLibreTinyOTABackend>(); }
 
-OTAResponseTypes ArduinoLibreTinyOTABackend::begin(size_t image_size) {
+OTAResponseTypes ArduinoLibreTinyOTABackend::begin(size_t image_size, OTAType ota_type) {
+  if (ota_type != OTA_TYPE_UPDATE_APP) {
+    return OTA_RESPONSE_ERROR_UNSUPPORTED_OTA_TYPE;
+  }
   // Handle UPDATE_SIZE_UNKNOWN (0) which is used by web server OTA
   // where the exact firmware size is unknown due to multipart encoding
   if (image_size == 0) {

esphome/components/ota/ota_backend_arduino_libretiny.h

@@ -8,7 +8,7 @@ namespace esphome::ota {
 
 class ArduinoLibreTinyOTABackend final {
  public:
-  OTAResponseTypes begin(size_t image_size);
+  OTAResponseTypes begin(size_t image_size, OTAType ota_type = OTA_TYPE_UPDATE_APP);
   void set_update_md5(const char *md5);
   OTAResponseTypes write(uint8_t *data, size_t len);
   OTAResponseTypes end();

esphome/components/ota/ota_backend_arduino_rp2040.cpp

@@ -15,7 +15,10 @@ static const char *const TAG = "ota.arduino_rp2040";
 
 std::unique_ptr<ArduinoRP2040OTABackend> make_ota_backend() { return make_unique<ArduinoRP2040OTABackend>(); }
 
-OTAResponseTypes ArduinoRP2040OTABackend::begin(size_t image_size) {
+OTAResponseTypes ArduinoRP2040OTABackend::begin(size_t image_size, OTAType ota_type) {
+  if (ota_type != OTA_TYPE_UPDATE_APP) {
+    return OTA_RESPONSE_ERROR_UNSUPPORTED_OTA_TYPE;
+  }
   // OTA size of 0 is not currently handled, but
   // web_server is not supported for RP2040, so this is not an issue.
   bool ret = Update.begin(image_size, U_FLASH);

esphome/components/ota/ota_backend_arduino_rp2040.h

@@ -10,7 +10,7 @@ namespace esphome::ota {
 
 class ArduinoRP2040OTABackend final {
  public:
-  OTAResponseTypes begin(size_t image_size);
+  OTAResponseTypes begin(size_t image_size, OTAType ota_type = OTA_TYPE_UPDATE_APP);
   void set_update_md5(const char *md5);
   OTAResponseTypes write(uint8_t *data, size_t len);
   OTAResponseTypes end();

esphome/components/ota/ota_backend_esp8266.cpp

@@ -50,7 +50,10 @@ static const char *const TAG = "ota.esp8266";
 
 std::unique_ptr<ESP8266OTABackend> make_ota_backend() { return make_unique<ESP8266OTABackend>(); }
 
-OTAResponseTypes ESP8266OTABackend::begin(size_t image_size) {
+OTAResponseTypes ESP8266OTABackend::begin(size_t image_size, OTAType ota_type) {
+  if (ota_type != OTA_TYPE_UPDATE_APP) {
+    return OTA_RESPONSE_ERROR_UNSUPPORTED_OTA_TYPE;
+  }
   // Handle UPDATE_SIZE_UNKNOWN (0) by calculating available space
   if (image_size == 0) {
     // Round down to sector boundary: subtract one sector, then mask to sector alignment

esphome/components/ota/ota_backend_esp8266.h

@@ -14,7 +14,7 @@ namespace esphome::ota {
 /// by not having a global Update object in .bss.
 class ESP8266OTABackend final {
  public:
-  OTAResponseTypes begin(size_t image_size);
+  OTAResponseTypes begin(size_t image_size, OTAType ota_type = OTA_TYPE_UPDATE_APP);
   void set_update_md5(const char *md5);
   OTAResponseTypes write(uint8_t *data, size_t len);
   OTAResponseTypes end();

esphome/components/ota/ota_backend_esp_idf.cpp

@@ -16,7 +16,30 @@ static const char *const TAG = "ota.idf";
 
 std::unique_ptr<IDFOTABackend> make_ota_backend() { return make_unique<IDFOTABackend>(); }
 
-OTAResponseTypes IDFOTABackend::begin(size_t image_size) {
+OTAResponseTypes IDFOTABackend::begin(size_t image_size, ota::OTAType ota_type) {
+#ifdef USE_OTA_PARTITIONS
+  this->ota_type_ = ota_type;
+  if (this->ota_type_ == ota::OTA_TYPE_UPDATE_PARTITION_TABLE) {
+    // Reject any size other than ESP_PARTITION_TABLE_MAX_LEN: under- leaves stale bytes from the
+    // previous table; over- can't fit the reserved region.
+    if (image_size != ESP_PARTITION_TABLE_MAX_LEN) {
+      ESP_LOGE(TAG, "Wrong partition table size: expected %u bytes, got %zu", ESP_PARTITION_TABLE_MAX_LEN, image_size);
+      return OTA_RESPONSE_ERROR_PARTITION_TABLE_VERIFY;
+    }
+    memset(this->buf_, 0xFF, sizeof this->buf_);
+    this->buf_written_ = 0;
+    this->image_size_ = image_size;
+    this->md5_.init();
+    return OTA_RESPONSE_OK;
+  }
+  if (this->ota_type_ != ota::OTA_TYPE_UPDATE_APP) {
+    return OTA_RESPONSE_ERROR_UNSUPPORTED_OTA_TYPE;
+  }
+#else
+  if (ota_type != ota::OTA_TYPE_UPDATE_APP) {
+    return OTA_RESPONSE_ERROR_UNSUPPORTED_OTA_TYPE;
+  }
+#endif
 #ifdef USE_OTA_ROLLBACK
   // If we're starting an OTA, the current boot is good enough - mark it valid
   // to prevent rollback and allow the OTA to proceed even if the safe mode
@@ -52,6 +75,21 @@ void IDFOTABackend::set_update_md5(const char *expected_md5) {
 }
 
 OTAResponseTypes IDFOTABackend::write(uint8_t *data, size_t len) {
+#ifdef USE_OTA_PARTITIONS
+  if (this->ota_type_ == ota::OTA_TYPE_UPDATE_PARTITION_TABLE) {
+    if (len > PARTITION_TABLE_BUFFER_SIZE - this->buf_written_) {
+      ESP_LOGE(TAG, "Wrong partition table size");
+      return OTA_RESPONSE_ERROR_PARTITION_TABLE_VERIFY;
+    }
+    memcpy(this->buf_ + this->buf_written_, data, len);
+    this->buf_written_ += len;
+    this->md5_.add(data, len);
+    return OTA_RESPONSE_OK;
+  }
+  if (this->ota_type_ != ota::OTA_TYPE_UPDATE_APP) {
+    return OTA_RESPONSE_ERROR_UNSUPPORTED_OTA_TYPE;
+  }
+#endif
   esp_err_t err = esp_ota_write(this->update_handle_, data, len);
   this->md5_.add(data, len);
   if (err != ESP_OK) {
@@ -73,6 +111,14 @@ OTAResponseTypes IDFOTABackend::end() {
       return OTA_RESPONSE_ERROR_MD5_MISMATCH;
     }
   }
+#ifdef USE_OTA_PARTITIONS
+  if (this->ota_type_ == ota::OTA_TYPE_UPDATE_PARTITION_TABLE) {
+    return this->update_partition_table();
+  }
+  if (this->ota_type_ != ota::OTA_TYPE_UPDATE_APP) {
+    return OTA_RESPONSE_ERROR_UNSUPPORTED_OTA_TYPE;
+  }
+#endif
   esp_err_t err = esp_ota_end(this->update_handle_);
   this->update_handle_ = 0;
   if (err == ESP_OK) {
@@ -96,6 +142,14 @@ OTAResponseTypes IDFOTABackend::end() {
 }
 
 void IDFOTABackend::abort() {
+#ifdef USE_OTA_PARTITIONS
+  if (this->partition_table_part_ != nullptr) {
+    esp_partition_deregister_external(this->partition_table_part_);
+    this->partition_table_part_ = nullptr;
+  }
+#endif
+  // esp_ota_abort with handle 0 returns ESP_ERR_INVALID_ARG harmlessly, so this is safe whether
+  // or not an update is in flight.
   esp_ota_abort(this->update_handle_);
   this->update_handle_ = 0;
 }

esphome/components/ota/ota_backend_esp_idf.h

@@ -9,21 +9,54 @@
 
 namespace esphome::ota {
 
+#ifdef USE_OTA_PARTITIONS
+// Staging buffer holds the entire partition table for verification before any flash op.
+static constexpr size_t PARTITION_TABLE_BUFFER_SIZE = ESP_PARTITION_TABLE_MAX_LEN;  // 0xC00
+
+void get_running_app_position(uint32_t &offset, size_t &size);
+#endif
+
 class IDFOTABackend final {
  public:
-  OTAResponseTypes begin(size_t image_size);
+  OTAResponseTypes begin(size_t image_size, ota::OTAType ota_type = ota::OTA_TYPE_UPDATE_APP);
   void set_update_md5(const char *md5);
   OTAResponseTypes write(uint8_t *data, size_t len);
   OTAResponseTypes end();
   void abort();
   bool supports_compression() { return false; }
 
+ protected:
+#ifdef USE_OTA_PARTITIONS
+  // copy_dest_part non-null means the running app must be copied INTO this slot of the current
+  // table before the new partition table is committed. The destination is in the current table
+  // because that's where esp_partition_copy can write; once the new table replaces it, the same
+  // flash region becomes target_app_index in the new table.
+  struct PartitionTablePlan {
+    int target_app_index{-1};
+    const esp_partition_t *copy_dest_part{nullptr};
+  };
+
+  OTAResponseTypes validate_new_partition_table_(uint32_t running_app_offset, size_t running_app_size,
+                                                 PartitionTablePlan &plan);
+  OTAResponseTypes update_partition_table();
+#endif
+
  private:
   esp_ota_handle_t update_handle_{0};
   const esp_partition_t *partition_;
   md5::MD5Digest md5_{};
   char expected_bin_md5_[32];
   bool md5_set_{false};
+#ifdef USE_OTA_PARTITIONS
+  // Buffer first so it packs tightly after the preceding `bool md5_set_` with no alignment
+  // padding. Only resident during an active OTA: the backend is constructed per connection and
+  // destroyed on cleanup_connection_().
+  uint8_t buf_[PARTITION_TABLE_BUFFER_SIZE];
+  size_t buf_written_{0};
+  size_t image_size_{0};
+  const esp_partition_t *partition_table_part_{nullptr};
+  ota::OTAType ota_type_{ota::OTA_TYPE_UPDATE_APP};
+#endif
 };
 
 std::unique_ptr<IDFOTABackend> make_ota_backend();

esphome/components/ota/ota_backend_host.cpp

@@ -10,7 +10,9 @@ namespace esphome::ota {
 
 std::unique_ptr<HostOTABackend> make_ota_backend() { return make_unique<HostOTABackend>(); }
 
-OTAResponseTypes HostOTABackend::begin(size_t image_size) { return OTA_RESPONSE_ERROR_UPDATE_PREPARE; }
+OTAResponseTypes HostOTABackend::begin(size_t image_size, OTAType ota_type) {
+  return OTA_RESPONSE_ERROR_UPDATE_PREPARE;
+}
 
 void HostOTABackend::set_update_md5(const char *expected_md5) {}
 

esphome/components/ota/ota_backend_host.h

@@ -9,7 +9,7 @@ namespace esphome::ota {
 /// OTA triggers to compile for host platform during development.
 class HostOTABackend final {
  public:
-  OTAResponseTypes begin(size_t image_size);
+  OTAResponseTypes begin(size_t image_size, OTAType ota_type = OTA_TYPE_UPDATE_APP);
   void set_update_md5(const char *md5);
   OTAResponseTypes write(uint8_t *data, size_t len);
   OTAResponseTypes end();

esphome/components/ota/ota_partitions_esp_idf.cpp

@@ -0,0 +1,327 @@
+#ifdef USE_ESP32
+#include "ota_backend_esp_idf.h"
+
+#include "esphome/core/defines.h"
+
+#ifdef USE_OTA_PARTITIONS
+#include "esphome/components/watchdog/watchdog.h"
+#include "esphome/core/log.h"
+
+#include <esp_image_format.h>
+#include <esp_ota_ops.h>
+#include <nvs_flash.h>
+
+#include <cstring>
+
+namespace esphome::ota {
+
+static const char *const TAG = "ota.idf";
+
+static inline bool check_overlap(uint32_t a_offset, size_t a_size, uint32_t b_offset, size_t b_size) {
+  return (a_offset + a_size > b_offset && b_offset + b_size > a_offset);
+}
+
+// Wraps esp_partition_find/_get/_next/_release. Returns nullptr if no APP partition at `address`
+// is at least `min_size` bytes.
+static const esp_partition_t *find_app_partition_at(uint32_t address, size_t min_size) {
+  const esp_partition_t *found = nullptr;
+  esp_partition_iterator_t it = esp_partition_find(ESP_PARTITION_TYPE_APP, ESP_PARTITION_SUBTYPE_ANY, nullptr);
+  while (it != nullptr) {
+    const esp_partition_t *p = esp_partition_get(it);
+    if (p->address == address && p->size >= min_size) {
+      found = p;
+      break;
+    }
+    it = esp_partition_next(it);
+  }
+  esp_partition_iterator_release(it);
+  return found;
+}
+
+// Validates the staged partition table and picks the post-update boot slot. All non-destructive
+// checks live here; the destructive write is in update_partition_table().
+// Side effect: registers the live partition-table region as partition_table_part_ so the caller
+// can write to it; abort() releases it on error.
+OTAResponseTypes IDFOTABackend::validate_new_partition_table_(uint32_t running_app_offset, size_t running_app_size,
+                                                              PartitionTablePlan &plan) {
+  esp_err_t err = esp_partition_register_external(
+      nullptr, ESP_PRIMARY_PARTITION_TABLE_OFFSET, ESP_PARTITION_TABLE_SIZE, "PrimaryPrtTable",
+      ESP_PARTITION_TYPE_PARTITION_TABLE, ESP_PARTITION_SUBTYPE_PARTITION_TABLE_PRIMARY, &this->partition_table_part_);
+  if (err != ESP_OK) {
+    ESP_LOGE(TAG, "esp_partition_register_external failed (err=0x%X)", err);
+    return OTA_RESPONSE_ERROR_PARTITION_TABLE_VERIFY;
+  }
+
+  int num_partitions = 0;
+  const esp_partition_info_t *existing_partition_table = nullptr;
+  esp_partition_mmap_handle_t partition_table_map;
+  err = esp_partition_mmap(this->partition_table_part_, 0, ESP_PARTITION_TABLE_MAX_LEN, ESP_PARTITION_MMAP_DATA,
+                           reinterpret_cast<const void **>(&existing_partition_table), &partition_table_map);
+  if (err != ESP_OK) {
+    ESP_LOGE(TAG, "esp_partition_mmap failed (err=0x%X)", err);
+    return OTA_RESPONSE_ERROR_PARTITION_TABLE_VERIFY;
+  }
+  err = esp_partition_table_verify(existing_partition_table, true, &num_partitions);
+  esp_partition_munmap(partition_table_map);
+  if (err != ESP_OK) {
+    ESP_LOGE(TAG, "esp_partition_table_verify failed (existing partition table) (err=0x%X)", err);
+    return OTA_RESPONSE_ERROR_PARTITION_TABLE_VERIFY;
+  }
+
+  const esp_partition_info_t *new_partition_table = reinterpret_cast<const esp_partition_info_t *>(this->buf_);
+  err = esp_partition_table_verify(new_partition_table, true, &num_partitions);
+  if (err != ESP_OK) {
+    ESP_LOGE(TAG, "esp_partition_table_verify failed (new partition table) (err=0x%X)", err);
+    return OTA_RESPONSE_ERROR_PARTITION_TABLE_VERIFY;
+  }
+
+  // esp_partition_table_verify does not catch a missing MD5 entry, but the bootloader refuses
+  // to boot from a table without one.
+  bool checksum_found = false;
+  for (size_t i = 0; i < ESP_PARTITION_TABLE_MAX_ENTRIES; i++) {
+    if (new_partition_table[i].magic == ESP_PARTITION_MAGIC_MD5) {
+      checksum_found = true;
+      break;
+    }
+  }
+  if (!checksum_found) {
+    ESP_LOGE(TAG, "New partition table has no checksum");
+    return OTA_RESPONSE_ERROR_PARTITION_TABLE_VERIFY;
+  }
+
+  // Slot-selection policy when multiple slots can host the running app: pick the FIRST eligible
+  // slot in table order, preferring the no-copy path (matching offset) over the copy path.
+  // Deterministic and table-ordering-stable.
+  int app_partitions_found = 0;
+  int new_app_part_index = -1;
+  int new_app_part_index_with_copy = -1;
+  const esp_partition_t *app_copy_dest_part = nullptr;
+  bool otadata_partition_found = false;
+  bool otadata_overlap = false;
+  bool nvs_partition_found = false;
+  for (int i = 0; i < num_partitions; i++) {
+    const esp_partition_info_t *new_part = &new_partition_table[i];
+    if (new_part->type == ESP_PARTITION_TYPE_APP) {
+      app_partitions_found++;
+      if (new_part->pos.size >= running_app_size) {
+        if (new_part->pos.offset == running_app_offset) {
+          if (new_app_part_index == -1) {
+            new_app_part_index = i;
+          }
+        } else if (new_app_part_index_with_copy == -1 &&
+                   !check_overlap(running_app_offset, running_app_size, new_part->pos.offset, running_app_size)) {
+          // esp_partition_copy writes into a registered partition; need one at this offset in the
+          // current table.
+          const esp_partition_t *p = find_app_partition_at(new_part->pos.offset, running_app_size);
+          if (p != nullptr) {
+            new_app_part_index_with_copy = i;
+            app_copy_dest_part = p;
+          }
+        }
+      }
+    } else if (new_part->type == ESP_PARTITION_TYPE_DATA) {
+      if (new_part->subtype == ESP_PARTITION_SUBTYPE_DATA_OTA) {
+        otadata_partition_found = true;
+        otadata_overlap = check_overlap(running_app_offset, running_app_size, new_part->pos.offset, new_part->pos.size);
+      } else if (new_part->subtype == ESP_PARTITION_SUBTYPE_DATA_NVS &&
+                 strncmp(reinterpret_cast<const char *>(new_part->label), "nvs", sizeof(new_part->label)) == 0) {
+        nvs_partition_found = true;
+      }
+    }
+  }
+
+  if (new_app_part_index == -1 && new_app_part_index_with_copy == -1) {
+    // Most likely cause: the user picked the wrong migration .bin for their running app's size.
+    // Rejecting here is non-destructive (no flash op has run yet); the user can safely retry with
+    // a different .bin. Log enough info that they can pick the right method without guessing.
+    ESP_LOGE(TAG,
+             "Running app at 0x%X (%u bytes used) does not fit any compatible slot in the new "
+             "partition table. Pick a migration method whose size limit is at least %u bytes and "
+             "retry; no flash content was modified.",
+             running_app_offset, running_app_size, running_app_size);
+    return OTA_RESPONSE_ERROR_PARTITION_TABLE_VERIFY;
+  }
+  if (app_partitions_found < 2) {
+    ESP_LOGE(TAG, "New partition table needs at least 2 app partitions, found %d", app_partitions_found);
+    return OTA_RESPONSE_ERROR_PARTITION_TABLE_VERIFY;
+  }
+  if (!otadata_partition_found) {
+    ESP_LOGE(TAG, "New partition table is missing the required otadata partition");
+    return OTA_RESPONSE_ERROR_PARTITION_TABLE_VERIFY;
+  }
+  if (!nvs_partition_found) {
+    ESP_LOGE(TAG, "New partition table is missing the required nvs partition");
+    return OTA_RESPONSE_ERROR_PARTITION_TABLE_VERIFY;
+  }
+  if (otadata_overlap) {
+    ESP_LOGE(TAG,
+             "New otadata partition overlaps with the running app at 0x%X (size %u). The chosen "
+             "partition table is not compatible with this device's current flash layout; pick a "
+             "different migration method.",
+             running_app_offset, running_app_size);
+    return OTA_RESPONSE_ERROR_PARTITION_TABLE_VERIFY;
+  }
+
+  if (new_app_part_index != -1) {
+    plan.target_app_index = new_app_part_index;
+    plan.copy_dest_part = nullptr;
+  } else {
+    plan.target_app_index = new_app_part_index_with_copy;
+    plan.copy_dest_part = app_copy_dest_part;
+  }
+  return OTA_RESPONSE_OK;
+}
+
+OTAResponseTypes IDFOTABackend::update_partition_table() {
+  if (this->buf_written_ == 0 || this->image_size_ != this->buf_written_) {
+    ESP_LOGE(TAG, "Not enough data received");
+    return OTA_RESPONSE_ERROR_PARTITION_TABLE_VERIFY;
+  }
+
+  // Without a valid running-app size we cannot compute overlap or copy bounds. zero indicates
+  // esp_ota_get_running_partition() failed (e.g. cache unloaded by a previous aborted OTA).
+  uint32_t running_app_offset;
+  size_t running_app_size;
+  get_running_app_position(running_app_offset, running_app_size);
+  if (running_app_size == 0) {
+    ESP_LOGE(TAG, "Failed to determine running app position");
+    return OTA_RESPONSE_ERROR_PARTITION_TABLE_VERIFY;
+  }
+
+  PartitionTablePlan plan;
+  OTAResponseTypes validate_result = this->validate_new_partition_table_(running_app_offset, running_app_size, plan);
+  if (validate_result != OTA_RESPONSE_OK) {
+    return validate_result;
+  }
+
+  // ERROR severity so the warning shows up in default log filters; any failure past this point
+  // can leave the device unbootable until it is recovered with a serial flash.
+  ESP_LOGE(TAG, "Starting partition table update.\n"
+                "  DO NOT REMOVE POWER until the device reboots successfully.\n"
+                "  Loss of power during this operation may render the device unable to boot until\n"
+                "  it is recovered via a serial flash.");
+
+  // One guard over the whole critical section in case an IDF call takes longer than expected on
+  // some chip variant.
+  watchdog::WatchdogManager watchdog(15000);
+
+  esp_err_t err;
+  const esp_partition_info_t *new_partition_table = reinterpret_cast<const esp_partition_info_t *>(this->buf_);
+
+  if (plan.copy_dest_part != nullptr) {
+    // Resolve the source via running_app_offset rather than esp_ota_get_running_partition() in
+    // case a prior aborted partition-table OTA called esp_partition_unload_all() in this boot,
+    // which leaves esp_ota_get_running_partition() returning nullptr.
+    const esp_partition_t *running_app_part = find_app_partition_at(running_app_offset, running_app_size);
+    if (running_app_part == nullptr) {
+      ESP_LOGE(TAG, "Cannot resolve running app partition at offset 0x%X", running_app_offset);
+      return OTA_RESPONSE_ERROR_PARTITION_TABLE_UPDATE;
+    }
+    ESP_LOGD(TAG, "Copying running app from 0x%X to 0x%X (size: 0x%X)", running_app_part->address,
+             plan.copy_dest_part->address, running_app_size);
+    err = esp_partition_copy(plan.copy_dest_part, 0, running_app_part, 0, running_app_size);
+    if (err != ESP_OK) {
+      ESP_LOGE(TAG, "esp_partition_copy failed (err=0x%X)", err);
+      return OTA_RESPONSE_ERROR_PARTITION_TABLE_UPDATE;
+    }
+  }
+
+  // Deinit NVS only just before the first destructive write so verify/copy failure paths return
+  // with NVS still functional. From this point on, components that hold open NVS handles
+  // (e.g. preferences) will fail with ESP_ERR_NVS_INVALID_HANDLE on success or failure;
+  // nvs_flash_init() can re-init the subsystem but cannot revive existing handles. On the
+  // success path the device reboots immediately afterwards so this doesn't matter; on the
+  // failure path the user must reboot the device before retrying.
+  nvs_flash_deinit();
+
+  // Update the partition table
+  err = esp_ota_begin(this->partition_table_part_, ESP_PARTITION_TABLE_MAX_LEN, &this->update_handle_);
+  if (err != ESP_OK) {
+    esp_ota_abort(this->update_handle_);
+    this->update_handle_ = 0;
+    ESP_LOGE(TAG, "esp_ota_begin failed (err=0x%X)", err);
+    return OTA_RESPONSE_ERROR_PARTITION_TABLE_UPDATE;
+  }
+  err = esp_ota_write(this->update_handle_, this->buf_, ESP_PARTITION_TABLE_MAX_LEN);
+  if (err != ESP_OK) {
+    esp_ota_abort(this->update_handle_);
+    this->update_handle_ = 0;
+    ESP_LOGE(TAG, "esp_ota_write failed (err=0x%X)", err);
+    return OTA_RESPONSE_ERROR_PARTITION_TABLE_UPDATE;
+  }
+  err = esp_ota_end(this->update_handle_);
+  this->update_handle_ = 0;  // esp_ota_end releases the handle internally
+  if (err != ESP_OK) {
+    ESP_LOGE(TAG, "esp_ota_end failed (err=0x%X)", err);
+    return OTA_RESPONSE_ERROR_PARTITION_TABLE_UPDATE;
+  }
+  // unload first, then null the member pointer; if abort() ran between the two steps it would
+  // see a freed pointer. esp_partition_unload_all() invalidates partition_table_part_ too, so
+  // an explicit deregister would be redundant.
+  esp_partition_unload_all();
+  this->partition_table_part_ = nullptr;
+
+  // Write otadata to set the new boot partition
+  const esp_partition_info_t *new_part = &new_partition_table[plan.target_app_index];
+  const esp_partition_t *new_boot_partition = find_app_partition_at(new_part->pos.offset, 0);
+  if (new_boot_partition == nullptr) {
+    ESP_LOGE(TAG, "Selected app partition not found after partition table update");
+    return OTA_RESPONSE_ERROR_PARTITION_TABLE_UPDATE;
+  }
+  ESP_LOGD(TAG, "Setting next boot partition to 0x%X", new_boot_partition->address);
+  err = esp_ota_set_boot_partition(new_boot_partition);
+  if (err != ESP_OK) {
+    ESP_LOGE(TAG, "esp_ota_set_boot_partition failed (err=0x%X)", err);
+    return OTA_RESPONSE_ERROR_PARTITION_TABLE_UPDATE;
+  }
+  return OTA_RESPONSE_OK;
+}
+
+// Process-scoped cache. Cannot be a backend member: backends are per-connection but the cache
+// must outlive a connection that called esp_partition_unload_all(), after which
+// esp_ota_get_running_partition() no longer returns valid data.
+static bool s_running_app_initialized = false;
+static uint32_t s_running_app_cached_offset = 0;
+static size_t s_running_app_cached_size = 0;
+
+// Flag-gated rather than size==0 so a failed first call doesn't poison the cache.
+void get_running_app_position(uint32_t &offset, size_t &size) {
+  if (!s_running_app_initialized) {
+    const esp_partition_t *running_app_part = esp_ota_get_running_partition();
+    if (running_app_part == nullptr || running_app_part->erase_size == 0) {
+      // Surface zeros without committing to the cache so a later call has a chance to succeed.
+      offset = 0;
+      size = 0;
+      return;
+    }
+
+    uint32_t pending_offset = running_app_part->address;
+    size_t pending_size = running_app_part->size;
+
+    const esp_partition_pos_t running_app_pos = {
+        .offset = running_app_part->address,
+        .size = running_app_part->size,
+    };
+    esp_image_metadata_t image_metadata = {};
+    image_metadata.start_addr = running_app_part->address;
+    if (esp_image_verify(ESP_IMAGE_VERIFY_SILENT, &running_app_pos, &image_metadata) == ESP_OK &&
+        image_metadata.image_len < running_app_part->size) {
+      pending_size = image_metadata.image_len;
+    }
+    // Round up to a full flash sector so the copy spans complete erase blocks.
+    pending_size = ((pending_size + running_app_part->erase_size - 1) / running_app_part->erase_size) *
+                   running_app_part->erase_size;
+
+    s_running_app_cached_offset = pending_offset;
+    s_running_app_cached_size = pending_size;
+    s_running_app_initialized = true;
+  }
+
+  offset = s_running_app_cached_offset;
+  size = s_running_app_cached_size;
+}
+
+}  // namespace esphome::ota
+
+#endif  // USE_OTA_PARTITIONS
+#endif  // USE_ESP32

esphome/core/__init__.py

@@ -815,6 +815,17 @@ class EsphomeCore:
             return self.relative_pioenvs_path(self.name, "firmware.uf2")
         return self.relative_pioenvs_path(self.name, "firmware.bin")
 
+    @property
+    def partition_table_bin(self) -> Path:
+        # Native ESP-IDF (--native-idf): the partition table image is emitted under
+        # build/partition_table/partition-table.bin alongside firmware.bin. PlatformIO writes the
+        # equivalent file as partitions.bin in the env-specific .pioenvs directory.
+        if self.data.get(KEY_NATIVE_IDF):
+            return self.relative_build_path(
+                "build", "partition_table", "partition-table.bin"
+            )
+        return self.relative_pioenvs_path(self.name, "partitions.bin")
+
     @property
     def target_platform(self):
         return self.data[KEY_CORE][KEY_TARGET_PLATFORM]

esphome/espota2.py

@@ -16,6 +16,7 @@ from esphome.core import EsphomeError
 from esphome.helpers import ProgressBar, resolve_ip_address
 
 OTA_TYPE_UPDATE_APP = 0x00
+OTA_TYPE_UPDATE_PARTITION_TABLE = 0x01
 
 RESPONSE_OK = 0x00
 RESPONSE_REQUEST_AUTH = 0x01
@@ -46,6 +47,8 @@ RESPONSE_ERROR_MD5_MISMATCH = 0x8B
 RESPONSE_ERROR_RP2040_NOT_ENOUGH_SPACE = 0x8C
 RESPONSE_ERROR_SIGNATURE_INVALID = 0x8D
 RESPONSE_ERROR_UNSUPPORTED_OTA_TYPE = 0x8E
+RESPONSE_ERROR_PARTITION_TABLE_VERIFY = 0x8F
+RESPONSE_ERROR_PARTITION_TABLE_UPDATE = 0x90
 RESPONSE_ERROR_UNKNOWN = 0xFF
 
 OTA_VERSION_1_0 = 1
@@ -62,7 +65,9 @@ SERVER_FEATURE_SUPPORTS_PARTITION_ACCESS = 0x02
 # OTA types this client knows how to send. Future PRs that add bootloader/partition
 # updates extend this set. Anything outside the set is rejected up front so callers
 # of perform_ota/run_ota get a clear error instead of a post-auth 0x8E from the device.
-_SUPPORTED_OTA_TYPES: frozenset[int] = frozenset({OTA_TYPE_UPDATE_APP})
+_SUPPORTED_OTA_TYPES: frozenset[int] = frozenset(
+    {OTA_TYPE_UPDATE_APP, OTA_TYPE_UPDATE_PARTITION_TABLE}
+)
 
 UPLOAD_BLOCK_SIZE = 8192
 UPLOAD_BUFFER_SIZE = UPLOAD_BLOCK_SIZE * 8
@@ -128,6 +133,16 @@ _ERROR_MESSAGES: dict[int, str] = {
     RESPONSE_ERROR_UNSUPPORTED_OTA_TYPE: (
         "The requested OTA type is not supported by the device."
     ),
+    RESPONSE_ERROR_PARTITION_TABLE_VERIFY: (
+        "The partition table update could not be verified. No changes were "
+        "made to the flash content. Check the logs for more information and retry."
+    ),
+    RESPONSE_ERROR_PARTITION_TABLE_UPDATE: (
+        "An error occurred while updating the partition table. The device is now "
+        "in a degraded state (NVS handles are invalid; many components will fail) "
+        "and may not be able to boot. Check the logs, reboot the device, and "
+        "retry the update. If the device fails to boot, recover it via a serial flash."
+    ),
     RESPONSE_ERROR_UNKNOWN: "Unknown error from ESP",
 }
 

tests/components/ota/test-partition_access.esp32-idf.yaml

@@ -0,0 +1,5 @@
+ota:
+  - platform: esphome
+    allow_partition_access: true
+
+<<: !include common.yaml

tests/unit_tests/test_espota2.py

@@ -193,6 +193,14 @@ def test_receive_exactly_socket_error(mock_socket: Mock) -> None:
             espota2.RESPONSE_ERROR_UNSUPPORTED_OTA_TYPE,
             "Error: The requested OTA type is not supported by the device",
         ),
+        (
+            espota2.RESPONSE_ERROR_PARTITION_TABLE_VERIFY,
+            "Error: The partition table update could not be verified",
+        ),
+        (
+            espota2.RESPONSE_ERROR_PARTITION_TABLE_UPDATE,
+            "Error: An error occurred while updating the partition table",
+        ),
         (espota2.RESPONSE_ERROR_UNKNOWN, "Unknown error from ESP"),
     ],
 )
@@ -831,6 +839,66 @@ def test_perform_ota_extended_protocol_app(
     )
 
 
+@pytest.mark.usefixtures("mock_time")
+def test_perform_ota_successful_partition_table(
+    mock_socket: Mock, mock_file: io.BytesIO
+) -> None:
+    """Test OTA partition table update.
+
+    The mocked server advertises both COMPRESSION and PARTITION_ACCESS to exercise
+    the full extended-protocol negotiation path. Real IDFOTABackend devices return
+    ``supports_compression() == false`` and never set the COMPRESSION flag for a
+    partition-table OTA; the flag here is intentional protocol-coverage, not a
+    description of on-device behaviour.
+    """
+    recv_responses = [
+        bytes([espota2.RESPONSE_OK]),  # First byte of version response
+        bytes([espota2.OTA_VERSION_2_0]),  # Version number
+        bytes([espota2.RESPONSE_FEATURE_FLAGS]),  # Device supports extended protocol
+        bytes(
+            [
+                espota2.SERVER_FEATURE_SUPPORTS_COMPRESSION
+                | espota2.SERVER_FEATURE_SUPPORTS_PARTITION_ACCESS
+            ]
+        ),  # Device feature flags (compression flag is unrealistic; see docstring)
+        bytes([espota2.RESPONSE_AUTH_OK]),  # No auth required
+        bytes([espota2.RESPONSE_UPDATE_PREPARE_OK]),  # Binary size OK
+        bytes([espota2.RESPONSE_BIN_MD5_OK]),  # MD5 checksum OK
+        bytes([espota2.RESPONSE_CHUNK_OK]),  # Chunk OK
+        bytes([espota2.RESPONSE_RECEIVE_OK]),  # Receive OK
+        bytes([espota2.RESPONSE_UPDATE_END_OK]),  # Update end OK
+    ]
+
+    mock_socket.recv.side_effect = recv_responses
+
+    espota2.perform_ota(
+        mock_socket,
+        "testpass",
+        mock_file,
+        "partitions.bin",
+        espota2.OTA_TYPE_UPDATE_PARTITION_TABLE,
+    )
+
+    # Verify magic bytes were sent
+    assert mock_socket.sendall.call_args_list[0] == call(bytes(espota2.MAGIC_BYTES))
+
+    # Verify features were sent (compression + SHA256 support + extended protocol)
+    assert mock_socket.sendall.call_args_list[1] == call(
+        bytes(
+            [
+                espota2.CLIENT_FEATURE_SUPPORTS_COMPRESSION
+                | espota2.CLIENT_FEATURE_SUPPORTS_SHA256_AUTH
+                | espota2.CLIENT_FEATURE_SUPPORTS_EXTENDED_PROTOCOL
+            ]
+        )
+    )
+
+    # Verify ota type was sent
+    assert mock_socket.sendall.call_args_list[2] == call(
+        bytes([espota2.OTA_TYPE_UPDATE_PARTITION_TABLE])
+    )
+
+
 @pytest.mark.usefixtures("mock_time")
 def test_perform_ota_device_rejects_with_unsupported_ota_type(
     mock_socket: Mock, mock_file: io.BytesIO

tests/unit_tests/test_main.py

@@ -24,6 +24,7 @@ from esphome.__main__ import (
     _get_configured_xtal_freq,
     _make_crystal_freq_callback,
     _resolve_network_devices,
+    _validate_partition_table_binary,
     choose_upload_log_host,
     command_analyze_memory,
     command_bundle,
@@ -83,7 +84,7 @@ from esphome.const import (
     PLATFORM_RP2040,
 )
 from esphome.core import CORE, EsphomeError
-from esphome.espota2 import OTA_TYPE_UPDATE_APP
+from esphome.espota2 import OTA_TYPE_UPDATE_APP, OTA_TYPE_UPDATE_PARTITION_TABLE
 from esphome.util import BootselResult
 from esphome.zeroconf import _await_discovery, discover_mdns_devices
 
@@ -1112,6 +1113,7 @@ class MockArgs:
     reset: bool = False
     list_only: bool = False
     output: str | None = None
+    partition_table: bool = False
 
 
 def test_upload_program_serial_esp32(
@@ -1629,6 +1631,237 @@ def test_upload_program_ota_with_file_arg(
     )
 
 
+_PARTITION_TABLE_LEN = 0xC00
+
+
+def _make_partition_table_bytes() -> bytes:
+    """Build a minimal partition table image accepted by _validate_partition_table_binary."""
+    table = bytearray(b"\xff" * _PARTITION_TABLE_LEN)
+    # First entry: ESP_PARTITION_MAGIC (0x50AA) little-endian -> bytes 0xAA, 0x50.
+    table[0] = 0xAA
+    table[1] = 0x50
+    # MD5 checksum entry at offset 32: ESP_PARTITION_MAGIC_MD5 (0xEBEB) little-endian.
+    table[32] = 0xEB
+    table[33] = 0xEB
+    return bytes(table)
+
+
+def test_upload_program_ota_partition_table_with_file_arg(
+    mock_run_ota: Mock,
+    mock_get_port_type: Mock,
+    tmp_path: Path,
+) -> None:
+    """Test upload_program with OTA and partition table."""
+    setup_core(platform=PLATFORM_ESP32, tmp_path=tmp_path)
+
+    mock_get_port_type.return_value = "NETWORK"
+    mock_run_ota.return_value = (0, "192.168.1.100")
+
+    partition_file = tmp_path / "partitions.bin"
+    partition_file.write_bytes(_make_partition_table_bytes())
+
+    config = {
+        CONF_OTA: [
+            {
+                CONF_PLATFORM: CONF_ESPHOME,
+                CONF_PORT: 3232,
+                "allow_partition_access": True,
+            }
+        ]
+    }
+    args = MockArgs(file=str(partition_file), partition_table=True)
+    devices = ["192.168.1.100"]
+
+    exit_code, host = upload_program(config, args, devices)
+
+    assert exit_code == 0
+    assert host == "192.168.1.100"
+    mock_run_ota.assert_called_once_with(
+        ["192.168.1.100"],
+        3232,
+        None,
+        partition_file,
+        OTA_TYPE_UPDATE_PARTITION_TABLE,
+    )
+
+
+def test_upload_program_serial_partition_table(
+    mock_upload_using_esptool: Mock,
+    mock_get_port_type: Mock,
+) -> None:
+    """Test serial upload with partition table option (unsupported)."""
+    setup_core(platform=PLATFORM_ESP32)
+    mock_get_port_type.return_value = "SERIAL"
+    mock_upload_using_esptool.return_value = 0
+
+    config = {}
+    args = MockArgs(partition_table=True)
+    devices = ["/dev/ttyUSB0"]
+
+    with pytest.raises(
+        EsphomeError,
+        match="The option --partition-table can only be used for Over The Air updates",
+    ):
+        upload_program(config, args, devices)
+
+
+def test_upload_program_ota_partition_table_mqttip(
+    mock_run_ota: Mock,
+    mock_get_port_type: Mock,
+    tmp_path: Path,
+) -> None:
+    """--partition-table is allowed for MQTTIP devices; they resolve to a real IP at OTA time."""
+    setup_core(platform=PLATFORM_ESP32, tmp_path=tmp_path)
+
+    mock_get_port_type.return_value = "MQTTIP"
+    mock_run_ota.return_value = (0, "192.168.1.100")
+
+    partition_file = tmp_path / "partitions.bin"
+    partition_file.write_bytes(_make_partition_table_bytes())
+
+    config = {
+        CONF_OTA: [
+            {
+                CONF_PLATFORM: CONF_ESPHOME,
+                CONF_PORT: 3232,
+                "allow_partition_access": True,
+            }
+        ]
+    }
+    args = MockArgs(file=str(partition_file), partition_table=True)
+
+    with patch(
+        "esphome.__main__._resolve_network_devices", return_value=["192.168.1.100"]
+    ):
+        exit_code, host = upload_program(config, args, ["MQTTIP"])
+
+    assert exit_code == 0
+    assert host == "192.168.1.100"
+    mock_run_ota.assert_called_once_with(
+        ["192.168.1.100"],
+        3232,
+        None,
+        partition_file,
+        OTA_TYPE_UPDATE_PARTITION_TABLE,
+    )
+
+
+def test_validate_partition_table_binary_accepts_valid(tmp_path: Path) -> None:
+    f = tmp_path / "partitions.bin"
+    f.write_bytes(_make_partition_table_bytes())
+    _validate_partition_table_binary(f)
+
+
+_PARTITION_FIXTURE_DIR = Path(__file__).parent / "fixtures" / "partition_tables"
+
+
+@pytest.mark.parametrize(
+    "fixture",
+    [
+        # Stock ESP-IDF gen_esp32part.py output for an ESPHome build.
+        "esphome_default.bin",
+        # ESP-IDF Hello-world example partition table (vendored from espressif/esp-serial-flasher).
+        "esp_idf_hello_world.bin",
+        # Partition table shipped with esphome_dashboard's prebuilt firmware.
+        "esphome_dashboard_firmware.bin",
+    ],
+)
+def test_validate_partition_table_binary_accepts_real_binaries(fixture: str) -> None:
+    """Real-world partition-table binaries from ESP-IDF / ESPHome tooling pass validation."""
+    _validate_partition_table_binary(_PARTITION_FIXTURE_DIR / fixture)
+
+
+def test_validate_partition_table_binary_rejects_wrong_size(tmp_path: Path) -> None:
+    f = tmp_path / "partitions.bin"
+    f.write_bytes(b"\xaa\x50" + b"\xff" * 100)
+    with pytest.raises(EsphomeError, match="wrong size"):
+        _validate_partition_table_binary(f)
+
+
+def test_validate_partition_table_binary_rejects_wrong_magic(tmp_path: Path) -> None:
+    data = bytearray(_make_partition_table_bytes())
+    data[0] = 0x00
+    data[1] = 0x00
+    f = tmp_path / "partitions.bin"
+    f.write_bytes(bytes(data))
+    with pytest.raises(EsphomeError, match="partition magic"):
+        _validate_partition_table_binary(f)
+
+
+def test_validate_partition_table_binary_rejects_missing_md5(tmp_path: Path) -> None:
+    data = bytearray(_make_partition_table_bytes())
+    data[32] = 0xFF
+    data[33] = 0xFF
+    f = tmp_path / "partitions.bin"
+    f.write_bytes(bytes(data))
+    with pytest.raises(EsphomeError, match="missing the MD5 checksum entry"):
+        _validate_partition_table_binary(f)
+
+
+def test_validate_partition_table_binary_missing_file(tmp_path: Path) -> None:
+    with pytest.raises(EsphomeError, match="Cannot read partition table file"):
+        _validate_partition_table_binary(tmp_path / "does-not-exist.bin")
+
+
+def test_upload_program_ota_partition_table_invalid_file(
+    mock_run_ota: Mock,
+    mock_get_port_type: Mock,
+    tmp_path: Path,
+) -> None:
+    """--partition-table must fail before calling run_ota when the file is not a partition table."""
+    setup_core(platform=PLATFORM_ESP32, tmp_path=tmp_path)
+
+    mock_get_port_type.return_value = "NETWORK"
+
+    bad_file = tmp_path / "firmware.bin"
+    bad_file.write_bytes(b"\x00" * 4096)
+
+    config = {
+        CONF_OTA: [
+            {
+                CONF_PLATFORM: CONF_ESPHOME,
+                CONF_PORT: 3232,
+                "allow_partition_access": True,
+            }
+        ]
+    }
+    args = MockArgs(file=str(bad_file), partition_table=True)
+    devices = ["192.168.1.100"]
+
+    with pytest.raises(EsphomeError, match="wrong size"):
+        upload_program(config, args, devices)
+    mock_run_ota.assert_not_called()
+
+
+def test_upload_program_ota_partition_table_without_allow_flag(
+    mock_run_ota: Mock,
+    mock_get_port_type: Mock,
+    tmp_path: Path,
+) -> None:
+    """--partition-table must fail fast when allow_partition_access is not enabled in YAML."""
+    setup_core(platform=PLATFORM_ESP32, tmp_path=tmp_path)
+
+    mock_get_port_type.return_value = "NETWORK"
+
+    config = {
+        CONF_OTA: [
+            {
+                CONF_PLATFORM: CONF_ESPHOME,
+                CONF_PORT: 3232,
+            }
+        ]
+    }
+    args = MockArgs(file="partitions.bin", partition_table=True)
+    devices = ["192.168.1.100"]
+
+    with pytest.raises(
+        EsphomeError,
+        match="requires 'allow_partition_access: true'",
+    ):
+        upload_program(config, args, devices)
+    mock_run_ota.assert_not_called()
+
+
 def test_upload_program_ota_no_config(
     mock_get_port_type: Mock,
 ) -> None: