feat: add queryCountSuspend utility to Quartz and simplify ViewModels

Add INostrClient.queryCountSuspend() extension functions in Quartz that
wrap NIP-45 COUNT queries as suspend functions, managing subscription
lifecycle internally. Two overloads: single-relay and multi-relay.

Simplify BasicRelaySetupInfoModel and Nip65RelayListViewModel to use the
new utility — removes IRelayClientListener implementation, subId tracking
maps, onIncomingMessage handlers, and cleanup logic from both ViewModels.

https://claude.ai/code/session_016158D5mq5BygS1uBbLNbsA

amethyst/src/main/java/com/vitorpamplona/amethyst/ui/screen/loggedIn/relays/common/BasicRelaySetupInfoModel.kt

@@ -26,11 +26,7 @@ import com.vitorpamplona.amethyst.Amethyst
 import com.vitorpamplona.amethyst.model.Account
 import com.vitorpamplona.amethyst.service.replace
 import com.vitorpamplona.amethyst.ui.screen.loggedIn.AccountViewModel
-import com.vitorpamplona.quartz.nip01Core.relay.client.listeners.IRelayClientListener
-import com.vitorpamplona.quartz.nip01Core.relay.client.single.IRelayClient
-import com.vitorpamplona.quartz.nip01Core.relay.client.single.newSubId
-import com.vitorpamplona.quartz.nip01Core.relay.commands.toClient.CountMessage
-import com.vitorpamplona.quartz.nip01Core.relay.commands.toClient.Message
+import com.vitorpamplona.quartz.nip01Core.relay.client.accessories.queryCountSuspend
 import com.vitorpamplona.quartz.nip01Core.relay.normalizer.NormalizedRelayUrl
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.flow.MutableStateFlow
@@ -38,7 +34,7 @@ import kotlinx.coroutines.flow.asStateFlow
 import kotlinx.coroutines.flow.update
 import kotlinx.coroutines.launch
 
