package ai.openclaw.app.gateway

import kotlinx.coroutines.CancellationException
import kotlinx.coroutines.CompletableDeferred
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.Job
import kotlinx.coroutines.SupervisorJob
import kotlinx.coroutines.awaitCancellation
import kotlinx.coroutines.cancelAndJoin
import kotlinx.coroutines.delay
import kotlinx.coroutines.launch
import kotlinx.coroutines.runBlocking
import kotlinx.coroutines.withTimeout
import kotlinx.coroutines.withTimeoutOrNull
import kotlinx.serialization.json.Json
import kotlinx.serialization.json.JsonArray
import kotlinx.serialization.json.JsonObject
import kotlinx.serialization.json.jsonObject
import kotlinx.serialization.json.jsonPrimitive
import okhttp3.Response
import okhttp3.WebSocket
import okhttp3.WebSocketListener
import okhttp3.mockwebserver.Dispatcher
import okhttp3.mockwebserver.MockResponse
import okhttp3.mockwebserver.MockWebServer
import okhttp3.mockwebserver.RecordedRequest
import org.junit.Assert.assertEquals
import org.junit.Assert.assertFalse
import org.junit.Assert.assertNull
import org.junit.Assert.assertTrue
import org.junit.Test
import org.junit.runner.RunWith
import org.robolectric.RobolectricTestRunner
import org.robolectric.RuntimeEnvironment
import org.robolectric.annotation.Config
import java.util.concurrent.CopyOnWriteArrayList
import java.util.concurrent.CountDownLatch
import java.util.concurrent.TimeUnit
import java.util.concurrent.atomic.AtomicInteger
import java.util.concurrent.atomic.AtomicReference

private const val TEST_TIMEOUT_MS = 8_000L
private const val CONNECT_CHALLENGE_TS = 1_700_000_000_123L
private const val CONNECT_CHALLENGE_FRAME =
  """{"type":"event","event":"connect.challenge","payload":{"nonce":"android-test-nonce","ts":$CONNECT_CHALLENGE_TS}}"""

private class InMemoryDeviceAuthStore : DeviceAuthTokenStore {
  private val tokens = mutableMapOf<String, DeviceAuthEntry>()

  override fun loadEntry(
    gatewayId: String,
    deviceId: String,
    role: String,
  ): DeviceAuthEntry? = tokens["${gatewayId.trim()}|${deviceId.trim()}|${role.trim()}"]

  override fun saveToken(
    gatewayId: String,
    deviceId: String,
    role: String,
    token: String,
    scopes: List<String>,
  ) {
    tokens["${gatewayId.trim()}|${deviceId.trim()}|${role.trim()}"] =
      DeviceAuthEntry(
        token = token.trim(),
        role = role.trim(),
        scopes = scopes,
        updatedAtMs = System.currentTimeMillis(),
      )
  }

  override fun clearToken(
    gatewayId: String,
    deviceId: String,
    role: String,
  ) {
    tokens.remove("${gatewayId.trim()}|${deviceId.trim()}|${role.trim()}")
  }
}

private data class NodeHarness(
  val session: GatewaySession,
  val sessionJob: Job,
  val deviceAuthStore: InMemoryDeviceAuthStore,
)

private data class InvokeScenarioResult(
  val request: GatewaySession.InvokeRequest,
  val resultParams: JsonObject,
)

@RunWith(RobolectricTestRunner::class)
@Config(sdk = [34])
class GatewaySessionInvokeTest {
  @Test
  fun connect_usesGatewayChallengeTimestamp() =
    runBlocking {
      val json = testJson()
      val connected = CompletableDeferred<Unit>()
      val lastDisconnect = AtomicReference("")
      val server =
        startGatewayServer(json) { webSocket, id, method, frame ->
          if (method == "connect") {
            assertEquals(
              CONNECT_CHALLENGE_TS,
              frame["params"]
                ?.jsonObject
                ?.get("device")
                ?.jsonObject
                ?.get("signedAt")
                ?.jsonPrimitive
                ?.content
                ?.toLong(),
            )
            webSocket.send(connectResponseFrame(id))
          }
        }
      val harness =
        createNodeHarness(
          connected = connected,
          lastDisconnect = lastDisconnect,
        ) { GatewaySession.InvokeResult.ok("""{"handled":true}""") }

      try {
        connectNodeSession(harness.session, server.port)
        awaitConnectedOrThrow(connected, lastDisconnect, server)
      } finally {
        shutdownHarness(harness, server)
      }
    }

  @Test
  fun connect_rejectsChallengeWithoutTimestamp() =
    runBlocking {
      val json = testJson()
      val connected = CompletableDeferred<Unit>()
      val lastDisconnect = AtomicReference("")
      val connectRequests = AtomicInteger()
      val server =
        startGatewayServer(
          json = json,
          challengeFrame =
            """{"type":"event","event":"connect.challenge","payload":{"nonce":"android-test-nonce"}}""",
        ) { _, _, method, _ ->
          if (method == "connect") connectRequests.incrementAndGet()
        }
      val harness =
        createNodeHarness(
          connected = connected,
          lastDisconnect = lastDisconnect,
        ) { GatewaySession.InvokeResult.ok("""{"handled":true}""") }

      try {
        connectNodeSession(harness.session, server.port)
        withTimeout(TEST_TIMEOUT_MS) {
          while (lastDisconnect.get().isEmpty()) delay(10)
        }
        assertFalse(connected.isCompleted)
        assertEquals(0, connectRequests.get())
      } finally {
        shutdownHarness(harness, server)
      }
    }

  @Test
  fun canvasRoutePinsOnlyTheConnectedTlsEndpoint() {
    val fingerprint = "ab".repeat(32)
    val endpoint = GatewayEndpoint.manual(host = "gateway.example", port = 7443)

    assertEquals(
      fingerprint,
      gatewayTlsFingerprintForCanvasSurface(
        fingerprint = fingerprint,
        surfaceUrl = "https://gateway.example:7443/__openclaw__/cap/token",
        endpoint = endpoint,
        isTlsConnection = true,
      ),
    )
    assertNull(
      gatewayTlsFingerprintForCanvasSurface(
        fingerprint = fingerprint,
        surfaceUrl = "https://canvas.example:7443/__openclaw__/cap/token",
        endpoint = endpoint,
        isTlsConnection = true,
      ),
    )
    assertNull(
      gatewayTlsFingerprintForCanvasSurface(
        fingerprint = fingerprint,
        surfaceUrl = "https://gateway.example:9443/__openclaw__/cap/token",
        endpoint = endpoint,
        isTlsConnection = true,
      ),
    )
  }

  @Test
  fun refreshCanvasHostUrl_usesNodeRefreshMethod() =
    runBlocking {
      assertCanvasHostRefreshMethod(role = "node", expectedMethod = "node.pluginSurface.refresh")
    }

