From fc45e41e383011ba40da0191c608adc5295666ee Mon Sep 17 00:00:00 2001
From: Jake-B <jake-b@users.noreply.github.com>
Date: Mon, 5 Jan 2026 20:39:37 -0500
Subject: [PATCH] Additional OTA robustness improvements

---
 .../ESP32 OTA/BLE/ESP32BLEOTASheet.swift      |  1 +
 .../ESP32 OTA/WiFi/ESP32WifiOTASheet.swift    | 10 ++-------
 .../WiFi/ESP32WifiOTAViewModel.swift          | 22 ++++++++++++++++++-
 3 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/Meshtastic/Views/Settings/Firmware/ESP32 OTA/BLE/ESP32BLEOTASheet.swift b/Meshtastic/Views/Settings/Firmware/ESP32 OTA/BLE/ESP32BLEOTASheet.swift
index 2b7b2249..e9354d50 100644
--- a/Meshtastic/Views/Settings/Firmware/ESP32 OTA/BLE/ESP32BLEOTASheet.swift	
+++ b/Meshtastic/Views/Settings/Firmware/ESP32 OTA/BLE/ESP32BLEOTASheet.swift	
@@ -108,6 +108,7 @@ struct ESP32BLEOTASheet: View {
 			} label: {
 				Label("Retry", systemImage: "arrow.clockwise")
 					.frame(maxWidth: .infinity)
+					.foregroundStyle(.white) 
 			}
 			.buttonStyle(.borderedProminent)
 			.tint(.red)
diff --git a/Meshtastic/Views/Settings/Firmware/ESP32 OTA/WiFi/ESP32WifiOTASheet.swift b/Meshtastic/Views/Settings/Firmware/ESP32 OTA/WiFi/ESP32WifiOTASheet.swift
index c23c205d..8ad61112 100644
--- a/Meshtastic/Views/Settings/Firmware/ESP32 OTA/WiFi/ESP32WifiOTASheet.swift	
+++ b/Meshtastic/Views/Settings/Firmware/ESP32 OTA/WiFi/ESP32WifiOTASheet.swift	
@@ -62,14 +62,7 @@ struct ESP32WifiOTASheet: View {
 							beginWifiProcessButton()
 							
 						case .error:
-							Text("Error: \(ota.errorMessage, default: "Unknown")")
-								.frame(maxWidth: .infinity)
-								.multilineTextAlignment(.center)
-								.font(.headline)
-							
-							Button("Retry") {
-								ota.retry()
-							}
+							retryButton()
 							
 						default:
 							Text("\(ota.statusMessage, default: "")")
@@ -115,6 +108,7 @@ struct ESP32WifiOTASheet: View {
 			} label: {
 				Label("Retry", systemImage: "arrow.clockwise")
 					.frame(maxWidth: .infinity)
+					.foregroundStyle(.white)
 			}
 			.buttonStyle(.borderedProminent)
 			.tint(.red)
diff --git a/Meshtastic/Views/Settings/Firmware/ESP32 OTA/WiFi/ESP32WifiOTAViewModel.swift b/Meshtastic/Views/Settings/Firmware/ESP32 OTA/WiFi/ESP32WifiOTAViewModel.swift
index 4636f194..e819439d 100644
--- a/Meshtastic/Views/Settings/Firmware/ESP32 OTA/WiFi/ESP32WifiOTAViewModel.swift	
+++ b/Meshtastic/Views/Settings/Firmware/ESP32 OTA/WiFi/ESP32WifiOTAViewModel.swift	
@@ -63,6 +63,8 @@ class ESP32WifiOTAViewModel: ObservableObject {
 			try await connectAndUpload(endpoint: targetEndpoint, data: firmwareData)
 			
 			// 3. Success
+			// Explicitly force progress to 1.0 to ensure Checkmark appears
+			progress = 1.0
 			statusMessage = "Success!"
 			otaState = .completed
 			Logger.services.info("[ESP OTA] Update Complete")
@@ -101,7 +103,7 @@ class ESP32WifiOTAViewModel: ObservableObject {
 			let response: String
 			do {
 				// Timeout logic relies on the Reader unblocking immediately on cancel
-				response = try await withTimeout(seconds: 60.0) {
+				response = try await withTimeout(seconds: 30.0) {
 					try await reader.readLine()
 				}
 			} catch {
@@ -149,12 +151,18 @@ class ESP32WifiOTAViewModel: ObservableObject {
 					
 					offset += chunk.count
 					
+					// Update UI periodically to avoid flooding main thread
 					if offset % (self.chunkSize * 10) == 0 {
 						let percent = Double(offset) / Double(totalSize)
 						await self.updateUI { self.progress = percent }
 					}
 				}
 				
+				// MARK: - FIX: Force 100% on upload completion
+				// Because of the modulo operator above, the loop often ends at 99.xxx%.
+				// We force it to 1.0 here so the user sees "100%" while waiting for verification.
+				await self.updateUI { self.progress = 1.0 }
+				
 				isUploading.withLock { $0 = false }
 				Logger.services.info("[ESP OTA] Writer Task: All data sent.")
 			}
@@ -350,6 +358,18 @@ actor AsyncLineReader {
 			return try await withCheckedThrowingContinuation { continuation in
 				lock.withLock { $0 = continuation }
 				
+				// MARK: - FIX for Deadlock
+				// If `onCancel` ran before we set the lock (race condition), `Task.isCancelled`
+				// will be true here. We must abort immediately.
+				if Task.isCancelled {
+					lock.withLock { state in
+						guard let cont = state else { return }
+						state = nil
+						cont.resume(throwing: CancellationError())
+					}
+					return
+				}
+				
 				connection.receive(minimumIncompleteLength: 1, maximumLength: 65536) { data, _, _, error in
 					lock.withLock { state in
 						guard let cont = state else { return } // Already cancelled/resumed