Close the final chunk ack retry window and surface pending device errors

esphome/espota2.py

@@ -79,7 +79,9 @@ UPLOAD_BUFFER_SIZE = UPLOAD_BLOCK_SIZE * 8
 # Flaky Wi-Fi links often drop the first OTA attempt, and the device may need time
 # to clean up a half-open connection (its handshake watchdog runs at 20s) before it
 # accepts a new one, so wait between attempts instead of failing the upload outright.
-MAX_UPLOAD_ATTEMPTS = 3
+# Every resolved address is tried once, and this many extra attempts are shared
+# across the addresses on top of that.
+EXTRA_UPLOAD_ATTEMPTS = 2
 UPLOAD_RETRY_DELAY = 5.0
 
 _LOGGER = logging.getLogger(__name__)
@@ -181,6 +183,19 @@ class OTANetworkError(OTAError):
     """Network-level OTA failure (timeout, reset, closed connection); retrying may succeed."""
 
 
+def _committed_error(err: OTANetworkError) -> OTAError:
+    """Wrap a network failure that happened once the device had the full image.
+
+    Past that point the device commits and reboots on its own, so the failure
+    must not be retried; a re-upload could flash a device that already updated.
+    """
+    return OTAError(
+        f"{err} (the device may have already committed the update and "
+        f"be rebooting; check whether it comes back with the new "
+        f"firmware before uploading again)"
+    )
+
+
 def recv_decode(
     sock: socket.socket, amount: int, decode: bool = True
 ) -> bytes | list[int]:
@@ -473,11 +488,25 @@ def perform_ota(
 
             try:
                 sock.sendall(chunk)
-                if version >= OTA_VERSION_2_0:
-                    receive_exactly(sock, 1, "chunk result", RESPONSE_CHUNK_OK)
             except OSError as err:
+                # A send failure can hide an error byte the device reported
+                # just before dropping the connection; surface that as the
+                # real, non-retryable cause when it is available
+                with contextlib.suppress(OSError, OTANetworkError):
+                    sock.settimeout(1.0)
+                    check_error(recv_decode(sock, 1), None)
                 raise OTANetworkError(f"sending data: {err}") from err
 
+            if version >= OTA_VERSION_2_0:
+                try:
+                    receive_exactly(sock, 1, "chunk result", RESPONSE_CHUNK_OK)
+                except OTANetworkError as err:
+                    if offset < upload_size:
+                        raise
+                    # The device already had the complete image when this ack
+                    # was lost, so it may be committing; do not retry
+                    raise _committed_error(err) from err
+
             progress.update(offset / upload_size)
     except OTAError:
         # Terminate the progress bar line before the error is logged
@@ -499,11 +528,7 @@ def perform_ota(
         receive_exactly(sock, 1, "update receive result", RESPONSE_RECEIVE_OK)
         receive_exactly(sock, 1, "update end result", RESPONSE_UPDATE_END_OK)
     except OTANetworkError as err:
-        raise OTAError(
-            f"{err} (the device may have already committed the update and "
-            f"be rebooting; check whether it comes back with the new "
-            f"firmware before uploading again)"
-        ) from err
+        raise _committed_error(err) from err
 
     try:
         send_check(sock, RESPONSE_OK, "end acknowledgement")
@@ -551,13 +576,13 @@ def run_ota_impl_(
         _LOGGER.error("No addresses to connect to for %s", remote_host)
         return 1, None
 
-    # Every address is tried at least once and the budget grants
-    # MAX_UPLOAD_ATTEMPTS - 1 extra retries, cycling through the addresses.
-    # Wait before an attempt when the previous one actually reached the
-    # device, or when revisiting an address, so a flaky link can recover and
-    # the device can clean up a half-open connection (its handshake watchdog
-    # runs at 20s); moving on to the next address family stays immediate.
-    total_attempts = len(res) + MAX_UPLOAD_ATTEMPTS - 1
+    # Every address is tried at least once and EXTRA_UPLOAD_ATTEMPTS retries
+    # are shared across the addresses, cycling through them. Wait before an
+    # attempt when the previous one actually reached the device, or when
+    # revisiting an address, so a flaky link can recover and the device can
+    # clean up a half-open connection (its handshake watchdog runs at 20s);
+    # moving on to the next address family stays immediate.
+    total_attempts = len(res) + EXTRA_UPLOAD_ATTEMPTS
     last_error = ""
     reached_device = False
     for attempt in range(total_attempts):

tests/unit_tests/test_espota2.py

@@ -593,7 +593,10 @@ def _no_auth_handshake(version: int) -> list[bytes]:
 @pytest.mark.usefixtures("mock_time")
 def test_perform_ota_chunk_send_error(mock_socket: Mock, mock_file: io.BytesIO) -> None:
     """Test OTA raises the retryable OTANetworkError when sending a chunk fails."""
-    mock_socket.recv.side_effect = _no_auth_handshake(espota2.OTA_VERSION_2_0)
+    mock_socket.recv.side_effect = [
+        *_no_auth_handshake(espota2.OTA_VERSION_2_0),
+        OSError("Connection reset"),  # Probe for a pending error byte fails too
+    ]
     # Sends before the data phase: magic bytes, features, binary size, MD5;
     # fail on the fifth sendall, the first firmware chunk
     mock_socket.sendall.side_effect = [None] * 4 + [OSError("Broken pipe")]
@@ -602,6 +605,59 @@ def test_perform_ota_chunk_send_error(mock_socket: Mock, mock_file: io.BytesIO)
         espota2.perform_ota(mock_socket, None, mock_file, "test.bin")
 
 
+@pytest.mark.usefixtures("mock_time")
+def test_perform_ota_chunk_send_error_surfaces_device_error(
+    mock_socket: Mock, mock_file: io.BytesIO
+) -> None:
+    """Test a device error byte pending behind a send failure becomes the cause."""
+    mock_socket.recv.side_effect = [
+        *_no_auth_handshake(espota2.OTA_VERSION_1_0),
+        bytes([espota2.RESPONSE_ERROR_WRITING_FLASH]),  # Reason the device closed
+    ]
+    mock_socket.sendall.side_effect = [None] * 4 + [OSError("Broken pipe")]
+
+    with pytest.raises(
+        espota2.OTAError, match="Writing OTA data to flash memory failed"
+    ) as exc:
+        espota2.perform_ota(mock_socket, None, mock_file, "test.bin")
+
+    # The device-reported error is not retryable
+    assert not isinstance(exc.value, espota2.OTANetworkError)
+
+
+@pytest.mark.usefixtures("mock_time")
+def test_perform_ota_final_chunk_ack_failure_not_retryable(
+    mock_socket: Mock, mock_file: io.BytesIO
+) -> None:
+    """Test a lost ack for the final chunk is not retried."""
+    mock_socket.recv.side_effect = [
+        *_no_auth_handshake(espota2.OTA_VERSION_2_0),
+        OSError("Connection reset"),  # Ack for the only (final) chunk is lost
+    ]
+
+    with pytest.raises(espota2.OTAError, match="receiving chunk result") as exc:
+        espota2.perform_ota(mock_socket, None, mock_file, "test.bin")
+
+    # The device already had the whole image, so it may be committing
+    assert not isinstance(exc.value, espota2.OTANetworkError)
+
+
+@pytest.mark.usefixtures("mock_time")
+def test_perform_ota_intermediate_chunk_ack_failure_retryable(
+    mock_socket: Mock,
+) -> None:
+    """Test a lost ack for a non-final chunk stays retryable."""
+    # Two chunks: the firmware is larger than one upload block
+    big_file = io.BytesIO(b"x" * (espota2.UPLOAD_BLOCK_SIZE + 1))
+    mock_socket.recv.side_effect = [
+        *_no_auth_handshake(espota2.OTA_VERSION_2_0),
+        OSError("Connection reset"),  # Ack for the first of two chunks is lost
+    ]
+
+    with pytest.raises(espota2.OTANetworkError, match="receiving chunk result"):
+        espota2.perform_ota(mock_socket, None, big_file, "test.bin")
+
+
 @pytest.mark.usefixtures("mock_time")
 def test_perform_ota_post_commit_failure_not_retryable(
     mock_socket: Mock, mock_file: io.BytesIO
@@ -688,9 +744,9 @@ def test_run_ota_impl_connection_failed(
     assert result_host is None
     # A single address gets the whole attempt budget, with a delay before
     # each revisit
-    assert mock_socket.connect.call_count == espota2.MAX_UPLOAD_ATTEMPTS
-    assert mock_socket.close.call_count == espota2.MAX_UPLOAD_ATTEMPTS
-    assert mock_sleep.call_count == espota2.MAX_UPLOAD_ATTEMPTS - 1
+    assert mock_socket.connect.call_count == espota2.EXTRA_UPLOAD_ATTEMPTS + 1
+    assert mock_socket.close.call_count == espota2.EXTRA_UPLOAD_ATTEMPTS + 1
+    assert mock_sleep.call_count == espota2.EXTRA_UPLOAD_ATTEMPTS
     mock_sleep.assert_called_with(espota2.UPLOAD_RETRY_DELAY)
 
 
@@ -745,8 +801,8 @@ def test_run_ota_impl_network_error_exhausts_attempts(
 
     assert result_code == 1
     assert result_host is None
-    assert mock_perform_ota.call_count == espota2.MAX_UPLOAD_ATTEMPTS
-    assert mock_sleep.call_count == espota2.MAX_UPLOAD_ATTEMPTS - 1
+    assert mock_perform_ota.call_count == espota2.EXTRA_UPLOAD_ATTEMPTS + 1
+    assert mock_sleep.call_count == espota2.EXTRA_UPLOAD_ATTEMPTS
 
 
 @pytest.mark.usefixtures("mock_socket_constructor", "mock_resolve_ip_dual")
@@ -762,8 +818,8 @@ def test_run_ota_impl_multiple_addresses_cycle(
 
     assert result_code == 1
     assert result_host is None
-    # Each address is visited once, then the two spare attempts cycle back
-    # through them; the budget is shared, not per address
+    # Each address is visited once, then the EXTRA_UPLOAD_ATTEMPTS spare
+    # attempts cycle back through them; the budget is shared, not per address
     assert mock_socket.connect.call_args_list == [
         call(DUAL_STACK_SA6),
         call(DUAL_STACK_SA4),