  @Test
  fun refreshCanvasHostUrl_usesOperatorRefreshMethod() =
    runBlocking {
      assertCanvasHostRefreshMethod(role = "operator", expectedMethod = "plugin.surface.refresh")
    }

  private suspend fun assertCanvasHostRefreshMethod(
    role: String,
    expectedMethod: String,
  ) {
    val json = testJson()
    val connected = CompletableDeferred<Unit>()
    val lastDisconnect = AtomicReference("")
    val refreshRequests = AtomicInteger()
    val server =
      startGatewayServer(json) { webSocket, id, method, frame ->
        when (method) {
          "connect" ->
            webSocket.send(
              connectResponseFrame(
                id,
                pluginSurfaceUrls =
                  mapOf("canvas" to "http://127.0.0.1:18789/__openclaw__/cap/old-token"),
              ),
            )
          expectedMethod -> {
            refreshRequests.incrementAndGet()
            assertEquals(
              "canvas",
              frame["params"]
                ?.jsonObject
                ?.get("surface")
                ?.jsonPrimitive
                ?.content,
            )
            assertTrue(
              frame["params"]
                ?.jsonObject
                ?.get("observedUrl")
                ?.jsonPrimitive
                ?.content
                ?.endsWith("/old-token") == true,
            )
            webSocket.send(
              """{"type":"res","id":"$id","ok":true,"payload":{"surface":"canvas","pluginSurfaceUrls":{"canvas":"http://127.0.0.1:18789/__openclaw__/cap/new-token"}}}""",
            )
          }
        }
      }
    val harness =
      createNodeHarness(connected = connected, lastDisconnect = lastDisconnect) {
        GatewaySession.InvokeResult.ok("""{"handled":true}""")
      }

    try {
      connectNodeSession(
        session = harness.session,
        port = server.port,
        role = role,
        scopes = if (role == "operator") listOf("operator.read") else listOf("node:invoke"),
      )
      awaitConnectedOrThrow(connected, lastDisconnect, server)
      val oldUrl = requireNotNull(harness.session.currentCanvasHostUrl())
      assertTrue(oldUrl.endsWith("/old-token"))

      val refreshed = harness.session.refreshCanvasHostUrlIfCurrent(oldUrl)
      val lagging = harness.session.refreshCanvasHostUrlIfCurrent(oldUrl)

      assertTrue(refreshed?.endsWith("/new-token") == true)
      assertEquals(refreshed, harness.session.currentCanvasHostUrl())
      assertEquals(refreshed, lagging)
      assertEquals(1, refreshRequests.get())
    } finally {
      shutdownHarness(harness, server)
    }
  }

  @Test
  fun connect_advertisesCompatibleProtocolRange() =
    runBlocking {
      val json = testJson()
      val connected = CompletableDeferred<Unit>()
      val connectParams = CompletableDeferred<JsonObject>()
      val lastDisconnect = AtomicReference("")
      val server =
        startGatewayServer(json) { webSocket, id, method, frame ->
          when (method) {
            "connect" -> {
              if (!connectParams.isCompleted) {
                connectParams.complete(frame["params"]!!.jsonObject)
              }
              webSocket.send(connectResponseFrame(id))
              webSocket.close(1000, "done")
            }
          }
        }

      val harness =
        createNodeHarness(
          connected = connected,
          lastDisconnect = lastDisconnect,
        ) { GatewaySession.InvokeResult.ok("""{"handled":true}""") }

      try {
        connectNodeSession(harness.session, server.port)
        awaitConnectedOrThrow(connected, lastDisconnect, server)

        val params = withTimeout(TEST_TIMEOUT_MS) { connectParams.await() }
        assertEquals(
          GATEWAY_MIN_PROTOCOL_VERSION,
          params["minProtocol"]?.jsonPrimitive?.content?.toInt(),
        )
        assertEquals(
          GATEWAY_PROTOCOL_VERSION,
          params["maxProtocol"]?.jsonPrimitive?.content?.toInt(),
        )
      } finally {
        shutdownHarness(harness, server)
      }
    }

  @Test
  fun disconnectFailsPendingRpcWithUnknownOutcomeWithoutWaitingForTimeout() {
    runBlocking {
      val json = testJson()
      val connected = CompletableDeferred<Unit>()
      val slowRequestSeen = CompletableDeferred<Unit>()
      val requestResult = CompletableDeferred<Result<GatewaySession.RpcResult>>()
      val lastDisconnect = AtomicReference("")
      val serverWebSocket = AtomicReference<WebSocket?>(null)
      val server =
        startGatewayServer(json) { webSocket, id, method, _ ->
          serverWebSocket.set(webSocket)
          when (method) {
            "connect" -> webSocket.send(connectResponseFrame(id))
            "slow.method" -> {
              if (!slowRequestSeen.isCompleted) slowRequestSeen.complete(Unit)
            }
          }
        }

      val harness =
        createNodeHarness(
          connected = connected,
          lastDisconnect = lastDisconnect,
        ) { GatewaySession.InvokeResult.ok("""{"handled":true}""") }
      var requestJob: Job? = null

      try {
        connectNodeSession(harness.session, server.port)
        awaitConnectedOrThrow(connected, lastDisconnect, server)
        requestJob =
          launch {
            requestResult.complete(
              runCatching {
                harness.session.requestDetailed("slow.method", null, timeoutMs = 30_000)
              },
            )
          }
        withTimeout(TEST_TIMEOUT_MS) { slowRequestSeen.await() }

        harness.session.disconnect()

        val result = withTimeout(2_000) { requestResult.await() }
        assertEquals(true, result.exceptionOrNull() is GatewayRequestOutcomeUnknown)
        serverWebSocket.get()?.close(1000, "done")
        withTimeoutOrNull(2_000) {
          while (lastDisconnect.get().isEmpty()) delay(10)
        }
      } finally {
        requestJob?.cancelAndJoin()
        runCatching { serverWebSocket.get()?.close(1000, "done") }
        delay(100)
        harness.session.disconnect()
        harness.sessionJob.cancelAndJoin()
        server.shutdown()
      }
    }
  }

