feat: WebSocket heartbeat, multi-vault subscription, and connection s…

android/app/src/main/java/com/ethosprotocol/services/VaultEventSocket.kt

@@ -14,11 +14,17 @@ import javax.inject.Inject
 import javax.inject.Singleton
 import kotlin.random.Random
 import kotlinx.coroutines.CancellationException
+import kotlinx.coroutines.cancel
+import kotlinx.coroutines.coroutineScope
 import kotlinx.coroutines.currentCoroutineContext
 import kotlinx.coroutines.delay
 import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.asStateFlow
 import kotlinx.coroutines.flow.flow
 import kotlinx.coroutines.isActive
+import kotlinx.coroutines.launch
 import kotlinx.serialization.json.Json
 
 // Reconnect/backoff schedule for VaultEventSocket. Unlike RetryPolicy (bounded
@@ -49,6 +55,11 @@ data class ReconnectBackoff(
     }
 }
 
+// Connection state for the WebSocket. FALLBACK_TO_POLLING is included for parity
+// with the iOS API contract; on Android the socket retries indefinitely unless
+// maxReconnectAttempts is set to a finite value.
+enum class ConnectionState { DISCONNECTED, CONNECTING, CONNECTED, FALLBACK_TO_POLLING }
+
 // Client for the `wss://.../ws?vault_id={id}` endpoint (shared/api-contract.md).
 // Reuses ApiClient's HttpClient/WebSockets plugin rather than a second client.
 // A dropped or failed connection reconnects with exponential backoff
