fix: auto reconnect

core/src/commonMain/kotlin/dev/kdriver/core/connection/DefaultConnection.kt

@@ -17,6 +17,7 @@
 import kotlinx.coroutines.sync.withLock
 import kotlinx.serialization.json.JsonElement
 import kotlinx.serialization.json.jsonPrimitive
+import kotlin.concurrent.Volatile
 import kotlin.reflect.KClass
 import kotlin.time.Clock
 import kotlin.time.Duration.Companion.milliseconds
@@ -31,9 +32,29 @@
     override var owner: Browser? = null,
 ) : OwnedConnection {
 
+    companion object {
+        // One initial send plus one resend after a transparent reconnect. The resend is only ever reached
+        // when a send fails (bytes never left the client), so it cannot double-execute a command.
+        private const val SEND_ATTEMPTS = 2
+    }
+
     private val logger = KtorSimpleLogger("Connection")
 
-    private val transport: WebSocketTransport by lazy { createTransport() }
+    // Recreated on every (re)connect. A dropped socket can keep reporting a stale `isActive` for a
+    // short window, so once the receive loop has observed a disconnect we throw the whole transport
+    // away and build a fresh one rather than trusting it (see [connect]). @Volatile because it is
+    // read off-lock on the fast path of [connect], from the receive-loop coroutine in [onTransportGone],
+    // and in [close].
+    @Volatile
+    private var transport: WebSocketTransport? = null
+
+    @Volatile
+    private var disconnected = false
+
+    // Set once [close] has run; terminal. [connect] refuses to re-open a closed connection so a
+    // command issued after close fails instead of silently resurrecting the socket.
+    @Volatile
+    private var closed = false
 
     private var socketSubscription: Job? = null
 
@@ -67,22 +88,43 @@
     @InternalCdpApi
     override val generatedDomains: MutableMap<KClass<out Domain>, Domain> = mutableMapOf()
 
-    private suspend fun connect() {
-        if (transport.isActive) return
+    private suspend fun connect(): WebSocketTransport {
+        if (closed) throw ConnectionClosedException("Connection closed")
+        transport?.let { if (it.isActive && !disconnected) return it }
         // Guard so concurrent first commands don't each open a session (which would leak the
         // duplicate sockets/listeners). Double-checked: skip the lock once connected (ISSUE-4).
-        connectMutex.withLock {
-            if (transport.isActive) return@withLock
-            transport.connect()
-            startListening()
+        return connectMutex.withLock {
+            if (closed) throw ConnectionClosedException("Connection closed")
+            transport?.let { if (it.isActive && !disconnected) return@withLock it }
+            // Once the receive loop has reported a disconnect we don't trust the old transport's
+            // `isActive` (it can lag the real socket close), so dispose it and build a fresh one. This
+            // makes a command issued after a drop reconnect deterministically instead of reusing a
+            // dead socket and failing (or hanging until the command timeout).
+            if (disconnected) {
+                transport?.let { old -> runCatching { old.close() } }
+                transport = null
+                disconnected = false
+                // A reconnect opens a *new* CDP session, so the browser-side preparation applied on
+                // the old one is gone: re-run it on the next command. NOTE: domains the caller enabled
+                // themselves (Network.enable, Page.enable, event subscriptions, …) are NOT restored —
+                // restoring arbitrary caller state on reconnect is out of scope here.
+                prepareHeadlessDone = false
+                prepareExpertDone = false
+            }
+            val t = transport ?: createTransport().also { transport = it }
+            if (!t.isActive) {
+                t.connect()
+                startListening(t)
+            }
+            t
         }
     }
 
-    private fun startListening() {
+    private fun startListening(t: WebSocketTransport) {
         socketSubscription?.cancel()
         socketSubscription = messageListeningScope.launch {
             try {
-                transport.incoming().collect { text ->
+                t.incoming().collect { text ->
                     try {
                         logger.debug("WS < CDP: ${text.take(owner?.config?.debugStringLimit ?: Defaults.DEBUG_STRING_LIMIT)}")
                         val received = Serialization.json.decodeFromString<Message>(text)
@@ -99,16 +141,27 @@
                 }
                 // incoming() completed without error => the socket was closed. Fail any in-flight
                 // commands so their callers observe the disconnect instead of hanging (ISSUE-3).
-                failPendingRequests(ConnectionClosedException())
+                onTransportGone(t, ConnectionClosedException())
             } catch (e: CancellationException) {
                 throw e
             } catch (e: Exception) {
                 logger.error("WebSocket receive loop terminated: {}", e)
-                failPendingRequests(ConnectionClosedException(cause = e))
+                onTransportGone(t, ConnectionClosedException(cause = e))
             }
         }
     }
 
+    /**
+     * Reacts to a receive loop ending because its socket went away: marks the connection for a
+     * reconnect and fails any in-flight commands. Ignored if [t] is no longer the current transport
+     * (a stale loop from a transport we've already replaced must not disturb the new one).
+     */
+    private suspend fun onTransportGone(t: WebSocketTransport, cause: Throwable) {
+        if (transport !== t) return
+        disconnected = true
+        failPendingRequests(cause)
+    }
+
     /**
      * Completes every in-flight request waiter exceptionally and clears the registry, so callers
      * parked in [callCommand] observe a failure rather than hanging when the connection goes away.
@@ -136,35 +189,79 @@
         }
 
         val requestId = currentIdMutex.withLock { currentId++ }
-        // Register the response waiter *before* sending, so a reply that arrives before we start
-        // awaiting is still captured (the receive loop completes this deferred). Awaiting the
-        // response via a replay-0 shared flow after sending could miss it and hang (ISSUE-1).
-        val deferred = CompletableDeferred<Message.Response>()
-        pendingRequestsMutex.withLock { pendingRequests[requestId] = deferred }
-        try {
-            val jsonString = Serialization.json.encodeToString(Request(requestId, method, parameter))
-            transport.send(jsonString)
-            logger.debug("WS > CDP: ${jsonString.take(owner?.config?.debugStringLimit ?: Defaults.DEBUG_STRING_LIMIT)}")
+        val jsonString = Serialization.json.encodeToString(Request(requestId, method, parameter))
+        val timeout = owner?.config?.commandTimeout ?: Defaults.COMMAND_TIMEOUT
 
-            val timeout = owner?.config?.commandTimeout ?: Defaults.COMMAND_TIMEOUT
-            // A non-null Message.Response means success; null can only come from the timeout, so
-            // there's no ambiguity with a legitimate value. A timeout <= 0 waits indefinitely.
-            val result =
-                if (timeout > 0) withTimeoutOrNull(timeout.milliseconds) { deferred.await() }
-                    ?: throw CommandTimeoutException(method, requestId, timeout)
-                else deferred.await()
-            result.error?.throwAsException(method)
-            return result.result
-        } finally {
-            pendingRequestsMutex.withLock { pendingRequests.remove(requestId) }
+        // A dropped socket can still report itself active for a short window, so the send below can
+        // fail (or the next command can land on a dead transport). A failed send means the bytes
+        // never left the client, so reconnecting and resending is safe — no risk of executing the
+        // command twice. We only retry the *send*: once a reply is awaited we never resend, because
+        // the command may already have run (that case surfaces as a ConnectionClosedException).
+        var sendError: Throwable? = null
+        repeat(SEND_ATTEMPTS) {
+            val transport = connect()
+            // Register the response waiter *before* sending, so a reply that arrives before we start
+            // awaiting is still captured (the receive loop completes this deferred). Awaiting the
+            // response via a replay-0 shared flow after sending could miss it and hang (ISSUE-1).
+            val deferred = CompletableDeferred<Message.Response>()
+            pendingRequestsMutex.withLock { pendingRequests[requestId] = deferred }
+
+            val sent = try {
+                transport.send(jsonString)
+                true
+            } catch (e: CancellationException) {
+                // A send on a dead Ktor session surfaces as a channel CancellationException even
+                // though *this* coroutine isn't cancelled; only the latter must propagate.
+                if (!currentCoroutineContext().isActive) throw e
+                sendError = e
+                false
+            } catch (e: Exception) {
+                sendError = e
+                false
+            }
+            if (!sent) {
+                // Treat the failed send as this transport going away: mark it for reconnect and fail
+                // any *other* commands in flight on it (ISSUE-3), but only if it is still the current
+                // transport — a send that fails late on an already-replaced transport must not tear
+                // down the healthy one that superseded it. onTransportGone clears our own just-added
+                // waiter too; we resend on the next iteration.
+                onTransportGone(transport, ConnectionClosedException(cause = sendError))
+                pendingRequestsMutex.withLock { pendingRequests.remove(requestId) }
+                return@repeat // reconnect and resend on the next iteration
+            }
+
+            logger.debug("WS > CDP: ${jsonString.take(owner?.config?.debugStringLimit ?: Defaults.DEBUG_STRING_LIMIT)}")
+            try {
+                // A non-null Message.Response means success; null can only come from the timeout, so
+                // there's no ambiguity with a legitimate value. A timeout <= 0 waits indefinitely.
+                val result =
+                    if (timeout > 0) withTimeoutOrNull(timeout.milliseconds) { deferred.await() }
+                        ?: throw CommandTimeoutException(method, requestId, timeout)
+                    else deferred.await()
+                result.error?.throwAsException(method)
+                return result.result
+            } finally {
+                pendingRequestsMutex.withLock { pendingRequests.remove(requestId) }
+            }
         }
+        // Every send attempt failed: the connection is gone. Surface it as a ConnectionClosedException
+        // (never the raw send failure, which may be a CancellationException a caller would mistake for
+        // its own cancellation), keeping the original send error as the cause.
+        throw ConnectionClosedException(cause = sendError)
     }
 
     @InternalCdpApi
     override suspend fun close() {
-        transport.close()
-        socketSubscription?.cancel()
-        socketSubscription = null
+        // Take connectMutex so close is ordered against an in-progress (re)connect: without it, a
+        // concurrent connect() could reassign `transport`/`socketSubscription` right after we cleared
+        // them, resurrecting a live socket and receive loop that outlive close().
+        connectMutex.withLock {
+            closed = true
+            transport?.close()
+            transport = null
+            socketSubscription?.cancel()
+            socketSubscription = null
+        }
         // Fail any commands still awaiting a reply that will now never come (ISSUE-3).
         failPendingRequests(ConnectionClosedException("Connection closed"))
     }

core/src/commonMain/kotlin/dev/kdriver/core/connection/KtorWebSocketTransport.kt

@@ -5,9 +5,9 @@
 import io.ktor.client.plugins.websocket.*
 import io.ktor.http.*
 import io.ktor.websocket.*
-import kotlinx.coroutines.isActive
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.flow
+import kotlinx.coroutines.isActive
 
 /**
  * Default [WebSocketTransport] backed by a Ktor WebSocket session.
@@ -39,7 +39,10 @@
     }
 
     override suspend fun send(message: String) {
-        session?.send(message)
+        // Fail loudly instead of silently succeeding when there is no session: a false success would
+        // make the caller wait for a reply to a frame that was never sent.
+        val session = session ?: throw IllegalStateException("WebSocket session is not connected")
+        session.send(message)
     }
 
     override fun incoming(): Flow<String> = flow {

core/src/jvmTest/kotlin/dev/kdriver/core/connection/CommandTimeoutTest.kt

@@ -43,7 +43,7 @@ class CommandTimeoutTest {
     @Test
     fun callCommand_failsWithCommandTimeout_whenNoResponse() = runTest(UnconfinedTestDispatcher()) {
         val transport = SilentTransport()
-        val connection = object : DefaultConnection("ws://stub/devtools/page/stub", this) {
+        val connection = object : DefaultConnection("ws://stub/devtools/page/stub", this@runTest) {
             override fun createTransport(): WebSocketTransport = transport
         }
 

core/src/jvmTest/kotlin/dev/kdriver/core/connection/ReconnectionIntegrationTest.kt

@@ -0,0 +1,100 @@
+package dev.kdriver.core.connection
+
+import dev.kdriver.cdp.CommandMode
+import dev.kdriver.core.browser.createBrowser
+import dev.kdriver.core.exceptions.CommandTimeoutException
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.withTimeout
+import kotlinx.serialization.json.*
+import kotlin.test.*
+
+/**
+ * End-to-end reconnection behaviour against a real Chrome, with a [TcpProxy] in the middle so the
+ * test can cut or freeze the socket at will.
+ *
+ * A real browser is launched (the "host"), and the test opens its own [DefaultConnection] to one of
+ * its page targets *through the proxy* — the per-target WebSocket URL is built from host:port, so it
+ * always goes through the proxy regardless of what Chrome advertises for the browser endpoint.
+ */
+class ReconnectionIntegrationTest {
+
+    private class ProxiedConnection(
+        private val wsUrl: String,
+        scope: CoroutineScope,
+    ) : DefaultConnection(wsUrl, scope) {
+        override fun createTransport(): WebSocketTransport = KtorWebSocketTransport(wsUrl)
+    }
+
+    private fun evalParam(expression: String): JsonObject = buildJsonObject {
+        put("expression", expression)
+        put("returnByValue", true)
+    }
+
+    private suspend fun DefaultConnection.evaluate(expression: String): Int {
+        val result = callCommand("Runtime.evaluate", evalParam(expression), CommandMode.ONE_SHOT)
+        return result!!.jsonObject["result"]!!.jsonObject["value"]!!.jsonPrimitive.int
+    }
+
+    @Test
+    fun reconnects_afterHardCut_throughProxy() = runBlocking {
+        val host = createBrowser(this, headless = true, sandbox = false)
+        val proxy = TcpProxy("127.0.0.1", host.config.port!!, this).also { it.start() }
+        val targetId = assertNotNull(host.mainTab?.targetId, "host browser should have a page target")
+        val wsUrl = "ws://127.0.0.1:${proxy.port}/devtools/page/$targetId"
+        val conn = ProxiedConnection(wsUrl, this)
+
+        try {
+            // Sanity: the proxied connection works end-to-end.
+            assertEquals(2, conn.evaluate("1 + 1"))
+
+            // Hard cut: the socket is closed under the connection.
+            proxy.severAll()
+
+            // The next command must transparently reconnect through the proxy and succeed, without
+            // the caller having to do anything.
+            assertEquals(4, withTimeout(15_000) { conn.evaluate("2 + 2") })
+        } finally {
+            conn.close()
+            proxy.stop()
+            host.stop()
+        }
+    }
+
+    @Test
+    fun halfOpenSocket_isBoundedByCommandTimeout_andRecovers() = runBlocking {
+        // Short command timeout so a half-open socket fails fast instead of hanging for 30s.
+        val host = createBrowser(this) {
+            headless = true
+            sandbox = false
+            commandTimeout = 3_000
+        }
+        val proxy = TcpProxy("127.0.0.1", host.config.port!!, this).also { it.start() }
+        val targetId = assertNotNull(host.mainTab?.targetId, "host browser should have a page target")
+        val wsUrl = "ws://127.0.0.1:${proxy.port}/devtools/page/$targetId"
+        // owner = host so the connection picks up the short commandTimeout from its config.
+        val conn = ProxiedConnection(wsUrl, this).also { it.owner = host }
+
+        try {
+            assertEquals(2, conn.evaluate("1 + 1"))
+
+            // Half-open: sockets stay open but no application bytes flow, so the reply never comes
+            // and the socket never closes. This is the case keep-alive pings would be needed to
+            // detect; for now it is merely *bounded* by the command timeout rather than hanging.
+            proxy.freeze()
+            val start = System.currentTimeMillis()
+            val error = runCatching { conn.evaluate("2 + 2") }.exceptionOrNull()
+            val elapsed = System.currentTimeMillis() - start
+            assertIs<CommandTimeoutException>(error, "a half-open socket must fail with a command timeout, not hang")
+            assertTrue(elapsed in 3_000..15_000, "should fail around the 3s command timeout (was ${elapsed}ms)")
+
+            // Once traffic flows again the same session is healthy, so commands resume working.
+            proxy.unfreeze()
+            assertEquals(9, withTimeout(15_000) { conn.evaluate("3 + 3 + 3") })
+        } finally {
+            conn.close()
+            proxy.stop()
+            host.stop()
+        }
+    }
+}