  @Test
  fun disconnectReportsUnknownOutcomeForFireAndForgetRpc() {
    runBlocking {
      val json = testJson()
      val connected = CompletableDeferred<Unit>()
      val requestSeen = CompletableDeferred<Unit>()
      val requestError = CompletableDeferred<GatewaySession.ErrorShape>()
      val lastDisconnect = AtomicReference("")
      val serverWebSocket = AtomicReference<WebSocket?>(null)
      val server =
        startGatewayServer(json) { webSocket, id, method, _ ->
          serverWebSocket.set(webSocket)
          when (method) {
            "connect" -> webSocket.send(connectResponseFrame(id))
            "fire.and.forget" -> requestSeen.complete(Unit)
          }
        }
      val harness =
        createNodeHarness(
          connected = connected,
          lastDisconnect = lastDisconnect,
        ) { GatewaySession.InvokeResult.ok("""{"handled":true}""") }

      try {
        connectNodeSession(harness.session, server.port)
        awaitConnectedOrThrow(connected, lastDisconnect, server)
        harness.session.sendRequestFrame(
          method = "fire.and.forget",
          paramsJson = null,
          timeoutMs = 30_000,
          onError = { requestError.complete(it) },
        )
        withTimeout(TEST_TIMEOUT_MS) { requestSeen.await() }

        harness.session.disconnect()

        val error = withTimeout(2_000) { requestError.await() }
        assertEquals("UNAVAILABLE", error.code)
        assertEquals("Gateway disconnected before response", error.message)
        serverWebSocket.get()?.close(1000, "done")
      } finally {
        runCatching { serverWebSocket.get()?.close(1000, "done") }
        delay(100)
        harness.session.disconnect()
        harness.sessionJob.cancelAndJoin()
        server.shutdown()
      }
    }
  }

  @Test
  fun eventsAreDispatchedInWebSocketFrameOrder() =
    runBlocking {
      val json = testJson()
      val connected = CompletableDeferred<Unit>()
      val firstEventStarted = CompletableDeferred<Unit>()
      val releaseFirstEvent = CompletableDeferred<Unit>()
      val secondEventHandled = CompletableDeferred<Unit>()
      val events = CopyOnWriteArrayList<String>()
      val lastDisconnect = AtomicReference("")
      val serverWebSocket = AtomicReference<WebSocket?>(null)
      val server =
        startGatewayServer(json) { webSocket, id, method, _ ->
          serverWebSocket.set(webSocket)
          if (method == "connect") {
            webSocket.send(connectResponseFrame(id))
            webSocket.send("""{"type":"event","event":"voice.first","payload":{}}""")
            webSocket.send("""{"type":"event","event":"voice.second","payload":{}}""")
          }
        }

      val harness =
        createNodeHarness(
          connected = connected,
          lastDisconnect = lastDisconnect,
          onEvent = { event, _ ->
            if (event == "voice.first") {
              firstEventStarted.complete(Unit)
              runBlocking { releaseFirstEvent.await() }
            }
            events += event
            if (event == "voice.second") {
              secondEventHandled.complete(Unit)
            }
          },
        ) { GatewaySession.InvokeResult.ok("""{"handled":true}""") }

      try {
        connectNodeSession(harness.session, server.port)
        awaitConnectedOrThrow(connected, lastDisconnect, server)
        withTimeout(TEST_TIMEOUT_MS) { firstEventStarted.await() }

        assertNull(withTimeoutOrNull(200) { secondEventHandled.await() })

        releaseFirstEvent.complete(Unit)
        withTimeout(TEST_TIMEOUT_MS) { secondEventHandled.await() }
        assertEquals(listOf("voice.first", "voice.second"), events.toList())
      } finally {
        releaseFirstEvent.complete(Unit)
        runCatching { serverWebSocket.get()?.close(1000, "done") }
        delay(100)
        shutdownHarness(harness, server)
      }
    }

  @Test
  fun explicitNullPayloadsRemainPresentForResponsesAndEvents() =
    runBlocking {
      val json = testJson()
      val connected = CompletableDeferred<Unit>()
      val eventPayload = CompletableDeferred<String?>()
      val lastDisconnect = AtomicReference("")
      val server =
        startGatewayServer(json) { webSocket, id, method, _ ->
          when (method) {
            "connect" -> {
              webSocket.send(connectResponseFrame(id))
              webSocket.send("""{"type":"event","event":"health","payload":null}""")
            }
            "test.null-payload" ->
              webSocket.send("""{"type":"res","id":"$id","ok":true,"payload":null}""")
          }
        }
      val harness =
        createNodeHarness(
          connected = connected,
          lastDisconnect = lastDisconnect,
          onEvent = { event, payload ->
            if (event == GatewayEvent.Health.rawValue) eventPayload.complete(payload)
          },
        ) { GatewaySession.InvokeResult.ok("""{"handled":true}""") }

      try {
        connectNodeSession(harness.session, server.port)
        awaitConnectedOrThrow(connected, lastDisconnect, server)

        val response = harness.session.requestDetailed("test.null-payload", null)

        assertEquals("null", response.payloadJson)
        assertEquals("null", withTimeout(TEST_TIMEOUT_MS) { eventPayload.await() })
      } finally {
        shutdownHarness(harness, server)
      }
    }

  @Test
  fun connect_usesBootstrapTokenWhenSharedAndDeviceTokensAreAbsent() =
    runBlocking {
      val json = testJson()
      val connected = CompletableDeferred<Unit>()
      val connectAuth = CompletableDeferred<JsonObject?>()
      val lastDisconnect = AtomicReference("")
      val server =
        startGatewayServer(json) { webSocket, id, method, frame ->
          when (method) {
            "connect" -> {
              if (!connectAuth.isCompleted) {
                connectAuth.complete(frame["params"]?.jsonObject?.get("auth")?.jsonObject)
              }
              webSocket.send(connectResponseFrame(id))
              webSocket.close(1000, "done")
            }
          }
        }

      val harness =
        createNodeHarness(
          connected = connected,
          lastDisconnect = lastDisconnect,
        ) { GatewaySession.InvokeResult.ok("""{"handled":true}""") }

      try {
        connectNodeSession(
          session = harness.session,
          port = server.port,
          token = null,
          bootstrapToken = "bootstrap-token",
        )
        awaitConnectedOrThrow(connected, lastDisconnect, server)

        val auth = withTimeout(TEST_TIMEOUT_MS) { connectAuth.await() }
        assertEquals("bootstrap-token", auth?.get("bootstrapToken")?.jsonPrimitive?.content)
        assertNull(auth?.get("token"))
      } finally {
        shutdownHarness(harness, server)
      }
    }

  @Test
  fun connect_prefersStoredDeviceTokenOverBootstrapToken() =
    runBlocking {
      val json = testJson()
      val connected = CompletableDeferred<Unit>()
      val connectAuth = CompletableDeferred<JsonObject?>()
      val lastDisconnect = AtomicReference("")
      val server =
        startGatewayServer(json) { webSocket, id, method, frame ->
          when (method) {
            "connect" -> {
              if (!connectAuth.isCompleted) {
                connectAuth.complete(frame["params"]?.jsonObject?.get("auth")?.jsonObject)
              }
              webSocket.send(connectResponseFrame(id))
              webSocket.close(1000, "done")
            }
          }
        }

      val harness =
        createNodeHarness(
          connected = connected,
          lastDisconnect = lastDisconnect,
        ) { GatewaySession.InvokeResult.ok("""{"handled":true}""") }

      try {
        val deviceId = testDeviceIdentityStore(RuntimeEnvironment.getApplication()).loadOrCreate().deviceId
        harness.deviceAuthStore.saveToken(gatewayIdForPort(server.port), deviceId, "node", "device-token")

        connectNodeSession(
          session = harness.session,
          port = server.port,
          token = null,
          bootstrapToken = "bootstrap-token",
        )
        awaitConnectedOrThrow(connected, lastDisconnect, server)

        val auth = withTimeout(TEST_TIMEOUT_MS) { connectAuth.await() }
        assertEquals("device-token", auth?.get("token")?.jsonPrimitive?.content)
        assertNull(auth?.get("bootstrapToken"))
      } finally {
        shutdownHarness(harness, server)
      }
    }