-abstract class BasicRelaySetupInfoModel : ViewModel(), IRelayClientListener {
+abstract class BasicRelaySetupInfoModel : ViewModel() {
     lateinit var accountViewModel: AccountViewModel
     lateinit var account: Account
 
@@ -48,9 +44,6 @@ abstract class BasicRelaySetupInfoModel : ViewModel(), IRelayClientListener {
     private val _countResults = MutableStateFlow<Map<NormalizedRelayUrl, RelayCountResult>>(emptyMap())
     val countResults = _countResults.asStateFlow()
 
-    private val subIdToRelay = mutableMapOf<String, Pair<NormalizedRelayUrl, Int>>()
-    private val relayQueryInfos = mutableMapOf<NormalizedRelayUrl, List<CountQueryInfo>>()
-
     var hasModified = false
 
     fun init(accountViewModel: AccountViewModel) {
@@ -94,82 +87,39 @@ abstract class BasicRelaySetupInfoModel : ViewModel(), IRelayClientListener {
     }
 
     private fun loadCounts() {
-        val client = Amethyst.instance.client
-        cleanupCounts()
+        _countResults.value = emptyMap()
 
+        val client = Amethyst.instance.client
         val relayList = _relays.value
         if (relayList.isEmpty()) return
 
-        val hasFilters = relayList.any { countFilters(it.relay).isNotEmpty() }
-        if (!hasFilters) return
-
-        client.subscribe(this)
-
         relayList.forEach { item ->
             val filters = countFilters(item.relay)
             if (filters.isEmpty()) return@forEach
 
-            val queryInfos = mutableListOf<CountQueryInfo>()
-
-            filters.forEachIndexed { index, countFilter ->
-                val subId = newSubId()
-                subIdToRelay[subId] = Pair(item.relay, index)
-                queryInfos.add(CountQueryInfo(subId, countFilter.label, index))
-
-                client.queryCount(
-                    subId = subId,
-                    filters = mapOf(item.relay to listOf(countFilter.filter)),
-                )
-            }
-
-            relayQueryInfos[item.relay] = queryInfos
-        }
-    }
-
-    override fun onIncomingMessage(
-        relay: IRelayClient,
-        msgStr: String,
-        msg: Message,
-    ) {
-        if (msg is CountMessage) {
-            val (relayUrl, filterIndex) = subIdToRelay[msg.queryId] ?: return
-            val queryInfos = relayQueryInfos[relayUrl] ?: return
-            val queryInfo = queryInfos.find { it.filterIndex == filterIndex } ?: return
-
-            _countResults.update { currentMap ->
-                val currentResult = currentMap[relayUrl] ?: RelayCountResult()
-                val updatedEntries = currentResult.counts.toMutableList()
-
-                val existingIndex = updatedEntries.indexOfFirst { it.label == queryInfo.label }
-                val newEntry =
-                    RelayCountResult.CountEntry(
-                        label = queryInfo.label,
-                        count = msg.result.count,
-                        approximate = msg.result.approximate,
-                    )
-
-                if (existingIndex >= 0) {
-                    updatedEntries[existingIndex] = newEntry
-                } else {
-                    updatedEntries.add(newEntry)
+            filters.forEach { countFilter ->
+                viewModelScope.launch(Dispatchers.IO) {
+                    val result = client.queryCountSuspend(item.relay, countFilter.filter)
+                    if (result != null) {
+                        _countResults.update { currentMap ->
+                            val current = currentMap[item.relay] ?: RelayCountResult()
+                            val entries = current.counts.toMutableList()
+                            val newEntry =
+                                RelayCountResult.CountEntry(
+                                    label = countFilter.label,
+                                    count = result.count,
+                                    approximate = result.approximate,
+                                )
+                            val existing = entries.indexOfFirst { it.label == countFilter.label }
+                            if (existing >= 0) entries[existing] = newEntry else entries.add(newEntry)
+                            currentMap + (item.relay to RelayCountResult(entries))
+                        }
+                    }
                 }
-
-                currentMap + (relayUrl to RelayCountResult(updatedEntries))
             }
         }
     }
 
-    private fun cleanupCounts() {
-        val client = Amethyst.instance.client
-        subIdToRelay.keys.forEach { subId ->
-            client.close(subId)
-        }
-        subIdToRelay.clear()
-        relayQueryInfos.clear()
-        _countResults.value = emptyMap()
-        client.unsubscribe(this)
-    }
-
     open fun relayListBuilder(): List<BasicRelaySetupInfo> {
         val relayList = getRelayList() ?: emptyList()
 
@@ -213,15 +163,4 @@ abstract class BasicRelaySetupInfoModel : ViewModel(), IRelayClientListener {
     ) {
         _relays.update { it.replace(relay, relay.copy(paidRelay = paid)) }
     }
-
-    override fun onCleared() {
-        cleanupCounts()
-        super.onCleared()
-    }
-
-    private data class CountQueryInfo(
-        val subId: String,
-        val label: String,
-        val filterIndex: Int,
-    )
 }

amethyst/src/main/java/com/vitorpamplona/amethyst/ui/screen/loggedIn/relays/nip65/Nip65RelayListViewModel.kt

@@ -30,11 +30,7 @@ import com.vitorpamplona.amethyst.ui.screen.loggedIn.AccountViewModel
 import com.vitorpamplona.amethyst.ui.screen.loggedIn.relays.common.BasicRelaySetupInfo
 import com.vitorpamplona.amethyst.ui.screen.loggedIn.relays.common.RelayCountResult
 import com.vitorpamplona.amethyst.ui.screen.loggedIn.relays.common.relaySetupInfoBuilder
-import com.vitorpamplona.quartz.nip01Core.relay.client.listeners.IRelayClientListener
-import com.vitorpamplona.quartz.nip01Core.relay.client.single.IRelayClient
-import com.vitorpamplona.quartz.nip01Core.relay.client.single.newSubId
-import com.vitorpamplona.quartz.nip01Core.relay.commands.toClient.CountMessage
-import com.vitorpamplona.quartz.nip01Core.relay.commands.toClient.Message
+import com.vitorpamplona.quartz.nip01Core.relay.client.accessories.queryCountSuspend
 import com.vitorpamplona.quartz.nip01Core.relay.filters.Filter
 import com.vitorpamplona.quartz.nip01Core.relay.normalizer.NormalizedRelayUrl
 import com.vitorpamplona.quartz.nip65RelayList.tags.AdvertisedRelayInfo
@@ -46,7 +42,7 @@ import kotlinx.coroutines.flow.update
 import kotlinx.coroutines.launch
 
 @Stable
-class Nip65RelayListViewModel : ViewModel(), IRelayClientListener {
+class Nip65RelayListViewModel : ViewModel() {
     private lateinit var accountViewModel: AccountViewModel
     private lateinit var account: Account
 
@@ -62,8 +58,6 @@ class Nip65RelayListViewModel : ViewModel(), IRelayClientListener {
     private val _notifCountResults = MutableStateFlow<Map<NormalizedRelayUrl, RelayCountResult>>(emptyMap())
     val notifCountResults = _notifCountResults.asStateFlow()
 
-    private val subIdToRelay = mutableMapOf<String, Pair<NormalizedRelayUrl, Boolean>>()
-
     var hasModified = false
 
     fun init(accountViewModel: AccountViewModel) {
@@ -129,71 +123,48 @@ class Nip65RelayListViewModel : ViewModel(), IRelayClientListener {
     }
 
     private fun loadCounts() {
-        val client = Amethyst.instance.client
-        cleanupCounts()
-
-        val homeList = _homeRelays.value
-        val notifList = _notificationRelays.value
-
-        if (homeList.isEmpty() && notifList.isEmpty()) return
-
-        client.subscribe(this)
-
-        homeList.forEach { item ->
-            val subId = newSubId()
-            subIdToRelay[subId] = Pair(item.relay, true)
-            client.queryCount(
-                subId = subId,
-                filters = mapOf(item.relay to listOf(Filter(authors = listOf(account.pubKey)))),
-            )
-        }
-
-        notifList.forEach { item ->
-            val subId = newSubId()
-            subIdToRelay[subId] = Pair(item.relay, false)
-            client.queryCount(
-                subId = subId,
-                filters = mapOf(item.relay to listOf(Filter(tags = mapOf("p" to listOf(account.pubKey))))),
-            )
-        }
-    }
-
-    override fun onIncomingMessage(
-        relay: IRelayClient,
-        msgStr: String,
-        msg: Message,
-    ) {
-        if (msg is CountMessage) {
-            val (relayUrl, isHome) = subIdToRelay[msg.queryId] ?: return
-
-            val newResult =
-                RelayCountResult(
-                    listOf(
-                        RelayCountResult.CountEntry(
-                            label = "events",
-                            count = msg.result.count,
-                            approximate = msg.result.approximate,
-                        ),
-                    ),
-                )
-
-            if (isHome) {
-                _homeCountResults.update { it + (relayUrl to newResult) }
-            } else {
-                _notifCountResults.update { it + (relayUrl to newResult) }
-            }
-        }
-    }
-
-    private fun cleanupCounts() {
-        val client = Amethyst.instance.client
-        subIdToRelay.keys.forEach { subId ->
-            client.close(subId)
-        }
-        subIdToRelay.clear()
         _homeCountResults.value = emptyMap()
         _notifCountResults.value = emptyMap()
-        client.unsubscribe(this)
+
+        val client = Amethyst.instance.client
+
+        _homeRelays.value.forEach { item ->
+            viewModelScope.launch(Dispatchers.IO) {
+                val result = client.queryCountSuspend(item.relay, Filter(authors = listOf(account.pubKey)))
+                if (result != null) {
+                    val countResult =
+                        RelayCountResult(
+                            listOf(
+                                RelayCountResult.CountEntry(
+                                    label = "events",
+                                    count = result.count,
+                                    approximate = result.approximate,
+                                ),
+                            ),
+                        )
+                    _homeCountResults.update { it + (item.relay to countResult) }
+                }
+            }
+        }
+
+        _notificationRelays.value.forEach { item ->
+            viewModelScope.launch(Dispatchers.IO) {
+                val result = client.queryCountSuspend(item.relay, Filter(tags = mapOf("p" to listOf(account.pubKey))))
+                if (result != null) {
+                    val countResult =
+                        RelayCountResult(
+                            listOf(
+                                RelayCountResult.CountEntry(
+                                    label = "events",
+                                    count = result.count,
+                                    approximate = result.approximate,
+                                ),
+                            ),
+                        )
+                    _notifCountResults.update { it + (item.relay to countResult) }
+                }
+            }
+        }
     }
 
     fun clear() {
@@ -266,9 +237,4 @@ class Nip65RelayListViewModel : ViewModel(), IRelayClientListener {
     ) {
         _notificationRelays.update { it.replace(relay, relay.copy(paidRelay = paid)) }
     }
-
-    override fun onCleared() {
-        cleanupCounts()
-        super.onCleared()
-    }
 }

quartz/src/commonMain/kotlin/com/vitorpamplona/quartz/nip01Core/relay/client/accessories/NostrClientCountExt.kt

@@ -0,0 +1,136 @@
+/*
+ * Copyright (c) 2025 Vitor Pamplona
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy of
+ * this software and associated documentation files (the "Software"), to deal in
+ * the Software without restriction, including without limitation the rights to use,
+ * copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the
+ * Software, and to permit persons to whom the Software is furnished to do so,
+ * subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
+ * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
+ * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN
+ * AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
+ * WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+ */
+package com.vitorpamplona.quartz.nip01Core.relay.client.accessories
+
+import com.vitorpamplona.quartz.nip01Core.relay.client.INostrClient
+import com.vitorpamplona.quartz.nip01Core.relay.client.listeners.IRelayClientListener
+import com.vitorpamplona.quartz.nip01Core.relay.client.single.IRelayClient
+import com.vitorpamplona.quartz.nip01Core.relay.client.single.newSubId
+import com.vitorpamplona.quartz.nip01Core.relay.commands.toClient.CountMessage
+import com.vitorpamplona.quartz.nip01Core.relay.commands.toClient.CountResult
+import com.vitorpamplona.quartz.nip01Core.relay.commands.toClient.Message
+import com.vitorpamplona.quartz.nip01Core.relay.filters.Filter
+import com.vitorpamplona.quartz.nip01Core.relay.normalizer.NormalizedRelayUrl
+import kotlinx.coroutines.channels.Channel
+import kotlinx.coroutines.channels.Channel.Factory.UNLIMITED
+import kotlinx.coroutines.withTimeoutOrNull
+
+/**
+ * Sends a NIP-45 COUNT query to a single relay and suspends until
+ * the result arrives or the timeout expires.
+ *
+ * @param relay Target relay to query.
+ * @param filter The filter to count against.
+ * @param timeoutMs How long to wait for a response (default 15 s).
+ * @return The [CountResult], or `null` on timeout.
+ */
+suspend fun INostrClient.queryCountSuspend(
+    relay: NormalizedRelayUrl,
+    filter: Filter,
+    timeoutMs: Long = 15_000,
+): CountResult? {
+    val subId = newSubId()
+    val resultChannel = Channel<CountResult>(UNLIMITED)
+
+    val listener =
+        object : IRelayClientListener {
+            override fun onIncomingMessage(
+                relay: IRelayClient,
+                msgStr: String,
+                msg: Message,
+            ) {
+                if (msg is CountMessage && msg.queryId == subId) {
+                    resultChannel.trySend(msg.result)
+                }
+            }
+        }
+
+    subscribe(listener)
+
+    queryCount(subId = subId, filters = mapOf(relay to listOf(filter)))
+
+    val result =
+        withTimeoutOrNull(timeoutMs) {
+            resultChannel.receive()
+        }
+
+    close(subId)
+    unsubscribe(listener)
+    resultChannel.close()
+
+    return result
+}
+
+/**
+ * Sends NIP-45 COUNT queries to multiple relays in parallel
+ * (one filter per relay) and suspends until all results arrive
+ * or the timeout expires.
+ *
+ * @param filters Map of relay -> filter to count.
+ * @param timeoutMs How long to wait for all responses (default 15 s).
+ * @return Map of relay -> [CountResult] for every relay that responded in time.
+ */
+suspend fun INostrClient.queryCountSuspend(
+    filters: Map<NormalizedRelayUrl, List<Filter>>,
+    timeoutMs: Long = 15_000,
+): Map<NormalizedRelayUrl, CountResult> {
+    if (filters.isEmpty()) return emptyMap()
+
+    val subIdToRelay = mutableMapOf<String, NormalizedRelayUrl>()
+    val resultChannel = Channel<Pair<NormalizedRelayUrl, CountResult>>(UNLIMITED)
+
+    val listener =
+        object : IRelayClientListener {
+            override fun onIncomingMessage(
+                relay: IRelayClient,
+                msgStr: String,
+                msg: Message,
+            ) {
+                if (msg is CountMessage) {
+                    val relayUrl = subIdToRelay[msg.queryId] ?: return
+                    resultChannel.trySend(relayUrl to msg.result)
+                }
+            }
+        }
+
+    subscribe(listener)
+
+    filters.forEach { (relay, filterList) ->
+        val subId = newSubId()
+        subIdToRelay[subId] = relay
+        queryCount(subId = subId, filters = mapOf(relay to filterList))
+    }
+
+    val results = mutableMapOf<NormalizedRelayUrl, CountResult>()
+
+    withTimeoutOrNull(timeoutMs) {
+        while (results.size < filters.size) {
+            val (relay, result) = resultChannel.receive()
+            results[relay] = result
+        }
+    }
+
+    subIdToRelay.keys.forEach { close(it) }
+    unsubscribe(listener)
+    resultChannel.close()
+
+    return results
+}