@@ -58,14 +69,19 @@ data class ReconnectBackoff(
 class VaultEventSocket(
     private val apiClient: ApiClient,
     private val tokenProvider: TokenProvider,
-    private val backoff: ReconnectBackoff = ReconnectBackoff.default
+    private val backoff: ReconnectBackoff = ReconnectBackoff.default,
+    val heartbeatIntervalMillis: Long = 30_000L,
+    val maxReconnectAttempts: Int = Int.MAX_VALUE
 ) {
     // Distinct @Inject constructor for the same reason as ApiClient's: Dagger's
     // codegen calls the full-arg constructor directly and ignores Kotlin default
     // values, so Hilt needs an explicit constructor matching only its bindings.
     @Inject constructor(apiClient: ApiClient, tokenProvider: TokenProvider) :
         this(apiClient, tokenProvider, ReconnectBackoff.default)
 
+    private val _connectionState = MutableStateFlow(ConnectionState.DISCONNECTED)
+    val connectionState: StateFlow<ConnectionState> = _connectionState.asStateFlow()
+
     // internal (not private): tests replace this to simulate connect failures/frames
     // without a real server, matching how ApiClient exposes its engine to tests.
     internal var openSession: suspend (String) -> WebSocketSession = { vaultId ->
@@ -78,21 +94,107 @@ class VaultEventSocket(
     fun events(vaultId: String): Flow<VaultEvent> = flow {
         var attempt = 0
         while (currentCoroutineContext().isActive) {
+            _connectionState.value = ConnectionState.CONNECTING
             try {
                 val session = openSession(vaultId)
                 attempt = 0
-                for (frame in session.incoming) {
-                    if (frame is Frame.Text) {
-                        runCatching { Json.decodeFromString<VaultEvent>(frame.readText()) }
-                            .onSuccess { emit(it) }
+                _connectionState.value = ConnectionState.CONNECTED
+                coroutineScope {
+                    // Periodic client-side heartbeat to detect silent TCP drops.
+                    launch {
+                        while (isActive) {
+                            delay(heartbeatIntervalMillis)
+                            try {
+                                session.send(Frame.Text("""{ "type": "ping" }"""))
+                            } catch (_: Exception) {
+                                // Send failed — session is dead; cancel scope to stop incoming loop too.
+                                cancel()
+                            }
+                        }
+                    }
+                    for (frame in session.incoming) {
+                        if (frame is Frame.Text) {
+                            val text = frame.readText()
+                            runCatching { Json.decodeFromString<VaultEvent>(text) }
+                                .onSuccess { event ->
+                                    if (event.type == "ping") {
+                                        // Server keepalive — reply with pong.
+                                        runCatching { session.send(Frame.Text("""{ "type": "pong" }""")) }
+                                    }
+                                    emit(event)
+                                }
+                        }
                     }
+                    cancel() // incoming closed normally — stop heartbeat too
                 }
             } catch (e: CancellationException) {
                 throw e
             } catch (e: Exception) {
                 // Connection failed or dropped — fall through to backoff and reconnect.
+                _connectionState.value = ConnectionState.DISCONNECTED
             }
             if (!currentCoroutineContext().isActive) break
+            if (attempt >= maxReconnectAttempts) {
+                _connectionState.value = ConnectionState.FALLBACK_TO_POLLING
+                break
+            }
+            backoff.sleep(backoff.delayForAttempt(attempt))
+            attempt++
+        }
+    }
+
+    fun events(vaultIds: List<String>): Flow<VaultEvent> = flow {
+        if (vaultIds.isEmpty()) return@flow
+        // Use the first vault ID as the primary connection URL, then send a subscribe
+        // message for the rest once connected.
+        var attempt = 0
+        while (currentCoroutineContext().isActive) {
+            _connectionState.value = ConnectionState.CONNECTING
+            try {
+                val session = openSession(vaultIds[0])
+                attempt = 0
+                _connectionState.value = ConnectionState.CONNECTED
+                // Send multi-vault subscribe message post-connect
+                if (vaultIds.size > 1) {
+                    val idsJson = vaultIds.joinToString(",") { "\"$it\"" }
+                    runCatching { session.send(Frame.Text("""{ "type": "subscribe", "vault_ids": [$idsJson] }""")) }
+                }
+                coroutineScope {
+                    launch {
+                        while (isActive) {
+                            delay(heartbeatIntervalMillis)
+                            try {
+                                session.send(Frame.Text("""{ "type": "ping" }"""))
+                            } catch (_: Exception) {
+                                cancel()
+                            }
+                        }
+                    }
+                    for (frame in session.incoming) {
+                        if (frame is Frame.Text) {
+                            val text = frame.readText()
+                            runCatching { Json.decodeFromString<VaultEvent>(text) }
+                                .onSuccess { event ->
+                                    if (event.type == "ping") {
+                                        runCatching { session.send(Frame.Text("""{ "type": "pong" }""")) }
+                                    }
+                                    emit(event)
+                                }
+                        }
+                    }
+                    cancel()
+                }
+            } catch (e: CancellationException) {
+                throw e
+            } catch (e: Exception) {
+                // connection failed
+                _connectionState.value = ConnectionState.DISCONNECTED
+            }
+            if (!currentCoroutineContext().isActive) break
+            if (attempt >= maxReconnectAttempts) {
+                _connectionState.value = ConnectionState.FALLBACK_TO_POLLING
+                break
+            }
             backoff.sleep(backoff.delayForAttempt(attempt))
             attempt++
         }

android/app/src/main/java/com/ethosprotocol/ui/ViewModels.kt

@@ -352,7 +352,8 @@ data class VaultUiState(
     val hasMore: Boolean = false,
     val error: String? = null,
     val isOffline: Boolean = false,
-    val beneficiaryUpdated: Boolean = false
+    val beneficiaryUpdated: Boolean = false,
+    val socketConnectionState: com.ethosprotocol.services.ConnectionState = com.ethosprotocol.services.ConnectionState.DISCONNECTED
 )
 
 private const val PAGE_SIZE = 20
@@ -372,6 +373,14 @@ class VaultViewModel @Inject constructor(
     private var nextCursor: String? = null
     private val eventJobs = mutableMapOf<String, Job>()
 
+    init {
+        viewModelScope.launch {
+            vaultEventSocket.connectionState.collect { state ->
+                _state.update { it.copy(socketConnectionState = state) }
+            }
+        }
+    }
+
     fun load() = viewModelScope.launch {
         _state.update { it.copy(isLoading = true, error = null) }
         when (val result = apiClient.listVaults(limit = PAGE_SIZE)) {
@@ -463,13 +472,28 @@ class VaultViewModel @Inject constructor(
     // check-ins/deposits/withdrawals made elsewhere (another device, an expiry)
     // update this list in place instead of requiring a manual refresh.
     private fun subscribeToEvents(vaultIds: List<String>) {
+        // Cancel subscriptions for vaults no longer in the list
         val currentIds = vaultIds.toSet()
         eventJobs.keys.filter { it !in currentIds }.forEach { id -> eventJobs.remove(id)?.cancel() }
-        currentIds.filterNot { eventJobs.containsKey(it) }.forEach { id ->
-            eventJobs[id] = viewModelScope.launch {
-                vaultEventSocket.events(id).collect { event ->
-                    val updated = event.vault ?: return@collect
-                    _state.update { s -> s.copy(vaults = s.vaults.map { if (it.id == updated.id) updated else it }) }
+        // Subscribe new vaults. If more than one ID is requested, use a single multiplexed connection.
+        if (vaultIds.size > 1 && eventJobs.isEmpty()) {
+            // Multiplex all vaults over a single connection via subscribe message
+            val combinedKey = vaultIds.joinToString(",")
+            if (!eventJobs.containsKey(combinedKey)) {
+                eventJobs[combinedKey] = viewModelScope.launch {
+                    vaultEventSocket.events(vaultIds).collect { event ->
+                        val updated = event.vault ?: return@collect
+                        _state.update { s -> s.copy(vaults = s.vaults.map { if (it.id == updated.id) updated else it }) }
+                    }
+                }
+            }
+        } else {
+            currentIds.filterNot { eventJobs.containsKey(it) }.forEach { id ->
+                eventJobs[id] = viewModelScope.launch {
+                    vaultEventSocket.events(id).collect { event ->
+                        val updated = event.vault ?: return@collect
+                        _state.update { s -> s.copy(vaults = s.vaults.map { if (it.id == updated.id) updated else it }) }
+                    }
                 }
             }
         }

android/app/src/main/java/com/ethosprotocol/ui/screens/Screens.kt

@@ -1382,6 +1382,34 @@ private fun TwoFactorVerifyScreen(
     }
 }
 
+// MARK: - Connection Status Badge
+
+@Composable
+fun ConnectionStatusBadge(state: com.ethosprotocol.services.ConnectionState) {
+    val (label, color, icon) = when (state) {
+        com.ethosprotocol.services.ConnectionState.CONNECTED ->
+            Triple("Live", MaterialTheme.colorScheme.primary, Icons.Default.Wifi)
+        com.ethosprotocol.services.ConnectionState.CONNECTING ->
+            Triple("Connecting…", MaterialTheme.colorScheme.primary, Icons.Default.Wifi)
+        com.ethosprotocol.services.ConnectionState.DISCONNECTED ->
+            Triple("Reconnecting…", MaterialTheme.colorScheme.error, Icons.Default.WifiOff)
+        com.ethosprotocol.services.ConnectionState.FALLBACK_TO_POLLING ->
+            Triple("Polling", MaterialTheme.colorScheme.outline, Icons.Default.Refresh)
+    }
+    SuggestionChip(
+        onClick = {},
+        label = {
+            Row(verticalAlignment = Alignment.CenterVertically) {
+                Icon(icon, contentDescription = null, modifier = Modifier.size(12.dp))
+                Spacer(Modifier.width(4.dp))
+                Text(label, style = MaterialTheme.typography.labelSmall)
+            }
+        },
+        modifier = Modifier.testTag("connectionStatus_$label"),
+        colors = SuggestionChipDefaults.suggestionChipColors(labelColor = color)
+    )
+}
+
 // MARK: - Vault Detail Screen
 
 /**
@@ -1395,7 +1423,8 @@ private fun TwoFactorVerifyScreen(
 fun VaultDetailScreen(
     vaultId: String,
     onBack: () -> Unit,
-    twoFactorVm: TwoFactorViewModel = hiltViewModel()
+    twoFactorVm: TwoFactorViewModel = hiltViewModel(),
+    vaultVm: VaultViewModel = hiltViewModel()
 ) {
     val state by twoFactorVm.state.collectAsStateWithLifecycle()
     val context = LocalContext.current
@@ -1433,6 +1462,17 @@ fun VaultDetailScreen(
             Text("Two-Factor Authentication", style = MaterialTheme.typography.titleMedium)
             Spacer(Modifier.height(8.dp))
 
+            val socketState by vaultVm.state.collectAsStateWithLifecycle()
+            Row(
+                modifier = Modifier.fillMaxWidth().padding(bottom = 8.dp),
+                verticalAlignment = Alignment.CenterVertically
+            ) {
+                Text("Connection", style = MaterialTheme.typography.bodySmall,
+                    color = MaterialTheme.colorScheme.onSurfaceVariant,
+                    modifier = Modifier.weight(1f))
+                ConnectionStatusBadge(socketState.socketConnectionState)
+            }
+
             when {
                 state.isLoading -> CircularProgressIndicator()
                 state.error != null && state.status == null -> {