  @Test
  fun connect_reusesStoredDeviceTokenScopes() =
    runBlocking {
      val json = testJson()
      val connected = CompletableDeferred<Unit>()
      val connectParams = CompletableDeferred<JsonObject>()
      val lastDisconnect = AtomicReference("")
      val server =
        startGatewayServer(json) { webSocket, id, method, frame ->
          if (method == "connect") {
            if (!connectParams.isCompleted) {
              connectParams.complete(frame["params"]!!.jsonObject)
            }
            webSocket.send(connectResponseFrame(id))
            webSocket.close(1000, "done")
          }
        }

      val harness =
        createNodeHarness(
          connected = connected,
          lastDisconnect = lastDisconnect,
        ) { GatewaySession.InvokeResult.ok("""{"handled":true}""") }

      try {
        val deviceId = testDeviceIdentityStore(RuntimeEnvironment.getApplication()).loadOrCreate().deviceId
        harness.deviceAuthStore.saveToken(
          gatewayId = gatewayIdForPort(server.port),
          deviceId = deviceId,
          role = "operator",
          token = "operator-device-token",
          scopes = listOf("operator.pairing", "operator.write"),
        )

        connectNodeSession(
          session = harness.session,
          port = server.port,
          token = null,
          role = "operator",
          scopes = listOf("operator.approvals", "operator.read", "operator.write"),
        )
        awaitConnectedOrThrow(connected, lastDisconnect, server)

        val params = withTimeout(TEST_TIMEOUT_MS) { connectParams.await() }
        assertEquals(
          "operator-device-token",
          params["auth"]
            ?.jsonObject
            ?.get("token")
            ?.jsonPrimitive
            ?.content,
        )
        assertEquals(listOf("operator.pairing", "operator.write"), params.scopes())
      } finally {
        shutdownHarness(harness, server)
      }
    }

  @Test
  fun bootstrapConnect_requestsCanonicalLimitedOperatorHandoffScopes() =
    runBlocking {
      val json = testJson()
      val connected = CompletableDeferred<Unit>()
      val connectParams = CompletableDeferred<JsonObject>()
      val lastDisconnect = AtomicReference("")
      val server =
        startGatewayServer(json) { webSocket, id, method, frame ->
          if (method == "connect") {
            if (!connectParams.isCompleted) {
              connectParams.complete(frame["params"]!!.jsonObject)
            }
            webSocket.send(connectResponseFrame(id))
            webSocket.close(1000, "done")
          }
        }

      val harness =
        createNodeHarness(
          connected = connected,
          lastDisconnect = lastDisconnect,
        ) { GatewaySession.InvokeResult.ok("""{"handled":true}""") }

      try {
        connectNodeSession(
          session = harness.session,
          port = server.port,
          token = null,
          bootstrapToken = "setup-bootstrap-token",
          role = "operator",
          scopes =
            listOf(
              "operator.approvals",
              "operator.pairing",
              "operator.questions",
              "operator.read",
              "operator.talk.secrets",
              "operator.write",
            ),
        )
        awaitConnectedOrThrow(connected, lastDisconnect, server)

        val params = withTimeout(TEST_TIMEOUT_MS) { connectParams.await() }
        assertEquals(
          "setup-bootstrap-token",
          params["auth"]
            ?.jsonObject
            ?.get("bootstrapToken")
            ?.jsonPrimitive
            ?.content,
        )
        assertEquals(
          listOf(
            "operator.approvals",
            "operator.questions",
            "operator.read",
            "operator.talk.secrets",
            "operator.write",
          ),
          params.scopes(),
        )
      } finally {
        shutdownHarness(harness, server)
      }
    }

  @Test
  fun connect_retriesWithStoredDeviceTokenAfterSharedTokenMismatch() =
    runBlocking {
      val json = testJson()
      val connected = CompletableDeferred<Unit>()
      val firstConnectAuth = CompletableDeferred<JsonObject?>()
      val secondConnectAuth = CompletableDeferred<JsonObject?>()
      val connectAttempts = AtomicInteger(0)
      val lastDisconnect = AtomicReference("")
      val server =
        startGatewayServer(json) { webSocket, id, method, frame ->
          when (method) {
            "connect" -> {
              val auth = frame["params"]?.jsonObject?.get("auth")?.jsonObject
              when (connectAttempts.incrementAndGet()) {
                1 -> {
                  if (!firstConnectAuth.isCompleted) {
                    firstConnectAuth.complete(auth)
                  }
                  webSocket.send(
                    """{"type":"res","id":"$id","ok":false,"error":{"code":"INVALID_REQUEST","message":"unauthorized","details":{"code":"AUTH_TOKEN_MISMATCH","canRetryWithDeviceToken":true,"recommendedNextStep":"retry_with_device_token"}}}""",
                  )
                  webSocket.close(1000, "retry")
                }
                else -> {
                  if (!secondConnectAuth.isCompleted) {
                    secondConnectAuth.complete(auth)
                  }
                  webSocket.send(connectResponseFrame(id))
                }
              }
            }
          }
        }

      val harness =
        createNodeHarness(
          connected = connected,
          lastDisconnect = lastDisconnect,
        ) { GatewaySession.InvokeResult.ok("""{"handled":true}""") }

      try {
        val deviceId = testDeviceIdentityStore(RuntimeEnvironment.getApplication()).loadOrCreate().deviceId
        harness.deviceAuthStore.saveToken(gatewayIdForPort(server.port), deviceId, "node", "stored-device-token")

        connectNodeSession(
          session = harness.session,
          port = server.port,
          token = "shared-auth-token",
          bootstrapToken = null,
        )
        awaitConnectedOrThrow(connected, lastDisconnect, server)

        val firstAuth = withTimeout(TEST_TIMEOUT_MS) { firstConnectAuth.await() }
        val secondAuth = withTimeout(TEST_TIMEOUT_MS) { secondConnectAuth.await() }
        assertEquals("shared-auth-token", firstAuth?.get("token")?.jsonPrimitive?.content)
        assertNull(firstAuth?.get("deviceToken"))
        assertEquals("shared-auth-token", secondAuth?.get("token")?.jsonPrimitive?.content)
        assertEquals("stored-device-token", secondAuth?.get("deviceToken")?.jsonPrimitive?.content)
      } finally {
        shutdownHarness(harness, server)
      }
    }

