From 6bc3de12412d2169082d04f645592ef9a0c5b4ee Mon Sep 17 00:00:00 2001
From: PJ <ipriyanshujain@gmail.com>
Date: Thu, 23 Apr 2026 17:48:02 +0700
Subject: [PATCH] fix(sidecar): extract warmup + add one-shot WDA reconnect on
 IOException

---
 .../dev/sanderling/sidecar/DriverBackend.kt   | 47 +++++++++++++------
 1 file changed, 33 insertions(+), 14 deletions(-)

diff --git a/sidecar/src/main/kotlin/dev/sanderling/sidecar/DriverBackend.kt b/sidecar/src/main/kotlin/dev/sanderling/sidecar/DriverBackend.kt
index be6147e..517d820 100644
--- a/sidecar/src/main/kotlin/dev/sanderling/sidecar/DriverBackend.kt
+++ b/sidecar/src/main/kotlin/dev/sanderling/sidecar/DriverBackend.kt
@@ -487,7 +487,7 @@ private fun pngHeight(bytes: ByteArray): Int {
 }
 
 class IosDriverBackend(private val udid: String) : DriverBackend {
-    private val driver: maestro.drivers.IOSDriver
+    private lateinit var driver: maestro.drivers.IOSDriver
 
     init {
         val httpClient = xcuitest.api.OkHttpClientInstance.get()
@@ -509,7 +509,10 @@ class IosDriverBackend(private val udid: String) : DriverBackend {
         val device = ios.LocalIOSDevice(udid, xcTestDevice, simctlDevice, maestro.utils.NoopInsights)
         driver = maestro.drivers.IOSDriver(device, maestro.utils.NoopInsights, metrics)
         driver.open()
-        // warm up: absorbs WDA startup race (health check passes before accept() loop is stable)
+        warmup()
+    }
+
+    private fun warmup() {
         var warmupErr: Exception? = null
         repeat(3) { attempt ->
             try {
@@ -524,47 +527,63 @@ class IosDriverBackend(private val udid: String) : DriverBackend {
         warmupErr?.let { throw IllegalStateException("WDA warmup failed after 3 attempts: $it") }
     }
 
-    override fun launch(bundleId: String, clearState: Boolean, env: Map<String, String>) {
+    private fun <T> withReconnect(block: () -> T): T {
+        return try {
+            block()
+        } catch (e: java.io.IOException) {
+            try {
+                driver.open()
+                warmup()
+            } catch (reconnectErr: Exception) {
+                throw IllegalStateException("WDA reconnect failed: $reconnectErr", e)
+            }
+            block()
+        }
+    }
+
+    override fun launch(bundleId: String, clearState: Boolean, env: Map<String, String>) = withReconnect {
         runCatching { driver.stopApp(bundleId) }
         if (clearState) driver.clearAppState(bundleId)
         driver.launchApp(bundleId, env, java.util.UUID.randomUUID())
     }
 
-    override fun terminate(bundleId: String) = driver.stopApp(bundleId)
+    override fun terminate(bundleId: String) = withReconnect { driver.stopApp(bundleId) }
 
-    override fun tap(x: Int, y: Int) = driver.tap(maestro.Point(x, y))
+    override fun tap(x: Int, y: Int) = withReconnect { driver.tap(maestro.Point(x, y)) }
 
-    override fun tapSelector(selector: String) {
+    override fun tapSelector(selector: String) = withReconnect {
         val root = driver.contentDescriptor(false)
-        val bounds = findBoundsBySelector(root, selector) ?: return
+        val bounds = findBoundsBySelector(root, selector) ?: return@withReconnect
         driver.tap(maestro.Point((bounds[0] + bounds[2]) / 2, (bounds[1] + bounds[3]) / 2))
     }
 
-    override fun inputText(text: String) = driver.inputText(text)
+    override fun inputText(text: String) = withReconnect { driver.inputText(text) }
 
-    override fun swipe(fromX: Int, fromY: Int, toX: Int, toY: Int, durationMillis: Long) =
+    override fun swipe(fromX: Int, fromY: Int, toX: Int, toY: Int, durationMillis: Long) = withReconnect {
         driver.swipe(maestro.Point(fromX, fromY), maestro.Point(toX, toY), maxOf(durationMillis, 250L))
+    }
 
-    override fun pressKey(key: String) {
+    override fun pressKey(key: String) = withReconnect {
         StubDriverBackend.KEY_MAP[key]?.let { keyCode ->
             keyCodeToMaestro(keyCode)?.let { driver.pressKey(it) }
         }
     }
 
-    override fun screenshot(): Triple<ByteArray, Int, Int> {
+    override fun screenshot(): Triple<ByteArray, Int, Int> = withReconnect {
         val buf = okio.Buffer()
         driver.takeScreenshot(buf, false)
         val bytes = buf.readByteArray()
-        return Triple(bytes, pngWidth(bytes), pngHeight(bytes))
+        Triple(bytes, pngWidth(bytes), pngHeight(bytes))
     }
 
-    override fun hierarchy(): String =
+    override fun hierarchy(): String = withReconnect {
         com.fasterxml.jackson.module.kotlin.jacksonObjectMapper()
             .writeValueAsString(driver.contentDescriptor(false))
+    }
 
     override fun recentLogs(sinceUnixMillis: Long, minLevel: String): List<LogLine> = emptyList()
 
-    override fun waitForIdle(durationMillis: Long) {
+    override fun waitForIdle(durationMillis: Long) = withReconnect {
         driver.waitForAppToSettle(null, null, durationMillis.toInt())
     }