fix: restore correctness after reconnect

core/src/commonMain/kotlin/dev/kdriver/core/browser/DefaultBrowser.kt

@@ -295,6 +295,11 @@ open class DefaultBrowser(
             }
 
             connection.target.setDiscoverTargets(discover = true)
+            // The event collectors above survive a reconnect (they read the connection's shared flow),
+            // but the new CDP session won't emit target events until discovery is re-enabled on it.
+            connection.registerReconnectRestore("targetDiscovery") {
+                connection.target.setDiscoverTargets(discover = true)
+            }
         }
 
         updateTargets()

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

@@ -25,6 +25,29 @@ interface Connection : BrowserTarget, CDP {
     @InternalCdpApi
     suspend fun close()
 
+    /**
+     * Registers a [restore] action to be re-run automatically after the connection transparently
+     * reconnects.
+     *
+     * A reconnect opens a fresh CDP session that has lost all session-scoped state (enabled domains,
+     * overrides such as the user agent, scripts injected with `addScriptToEvaluateOnNewDocument`, …).
+     * Features that establish such state register a restorer here — keyed by [key] so it can be
+     * replaced or removed — and remove it with [unregisterReconnectRestore] when they tear that state
+     * down. Each restorer runs at most once per reconnect, before the command that triggered the
+     * restore proceeds; a concurrent command may still race ahead of an in-progress restore. A
+     * restorer that throws is logged and skipped rather than failing the command.
+     *
+     * A restorer must only re-issue CDP commands; it must not itself call
+     * [registerReconnectRestore]/[unregisterReconnectRestore].
+     */
+    suspend fun registerReconnectRestore(key: Any, restore: suspend () -> Unit)
+
+    /**
+     * Removes a restorer previously registered with [registerReconnectRestore]. No-op if [key] is
+     * not registered.
+     */
+    suspend fun unregisterReconnectRestore(key: Any)
+
     /**
      * Updates the target information by fetching it from the CDP.
      *

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

@@ -77,6 +77,14 @@
     private var prepareHeadlessDone = false
     private var prepareExpertDone = false
 
+    // Session state to re-apply after a reconnect (a new CDP session loses enabled domains, overrides,
+    // injected scripts, …). Guarded by [reconnectRestoreMutex]; [needsRestore] flips true on reconnect.
+    private val reconnectRestoreMutex = Mutex()
+    private val reconnectRestores = LinkedHashMap<Any, suspend () -> Unit>()
+
+    @Volatile
+    private var needsRestore = false
+
     private val allMessages = MutableSharedFlow<Message>(extraBufferCapacity = Channel.UNLIMITED)
 
     @InternalCdpApi
@@ -101,15 +109,21 @@
             // 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() } }
+                // Clear `transport` and cancel the old receive loop *before* closing the old transport,
+                // so the loop ending on that close sees `transport !== old` and no-ops instead of
+                // spuriously re-flagging `disconnected` (which would cascade into extra reconnects).
+                val old = transport
                 transport = null
+                socketSubscription?.cancel()
+                socketSubscription = null
+                old?.let { runCatching { it.close() } }
                 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.
+                // A reconnect opens a *new* CDP session, so all session state applied on the old one
+                // is gone. Re-run the built-in preparation, and flag the registered restorers (enabled
+                // domains, overrides, …) to be replayed before the next command proceeds.
                 prepareHeadlessDone = false
                 prepareExpertDone = false
+                needsRestore = true
             }
             val t = transport ?: createTransport().also { transport = it }
             if (!t.isActive) {
@@ -175,19 +189,54 @@
         pending.forEach { it.completeExceptionally(cause) }
     }
 
-    @InternalCdpApi
-    override suspend fun callCommand(method: String, parameter: JsonElement?, mode: CommandMode): JsonElement? {
-        connect()
+    override suspend fun registerReconnectRestore(key: Any, restore: suspend () -> Unit) {
+        reconnectRestoreMutex.withLock { reconnectRestores[key] = restore }
+    }
+
+    override suspend fun unregisterReconnectRestore(key: Any) {
+        reconnectRestoreMutex.withLock { reconnectRestores.remove(key) }
+    }
+
+    /**
+     * Re-applies the registered session state once after a reconnect. Snapshots the restorers and
+     * clears [needsRestore] under the lock, then runs them *outside* it: a restorer only re-issues
+     * CDP commands (whose [callCommand] sees [needsRestore] already false and so won't recurse here),
+     * and running off-lock avoids a deadlock if one of those commands itself triggers a reconnect.
+     */
+    private suspend fun runReconnectRestores() {
+        val restores = reconnectRestoreMutex.withLock {
+            if (!needsRestore) return
+            needsRestore = false
+            reconnectRestores.values.toList()
+        }
+        for (restore in restores) {
+            try {
+                restore()
+            } catch (e: CancellationException) {
+                throw e
+            } catch (e: Exception) {
+                logger.error("Reconnect restore failed: {}", e)
+            }
+        }
+    }
 
-        if (mode == CommandMode.DEFAULT) owner?.let { browser ->
-            // Serialize preparation so concurrent first commands run it once, not N times (ISSUE-4).
-            // prepare* issue ONE_SHOT commands, which skip this block, so prepareMutex isn't re-entered.
+    /**
+     * Runs the built-in browser-side preparation for a DEFAULT command, on the current session.
+     * prepare* issue ONE_SHOT commands, which skip this step, so [prepareMutex] isn't re-entered.
+     * Serialized so concurrent first commands run it once, not N times (ISSUE-4).
+     */
+    private suspend fun ensurePrepared(mode: CommandMode) {
+        if (mode != CommandMode.DEFAULT) return
+        owner?.let { browser ->
             prepareMutex.withLock {
                 if (browser.config.expert) prepareExpert()
                 if (browser.config.headless) prepareHeadless()
             }
         }
+    }
 
+    @InternalCdpApi
+    override suspend fun callCommand(method: String, parameter: JsonElement?, mode: CommandMode): JsonElement? {
         val requestId = currentIdMutex.withLock { currentId++ }
         val jsonString = Serialization.json.encodeToString(Request(requestId, method, parameter))
         val timeout = owner?.config?.commandTimeout ?: Defaults.COMMAND_TIMEOUT
@@ -199,6 +248,20 @@
         // the command may already have run (that case surfaces as a ConnectionClosedException).
         var sendError: Throwable? = null
         repeat(SEND_ATTEMPTS) {
+            connect()
+            // Restore + prepare the (possibly freshly reconnected) session *before* sending on it, so
+            // the command always runs against a fully re-established session — this covers a reconnect
+            // that happens here in the retry loop, not just one detected before it.
+            //
+            // Restore MUST run before (and outside of) ensurePrepared: prepare* issue their own inner
+            // commands, which would otherwise trigger the restore while prepareMutex is held, and a
+            // restorer's DEFAULT command re-entering the non-reentrant prepareMutex would deadlock.
+            // Running restore first is also correct precedence-wise — a restorer's own DEFAULT command
+            // re-runs prepare internally, so e.g. a user user-agent override still lands last.
+            if (needsRestore) runReconnectRestores()
+            ensurePrepared(mode)
+            // Fetch the transport only now: prepare/restore issue their own commands, which may have
+            // reconnected, so a transport captured before them could be the stale, disposed one.
             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
@@ -215,7 +278,7 @@
                 if (!currentCoroutineContext().isActive) throw e
                 sendError = e
                 false
             } catch (e: Exception) {
                 sendError = e
                 false
             }

core/src/commonMain/kotlin/dev/kdriver/core/network/BaseFetchInterception.kt

@@ -19,6 +19,12 @@ open class BaseFetchInterception(
     private var responseDeferred = CompletableDeferred<Fetch.RequestPausedParameter>()
     private var job: Job? = null
 
+    private val requestPattern = Fetch.RequestPattern(
+        urlPattern = urlPattern,
+        resourceType = resourceType,
+        requestStage = requestStage,
+    )
+
     private val handler: suspend (Fetch.RequestPausedParameter) -> Unit = handler@{ event ->
         responseDeferred.complete(event)
         job?.cancel()
@@ -32,18 +38,14 @@ open class BaseFetchInterception(
         job = coroutineScope.launch(start = CoroutineStart.UNDISPATCHED) {
             tab.fetch.requestPaused.collect { handler(it) }
         }
-        tab.fetch.enable(
-            listOf(
-                Fetch.RequestPattern(
-                    urlPattern = urlPattern,
-                    resourceType = resourceType,
-                    requestStage = requestStage
-                )
-            )
-        )
+        tab.fetch.enable(listOf(requestPattern))
+        // The requestPaused collector survives a reconnect (it reads the connection's shared flow),
+        // but Fetch must be re-enabled on the new CDP session for the browser to keep pausing requests.
+        tab.registerReconnectRestore(this) { tab.fetch.enable(listOf(requestPattern)) }
     }
 
     private suspend fun teardown() {
+        tab.unregisterReconnectRestore(this)
         job?.cancel()
         tab.fetch.disable()
     }

core/src/commonMain/kotlin/dev/kdriver/core/network/BaseRequestExpectation.kt

@@ -63,9 +63,13 @@ open class BaseRequestExpectation(
             tab.network.loadingFinished.collect { loadingFinishedHandler(it) }
         }
         tab.network.enable()
+        // The event collectors survive a reconnect; re-enable Network on the new session so the
+        // browser keeps emitting the request/response events this expectation waits for.
+        tab.registerReconnectRestore(this) { tab.network.enable() }
     }
 
-    private fun teardown() {
+    private suspend fun teardown() {
+        tab.unregisterReconnectRestore(this)
         requestJob?.cancel()
         responseJob?.cancel()
         loadingFinishedJob?.cancel()

core/src/commonMain/kotlin/dev/kdriver/core/tab/DefaultTab.kt

@@ -118,6 +118,15 @@ open class DefaultTab(
             acceptLanguage = acceptLanguage,
             platform = platform
         )
+        // A reconnect opens a new CDP session that would drop this override, so re-apply it then.
+        // Keyed so a later setUserAgent replaces (rather than stacks) the restorer.
+        registerReconnectRestore("userAgent") {
+            network.setUserAgentOverride(
+                userAgent = ua,
+                acceptLanguage = acceptLanguage,
+                platform = platform
+            )
+        }
     }
 
     override suspend fun getWindow(): dev.kdriver.cdp.domain.Browser.GetWindowForTargetReturn {

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

@@ -7,6 +7,7 @@ import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.runBlocking
 import kotlinx.coroutines.withTimeout
 import kotlinx.serialization.json.*
+import java.util.concurrent.atomic.AtomicInteger
 import kotlin.test.*
 
 /**
@@ -36,6 +37,21 @@ class ReconnectionIntegrationTest {
         return result!!.jsonObject["result"]!!.jsonObject["value"]!!.jsonPrimitive.int
     }
 
+    /**
+     * Evaluates, retrying through a transient reconnect. The single command *in flight* when the
+     * socket is cut may itself be the casualty (its send lands on the dying socket, then the reply
+     * never comes) — by design we never resend an awaited command. The connection self-heals, so the
+     * next attempt reconnects and succeeds, exactly as kdriver's high-level retry loops rely on.
+     */
+    private suspend fun DefaultConnection.evaluateRecovering(expression: String, timeoutMs: Long = 15_000): Int =
+        withTimeout(timeoutMs) {
+            while (true) {
+                val result = runCatching { evaluate(expression) }
+                if (result.isSuccess) return@withTimeout result.getOrThrow()
+            }
+            error("unreachable")
+        }
+
     @Test
     fun reconnects_afterHardCut_throughProxy() = runBlocking {
         val host = createBrowser(this, headless = true, sandbox = false)
@@ -51,9 +67,38 @@ class ReconnectionIntegrationTest {
             // 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") })
+            // The connection reconnects through the proxy on its own and commands succeed again.
+            assertEquals(4, conn.evaluateRecovering("2 + 2"))
+        } finally {
+            conn.close()
+            proxy.stop()
+            host.stop()
+        }
+    }
+
+    @Test
+    fun runsRegisteredRestore_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)
+        val restoreCount = AtomicInteger(0)
+
+        try {
+            assertEquals(2, conn.evaluate("1 + 1"))
+
+            // A feature (e.g. FetchInterception) registers a restorer to re-establish its session
+            // state after a reconnect. Here it just counts its runs.
+            conn.registerReconnectRestore("probe") { restoreCount.incrementAndGet() }
+
+            proxy.severAll()
+
+            // The connection reconnects once; the restorer must run exactly once as part of that
+            // single reconnect, no matter how many commands drive the recovery.
+            assertEquals(4, conn.evaluateRecovering("2 + 2"))
+            assertEquals(9, conn.evaluateRecovering("3 + 3 + 3"))
+            assertEquals(1, restoreCount.get(), "the registered restorer must run once per reconnect")
         } finally {
             conn.close()
             proxy.stop()