  @Test
  fun connect_storesPrimaryDeviceTokenFromSuccessfulSharedTokenConnect() =
    runBlocking {
      val json = testJson()
      val connected = CompletableDeferred<Unit>()
      val lastDisconnect = AtomicReference("")
      val server =
        startGatewayServer(json) { webSocket, id, method, _ ->
          when (method) {
            "connect" -> {
              webSocket.send(
                connectResponseFrame(
                  id,
                  authJson = """{"deviceToken":"shared-node-token","role":"node","scopes":[]}""",
                ),
              )
              webSocket.close(1000, "done")
            }
          }
        }

      val harness =
        createNodeHarness(
          connected = connected,
          lastDisconnect = lastDisconnect,
        ) { GatewaySession.InvokeResult.ok("""{"handled":true}""") }

      try {
        connectNodeSession(
          session = harness.session,
          port = server.port,
          token = "shared-auth-token",
          bootstrapToken = null,
        )
        awaitConnectedOrThrow(connected, lastDisconnect, server)

        val deviceId = testDeviceIdentityStore(RuntimeEnvironment.getApplication()).loadOrCreate().deviceId
        assertEquals("shared-node-token", harness.deviceAuthStore.loadToken(gatewayIdForPort(server.port), deviceId, "node"))
        assertNull(harness.deviceAuthStore.loadToken(gatewayIdForPort(server.port), deviceId, "operator"))
      } finally {
        shutdownHarness(harness, server)
      }
    }

  @Test
  fun bootstrapConnect_storesAdditionalBoundedDeviceTokensOnTrustedTransport() =
    runBlocking {
      val json = testJson()
      val connected = CompletableDeferred<Unit>()
      val lastDisconnect = AtomicReference("")
      val server =
        startGatewayServer(json) { webSocket, id, method, _ ->
          when (method) {
            "connect" -> {
              webSocket.send(
                connectResponseFrame(
                  id,
                  authJson =
                    """{"deviceToken":"bootstrap-node-token","role":"node","scopes":[],"deviceTokens":[{"deviceToken":"bootstrap-operator-token","role":"operator","scopes":["operator.admin","operator.approvals","operator.pairing","operator.read","operator.talk.secrets","operator.write"]}]}""",
                ),
              )
              webSocket.close(1000, "done")
            }
          }
        }

      val harness =
        createNodeHarness(
          connected = connected,
          lastDisconnect = lastDisconnect,
        ) { GatewaySession.InvokeResult.ok("""{"handled":true}""") }

      try {
        connectNodeSession(
          session = harness.session,
          port = server.port,
          token = null,
          bootstrapToken = "bootstrap-token",
        )
        awaitConnectedOrThrow(connected, lastDisconnect, server)

        val deviceId = testDeviceIdentityStore(RuntimeEnvironment.getApplication()).loadOrCreate().deviceId
        val nodeEntry = harness.deviceAuthStore.loadEntry(gatewayIdForPort(server.port), deviceId, "node")
        val operatorEntry = harness.deviceAuthStore.loadEntry(gatewayIdForPort(server.port), deviceId, "operator")
        assertEquals("bootstrap-node-token", nodeEntry?.token)
        assertEquals(emptyList<String>(), nodeEntry?.scopes)
        assertEquals("bootstrap-operator-token", operatorEntry?.token)
        assertEquals(
          listOf(
            "operator.admin",
            "operator.approvals",
            "operator.read",
            "operator.talk.secrets",
            "operator.write",
          ),
          operatorEntry?.scopes,
        )
      } finally {
        shutdownHarness(harness, server)
      }
    }

  @Test
  fun nonBootstrapConnect_ignoresAdditionalBootstrapDeviceTokens() =
    runBlocking {
      val json = testJson()
      val connected = CompletableDeferred<Unit>()
      val lastDisconnect = AtomicReference("")
      val server =
        startGatewayServer(json) { webSocket, id, method, _ ->
          when (method) {
            "connect" -> {
              webSocket.send(
                connectResponseFrame(
                  id,
                  authJson =
                    """{"deviceToken":"shared-node-token","role":"node","scopes":[],"deviceTokens":[{"deviceToken":"shared-operator-token","role":"operator","scopes":["operator.approvals","operator.read"]}]}""",
                ),
              )
              webSocket.close(1000, "done")
            }
          }
        }

      val harness =
        createNodeHarness(
          connected = connected,
          lastDisconnect = lastDisconnect,
        ) { GatewaySession.InvokeResult.ok("""{"handled":true}""") }

      try {
        connectNodeSession(
          session = harness.session,
          port = server.port,
          token = "shared-auth-token",
          bootstrapToken = null,
        )
        awaitConnectedOrThrow(connected, lastDisconnect, server)

        val deviceId = testDeviceIdentityStore(RuntimeEnvironment.getApplication()).loadOrCreate().deviceId
        assertEquals("shared-node-token", harness.deviceAuthStore.loadToken(gatewayIdForPort(server.port), deviceId, "node"))
        assertNull(harness.deviceAuthStore.loadToken(gatewayIdForPort(server.port), deviceId, "operator"))
      } finally {
        shutdownHarness(harness, server)
      }
    }

  @Test
  fun nodeInvokeRequest_roundTripsInvokeResult() =
    runBlocking {
      val handshakeOrigin = AtomicReference<String?>(null)
      val result =
        runInvokeScenario(
          invokeEventFrame =
            """{"type":"event","event":"node.invoke.request","payload":{"id":"invoke-1","nodeId":"node-1","command":"debug.ping","params":{"ping":"pong"},"timeoutMs":5000}}""",
          onHandshake = { request -> handshakeOrigin.compareAndSet(null, request.getHeader("Origin")) },
        ) {
          GatewaySession.InvokeResult.ok("""{"handled":true}""")
        }

      assertEquals("invoke-1", result.request.id)
      assertEquals("node-1", result.request.nodeId)
      assertEquals("debug.ping", result.request.command)
      assertEquals("""{"ping":"pong"}""", result.request.paramsJson)
      assertNull(handshakeOrigin.get())
      assertEquals("invoke-1", result.resultParams["id"]?.jsonPrimitive?.content)
      assertEquals("node-1", result.resultParams["nodeId"]?.jsonPrimitive?.content)
      assertEquals(
        true,
        result.resultParams["ok"]
          ?.jsonPrimitive
          ?.content
          ?.toBooleanStrict(),
      )
      assertEquals(
        true,
        result.resultParams["payload"]
          ?.jsonObject
          ?.get("handled")
          ?.jsonPrimitive
          ?.content
          ?.toBooleanStrict(),
      )
    }

  @Test
  fun nodeInvokeRequest_usesParamsJsonWhenProvided() =
    runBlocking {
      val result =
        runInvokeScenario(
          invokeEventFrame =
            """{"type":"event","event":"node.invoke.request","payload":{"id":"invoke-2","nodeId":"node-2","command":"debug.raw","paramsJSON":"{\"raw\":true}","params":{"ignored":1},"timeoutMs":5000}}""",
        ) {
          GatewaySession.InvokeResult.ok("""{"handled":true}""")
        }

      assertEquals("invoke-2", result.request.id)
      assertEquals("node-2", result.request.nodeId)
      assertEquals("debug.raw", result.request.command)
      assertEquals("""{"raw":true}""", result.request.paramsJson)
      assertEquals("invoke-2", result.resultParams["id"]?.jsonPrimitive?.content)
      assertEquals("node-2", result.resultParams["nodeId"]?.jsonPrimitive?.content)
      assertEquals(
        true,
        result.resultParams["ok"]
          ?.jsonPrimitive
          ?.content
          ?.toBooleanStrict(),
      )
    }

  @Test
  fun nodeInvokeRequest_mapsCodePrefixedErrorsIntoInvokeResult() =
    runBlocking {
      val result =
        runInvokeScenario(
          invokeEventFrame =
            """{"type":"event","event":"node.invoke.request","payload":{"id":"invoke-3","nodeId":"node-3","command":"camera.snap","params":{"facing":"front"},"timeoutMs":5000}}""",
        ) {
          throw IllegalStateException("CAMERA_PERMISSION_REQUIRED: grant Camera permission")
        }

      assertEquals("invoke-3", result.resultParams["id"]?.jsonPrimitive?.content)
      assertEquals("node-3", result.resultParams["nodeId"]?.jsonPrimitive?.content)
      assertEquals(
        false,
        result.resultParams["ok"]
          ?.jsonPrimitive
          ?.content
          ?.toBooleanStrict(),
      )
      assertEquals(
        "CAMERA_PERMISSION_REQUIRED",
        result.resultParams["error"]
          ?.jsonObject
          ?.get("code")
          ?.jsonPrimitive
          ?.content,
      )
      assertEquals(
        "grant Camera permission",
        result.resultParams["error"]
          ?.jsonObject
          ?.get("message")
          ?.jsonPrimitive
          ?.content,
      )
    }

  @Test
  fun nodeInvokeRequest_cancelsHandlerWhenExecutionTimeoutExpires() =
    runBlocking {
      val handlerCancelled = CompletableDeferred<Unit>()
      val result =
        runInvokeScenario(
          invokeEventFrame =
            """{"type":"event","event":"node.invoke.request","payload":{"id":"invoke-timeout","nodeId":"node-1","command":"camera.clip","timeoutMs":100}}""",
        ) {
          try {
            awaitCancellation()
          } finally {
            handlerCancelled.complete(Unit)
          }
        }

      withTimeout(TEST_TIMEOUT_MS) { handlerCancelled.await() }
      assertEquals(
        false,
        result.resultParams["ok"]
          ?.jsonPrimitive
          ?.content
          ?.toBooleanStrict(),
      )
      assertEquals(
        "TIMEOUT",
        result.resultParams["error"]
          ?.jsonObject
          ?.get("code")
          ?.jsonPrimitive
          ?.content,
      )
      assertEquals(
        "node invoke timed out",
        result.resultParams["error"]
          ?.jsonObject
          ?.get("message")
          ?.jsonPrimitive
          ?.content,
      )
    }

  @Test
  fun nodeInvokeRequest_sendsResultForHandlerOwnedTimeout() =
    runBlocking {
      val result =
        runInvokeScenario(
          invokeEventFrame =
            """{"type":"event","event":"node.invoke.request","payload":{"id":"handler-timeout","nodeId":"node-1","command":"camera.snap","timeoutMs":5000}}""",
        ) {
          withTimeout(10) { awaitCancellation() }
        }

      assertEquals(
        false,
        result.resultParams["ok"]
          ?.jsonPrimitive
          ?.content
          ?.toBooleanStrict(),
      )
      assertEquals(
        "TIMEOUT",
        result.resultParams["error"]
          ?.jsonObject
          ?.get("code")
          ?.jsonPrimitive
          ?.content,
      )
    }

  @Test
  fun nodeInvokeRequest_sendsTimeoutWhileBlockingHandlerIsStillRunning() =
    runBlocking {
      val releaseHandler = CountDownLatch(1)
      val handlerFinished = CompletableDeferred<Unit>()
      val result =
        runInvokeScenario(
          invokeEventFrame =
            """{"type":"event","event":"node.invoke.request","payload":{"id":"blocking-timeout","nodeId":"node-1","command":"camera.clip","timeoutMs":100}}""",
          afterResult = {
            assertFalse(handlerFinished.isCompleted)
            releaseHandler.countDown()
            withTimeout(TEST_TIMEOUT_MS) { handlerFinished.await() }
          },
        ) {
          try {
            check(releaseHandler.await(5, TimeUnit.SECONDS)) { "blocking handler was not released" }
            GatewaySession.InvokeResult.ok(null)
          } finally {
            handlerFinished.complete(Unit)
          }
        }

      assertEquals(
        false,
        result.resultParams["ok"]
          ?.jsonPrimitive
          ?.content
          ?.toBooleanStrict(),
      )
      assertEquals(
        "TIMEOUT",
        result.resultParams["error"]
          ?.jsonObject
          ?.get("code")
          ?.jsonPrimitive
          ?.content,
      )
    }

  @Test
  fun nodeInvokeRequest_doesNotSendResultAfterCancellation() =
    runBlocking {
      val json = testJson()
      val connected = CompletableDeferred<Unit>()
      val invokeStarted = CompletableDeferred<Unit>()
      val invokeResult = CompletableDeferred<Unit>()
      val lastDisconnect = AtomicReference("")
      val serverWebSocket = AtomicReference<WebSocket?>(null)
      val server =
        startGatewayServer(json) { webSocket, id, method, _ ->
          serverWebSocket.set(webSocket)
          when (method) {
            "connect" -> {
              webSocket.send(connectResponseFrame(id))
              webSocket.send(
                """{"type":"event","event":"node.invoke.request","payload":{"id":"invoke-cancelled","nodeId":"node-1","command":"camera.snap","timeoutMs":5000}}""",
              )
            }
            "node.invoke.result" -> invokeResult.complete(Unit)
          }
        }
      val harness =
        createNodeHarness(connected = connected, lastDisconnect = lastDisconnect) {
          invokeStarted.complete(Unit)
          throw CancellationException("cancelled")
        }

      try {
        connectNodeSession(harness.session, server.port)
        awaitConnectedOrThrow(connected, lastDisconnect, server)
        withTimeout(TEST_TIMEOUT_MS) { invokeStarted.await() }

        assertNull(withTimeoutOrNull(250) { invokeResult.await() })
      } finally {
        serverWebSocket.get()?.close(1000, "done")
        delay(100)
        shutdownHarness(harness, server)
      }
    }

  @Test
  fun sendNodeEventDetailed_sendsPresenceAlivePayloadAndReturnsStructuredResponse() =
    runBlocking {
      val json = testJson()
      val connected = CompletableDeferred<Unit>()
      val nodeEventParams = CompletableDeferred<JsonObject>()
      val lastDisconnect = AtomicReference("")
      val server =
        startGatewayServer(json) { webSocket, id, method, frame ->
          when (method) {
            "connect" -> {
              webSocket.send(connectResponseFrame(id))
            }
            "node.event" -> {
              if (!nodeEventParams.isCompleted) {
                nodeEventParams.complete(frame["params"]?.jsonObject ?: JsonObject(emptyMap()))
              }
              val payload =
                """{"ok":true,"event":"node.presence.alive","handled":true,"reason":"persisted"}"""
              webSocket.send(
                """{"type":"res","id":"$id","ok":true,"payload":$payload}""",
              )
              webSocket.close(1000, "done")
            }
          }
        }

      val harness =
        createNodeHarness(
          connected = connected,
          lastDisconnect = lastDisconnect,
        ) { GatewaySession.InvokeResult.ok("""{"handled":true}""") }

      try {
        connectNodeSession(harness.session, server.port)
        awaitConnectedOrThrow(connected, lastDisconnect, server)

        val result =
          harness.session.sendNodeEventDetailed(
            event = "node.presence.alive",
            payloadJson = """{"trigger":"connect","sentAtMs":123}""",
            timeoutMs = TEST_TIMEOUT_MS,
          )
        val params = withTimeout(TEST_TIMEOUT_MS) { nodeEventParams.await() }
        val response = json.parseToJsonElement(result.payloadJson.orEmpty()).jsonObject
        val payload = json.parseToJsonElement(params["payloadJSON"]?.jsonPrimitive?.content.orEmpty()).jsonObject

        assertEquals(true, result.ok)
        assertEquals("node.presence.alive", params["event"]?.jsonPrimitive?.content)
        assertEquals("connect", payload["trigger"]?.jsonPrimitive?.content)
        assertEquals("123", payload["sentAtMs"]?.jsonPrimitive?.content)
        assertEquals(true, response["handled"]?.jsonPrimitive?.content?.toBooleanStrict())
        assertEquals("persisted", response["reason"]?.jsonPrimitive?.content)
      } finally {
        shutdownHarness(harness, server)
      }
    }

  @Test
  fun sendNodeEvent_preservesCompletedRpcAsSuccessWhenGatewayReturnsError() =
    runBlocking {
      val json = testJson()
      val connected = CompletableDeferred<Unit>()
      val nodeEventParams = CompletableDeferred<JsonObject>()
      val lastDisconnect = AtomicReference("")
      val server =
        startGatewayServer(json) { webSocket, id, method, frame ->
          when (method) {
            "connect" -> {
              webSocket.send(connectResponseFrame(id))
            }
            "node.event" -> {
              if (!nodeEventParams.isCompleted) {
                nodeEventParams.complete(frame["params"]?.jsonObject ?: JsonObject(emptyMap()))
              }
              webSocket.send(
                """{"type":"res","id":"$id","ok":false,"error":{"code":"RATE_LIMITED","message":"slow down"}}""",
              )
              webSocket.close(1000, "done")
            }
          }
        }

      val harness =
        createNodeHarness(
          connected = connected,
          lastDisconnect = lastDisconnect,
        ) { GatewaySession.InvokeResult.ok("""{"handled":true}""") }

      try {
        connectNodeSession(harness.session, server.port)
        awaitConnectedOrThrow(connected, lastDisconnect, server)

        val sent =
          harness.session.sendNodeEvent(
            event = "agent.request",
            payloadJson = """{"message":"restore"}""",
          )
        val params = withTimeout(TEST_TIMEOUT_MS) { nodeEventParams.await() }

        assertEquals(true, sent)
        assertEquals("agent.request", params["event"]?.jsonPrimitive?.content)
      } finally {
        shutdownHarness(harness, server)
      }
    }

  @Test
  fun sendNodeEvent_waitsForCompletedConnectHandshake() =
    runBlocking {
      val json = testJson()
      val connected = CompletableDeferred<Unit>()
      val connectRequestSeen = CompletableDeferred<Unit>()
      val releaseConnectResponse = CompletableDeferred<Unit>()
      val nodeEvents = CopyOnWriteArrayList<String>()
      val eventAfterConnect = CompletableDeferred<Unit>()
      val lastDisconnect = AtomicReference("")
      val server =
        startGatewayServer(json) { webSocket, id, method, frame ->
          when (method) {
            "connect" -> {
              connectRequestSeen.complete(Unit)
              launch(Dispatchers.Default) {
                releaseConnectResponse.await()
                webSocket.send(connectResponseFrame(id))
              }
            }
            "node.event" -> {
              val event =
                frame["params"]
                  ?.jsonObject
                  ?.get("event")
                  ?.jsonPrimitive
                  ?.content
                  .orEmpty()
              nodeEvents += event
              eventAfterConnect.complete(Unit)
              webSocket.send(
                """{"type":"res","id":"$id","ok":true,"payload":{"ok":true}}""",
              )
              webSocket.close(1000, "done")
            }
          }
        }
      val harness =
        createNodeHarness(
          connected = connected,
          lastDisconnect = lastDisconnect,
        ) { GatewaySession.InvokeResult.ok("""{"handled":true}""") }

      try {
        connectNodeSession(harness.session, server.port)
        withTimeout(TEST_TIMEOUT_MS) { connectRequestSeen.await() }

        assertFalse(
          harness.session.sendNodeEvent(
            event = "notifications.changed",
            payloadJson = """{"change":"posted","key":"before"}""",
          ),
        )
        assertTrue(nodeEvents.isEmpty())

        releaseConnectResponse.complete(Unit)
        awaitConnectedOrThrow(connected, lastDisconnect, server)
        assertTrue(
          harness.session.sendNodeEvent(
            event = "notifications.changed",
            payloadJson = """{"change":"posted","key":"after"}""",
          ),
        )
        withTimeout(TEST_TIMEOUT_MS) { eventAfterConnect.await() }
        assertEquals(listOf("notifications.changed"), nodeEvents.toList())
      } finally {
        releaseConnectResponse.complete(Unit)
        shutdownHarness(harness, server)
      }
    }

  private fun testJson(): Json = Json { ignoreUnknownKeys = true }

  private fun JsonObject.scopes(): List<String> =
    (this["scopes"] as? JsonArray)
      ?.map { it.jsonPrimitive.content }
      ?: emptyList()

  private fun createNodeHarness(
    connected: CompletableDeferred<Unit>,
    lastDisconnect: AtomicReference<String>,
    onEvent: (event: String, payloadJson: String?) -> Unit = { _, _ -> },
    onInvoke: suspend (GatewaySession.InvokeRequest) -> GatewaySession.InvokeResult,
  ): NodeHarness {
    val app = RuntimeEnvironment.getApplication()
    val sessionJob = SupervisorJob()
    val deviceAuthStore = InMemoryDeviceAuthStore()
    val session =
      GatewaySession(
        scope = CoroutineScope(sessionJob + Dispatchers.Default),
        identityStore = testDeviceIdentityStore(app),
        deviceAuthStore = deviceAuthStore,
        onConnected = {
          if (!connected.isCompleted) connected.complete(Unit)
        },
        onDisconnected = { message ->
          lastDisconnect.set(message)
        },
        onEvent = onEvent,
        onInvoke = onInvoke,
      )

    return NodeHarness(session = session, sessionJob = sessionJob, deviceAuthStore = deviceAuthStore)
  }

  private suspend fun connectNodeSession(
    session: GatewaySession,
    port: Int,
    token: String? = "test-token",
    bootstrapToken: String? = null,
    role: String = "node",
    scopes: List<String> = listOf("node:invoke"),
  ) {
    session.connect(
      endpoint =
        GatewayEndpoint(
          stableId = gatewayIdForPort(port),
          name = "test",
          host = "127.0.0.1",
          port = port,
          tlsEnabled = false,
        ),
      token = token,
      bootstrapToken = bootstrapToken,
      password = null,
      options =
        GatewayConnectOptions(
          role = role,
          scopes = scopes,
          caps = emptyList(),
          commands = emptyList(),
          permissions = emptyMap(),
          client =
            GatewayClientInfo(
              id = "openclaw-android-test",
              displayName = "Android Test",
              version = "1.0.0-test",
              platform = "android",
              mode = role,
              instanceId = "android-test-instance",
              deviceFamily = "android",
              modelIdentifier = "test",
            ),
        ),
      tls = null,
    )
  }

  private fun gatewayIdForPort(port: Int): String = "manual|127.0.0.1|$port"

  private suspend fun awaitConnectedOrThrow(
    connected: CompletableDeferred<Unit>,
    lastDisconnect: AtomicReference<String>,
    server: MockWebServer,
  ) {
    val connectedWithinTimeout =
      withTimeoutOrNull(TEST_TIMEOUT_MS) {
        connected.await()
        true
      } == true
    if (!connectedWithinTimeout) {
      throw AssertionError("never connected; lastDisconnect=${lastDisconnect.get()}; requests=${server.requestCount}")
    }
  }

  private suspend fun shutdownHarness(
    harness: NodeHarness,
    server: MockWebServer,
  ) {
    harness.session.disconnect()
    harness.sessionJob.cancelAndJoin()
    server.shutdown()
  }

  private suspend fun runInvokeScenario(
    invokeEventFrame: String,
    onHandshake: ((RecordedRequest) -> Unit)? = null,
    afterResult: suspend (InvokeScenarioResult) -> Unit = {},
    onInvoke: suspend (GatewaySession.InvokeRequest) -> GatewaySession.InvokeResult,
  ): InvokeScenarioResult {
    val json = testJson()
    val connected = CompletableDeferred<Unit>()
    val invokeRequest = CompletableDeferred<GatewaySession.InvokeRequest>()
    val invokeResultParams = CompletableDeferred<String>()
    val lastDisconnect = AtomicReference("")
    val server =
      startGatewayServer(
        json = json,
        onHandshake = onHandshake,
      ) { webSocket, id, method, frame ->
        when (method) {
          "connect" -> {
            webSocket.send(connectResponseFrame(id))
            webSocket.send(invokeEventFrame)
          }
          "node.invoke.result" -> {
            if (!invokeResultParams.isCompleted) {
              invokeResultParams.complete(frame["params"]?.toString().orEmpty())
            }
            webSocket.send("""{"type":"res","id":"$id","ok":true,"payload":{"ok":true}}""")
            webSocket.close(1000, "done")
          }
        }
      }
    val harness =
      createNodeHarness(
        connected = connected,
        lastDisconnect = lastDisconnect,
      ) { req ->
        if (!invokeRequest.isCompleted) invokeRequest.complete(req)
        onInvoke(req)
      }

    try {
      connectNodeSession(harness.session, server.port)
      awaitConnectedOrThrow(connected, lastDisconnect, server)
      val request = withTimeout(TEST_TIMEOUT_MS) { invokeRequest.await() }
      val resultParamsJson = withTimeout(TEST_TIMEOUT_MS) { invokeResultParams.await() }
      val resultParams = json.parseToJsonElement(resultParamsJson).jsonObject
      val result = InvokeScenarioResult(request = request, resultParams = resultParams)
      afterResult(result)
      return result
    } finally {
      shutdownHarness(harness, server)
    }
  }

  private fun connectResponseFrame(
    id: String,
    pluginSurfaceUrls: Map<String, String> = emptyMap(),
    authJson: String? = null,
  ): String {
    val surfaces =
      pluginSurfaceUrls.entries
        .joinToString(",") { (key, value) -> """"$key":"$value"""" }
        .takeIf { it.isNotEmpty() }
        ?.let { """"pluginSurfaceUrls":{$it},""" }
        ?: ""
    val auth = authJson?.let { "\"auth\":$it," } ?: ""
    return """{"type":"res","id":"$id","ok":true,"payload":{$surfaces$auth"snapshot":{"sessionDefaults":{"mainSessionKey":"main"}}}}"""
  }

  private fun startGatewayServer(
    json: Json,
    challengeFrame: String = CONNECT_CHALLENGE_FRAME,
    onHandshake: ((RecordedRequest) -> Unit)? = null,
    onRequestFrame: (webSocket: WebSocket, id: String, method: String, frame: JsonObject) -> Unit,
  ): MockWebServer =
    MockWebServer().apply {
      dispatcher =
        object : Dispatcher() {
          override fun dispatch(request: RecordedRequest): MockResponse {
            onHandshake?.invoke(request)
            return MockResponse().withWebSocketUpgrade(
              object : WebSocketListener() {
                override fun onOpen(
                  webSocket: WebSocket,
                  response: Response,
                ) {
                  webSocket.send(challengeFrame)
                }

                override fun onMessage(
                  webSocket: WebSocket,
                  text: String,
                ) {
                  val frame = json.parseToJsonElement(text).jsonObject
                  if (frame["type"]?.jsonPrimitive?.content != "req") return
                  val id = frame["id"]?.jsonPrimitive?.content ?: return
                  val method = frame["method"]?.jsonPrimitive?.content ?: return
                  onRequestFrame(webSocket, id, method, frame)
                }
              },
            )
          }
        }
      start()
    }
}
