Improve Wear OS error handling on home, check websocket connection state

common/src/main/java/io/homeassistant/companion/android/common/data/websocket/WebSocketRepository.kt

@@ -15,6 +15,7 @@ import kotlinx.coroutines.flow.Flow
 
 @ExperimentalCoroutinesApi
 interface WebSocketRepository {
+    fun getConnectionState(): WebSocketState?
     suspend fun sendPing(): Boolean
     suspend fun getConfig(): GetConfigResponse?
     suspend fun getStates(): List<EntityResponse<Any>>?

common/src/main/java/io/homeassistant/companion/android/common/data/websocket/WebSocketState.kt

@@ -0,0 +1,5 @@
+package io.homeassistant.companion.android.common.data.websocket
+
+enum class WebSocketState {
+    AUTHENTICATING, ACTIVE, CLOSED_AUTH, CLOSED_OTHER
+}

common/src/main/java/io/homeassistant/companion/android/common/data/websocket/impl/WebSocketRepositoryImpl.kt

@@ -14,6 +14,7 @@ import io.homeassistant.companion.android.common.data.integration.ServiceData
 import io.homeassistant.companion.android.common.data.integration.impl.entities.EntityResponse
 import io.homeassistant.companion.android.common.data.url.UrlRepository
 import io.homeassistant.companion.android.common.data.websocket.WebSocketRepository
+import io.homeassistant.companion.android.common.data.websocket.WebSocketState
 import io.homeassistant.companion.android.common.data.websocket.impl.entities.AreaRegistryResponse
 import io.homeassistant.companion.android.common.data.websocket.impl.entities.AreaRegistryUpdatedEvent
 import io.homeassistant.companion.android.common.data.websocket.impl.entities.DeviceRegistryResponse
@@ -26,6 +27,7 @@ import io.homeassistant.companion.android.common.data.websocket.impl.entities.Ge
 import io.homeassistant.companion.android.common.data.websocket.impl.entities.SocketResponse
 import io.homeassistant.companion.android.common.data.websocket.impl.entities.StateChangedEvent
 import kotlinx.coroutines.CancellableContinuation
+import kotlinx.coroutines.CompletableDeferred
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.ExperimentalCoroutinesApi
@@ -78,15 +80,18 @@ class WebSocketRepositoryImpl @Inject constructor(
     private val responseCallbackJobs = mutableMapOf<Long, CancellableContinuation<SocketResponse>>()
     private val id = AtomicLong(1)
     private var connection: WebSocket? = null
+    private var connectionState: WebSocketState? = null
     private val connectedMutex = Mutex()
-    private var connected = Job()
+    private var connected = CompletableDeferred<Boolean>()
     private val eventSubscriptionMutex = Mutex()
     private val eventSubscriptionFlow = mutableMapOf<String, SharedFlow<*>>()
     private var eventSubscriptionProducerScope = mutableMapOf<String, ProducerScope<Any>>()
     private val notificationMutex = Mutex()
     private var notificationFlow: Flow<Map<String, Any>>? = null
     private var notificationProducerScope: ProducerScope<Map<String, Any>>? = null
 
+    override fun getConnectionState(): WebSocketState? = connectionState
+
     override suspend fun sendPing(): Boolean {
         val socketResponse = sendMessage(
             mapOf(
@@ -256,7 +261,7 @@ class WebSocketRepositoryImpl @Inject constructor(
     private suspend fun connect(): Boolean {
         connectedMutex.withLock {
             if (connection != null && connected.isCompleted) {
-                return true
+                return !connected.isCancelled
             }
 
             val url = urlRepository.getUrl()
@@ -275,6 +280,7 @@ class WebSocketRepositoryImpl @Inject constructor(
                 this
             ).also {
                 // Preemptively send auth
+                connectionState = WebSocketState.AUTHENTICATING
                 it.send(
                     mapper.writeValueAsString(
                         mapOf(
@@ -288,9 +294,7 @@ class WebSocketRepositoryImpl @Inject constructor(
             // Wait up to 30 seconds for auth response
             return true == withTimeoutOrNull(30000) {
                 return@withTimeoutOrNull try {
-                    connected.join()
-                    if (connected.isCancelled) throw AuthorizationException()
-                    true
+                    connected.await()
                 } catch (e: Exception) {
                     Log.e(TAG, "Unable to authenticate", e)
                     false
@@ -327,10 +331,13 @@ class WebSocketRepositoryImpl @Inject constructor(
         if (response?.result != null) mapper.convertValue(response.result) else null
 
     private fun handleAuthComplete(successful: Boolean) {
-        if (successful)
-            connected.complete()
-        else
-            connected.completeExceptionally(Exception("Authentication Error"))
+        if (successful) {
+            connectionState = WebSocketState.ACTIVE
+            connected.complete(true)
+        } else {
+            connectionState = WebSocketState.CLOSED_AUTH
+            connected.completeExceptionally(AuthorizationException())
+        }
     }
 
     private fun handleMessage(response: SocketResponse) {
@@ -378,8 +385,10 @@ class WebSocketRepositoryImpl @Inject constructor(
     private fun handleClosingSocket() {
         ioScope.launch {
             connectedMutex.withLock {
-                connected = Job()
+                connected = CompletableDeferred()
                 connection = null
+                if (connectionState != WebSocketState.CLOSED_AUTH)
+                    connectionState = WebSocketState.CLOSED_OTHER
             }
         }
         // If we still have flows flowing

common/src/main/res/values/strings.xml

@@ -171,6 +171,7 @@
     <string name="entity_widget_desc">Current state and attribute of any entity</string>
     <string name="error_auth_revoked">It appears that your authorization was revoked, please reconnect to Home Assistant.</string>
     <string name="error_connection_failed">Unable to connect to Home Assistant.</string>
+    <string name="error_loading_entities">Error loading entities</string>
     <string name="error_onboarding_connection_failed">Unable to connect to Home Assistant.</string>
     <string name="error_ssl">Unable to communicate with Home Assistant because of a SSL error.  Please ensure your certificate is valid.</string>
     <string name="error_with_registration">Please check to ensure you have the mobile_app\nintegration enabled on your home assistant instance.</string>
@@ -329,6 +330,8 @@
     <string name="nfc_write_tag_too_early">Please fill out the form first</string>
     <string name="no_notifications_summary">You have not received any notifications yet</string>
     <string name="no_notifications">No Notifications</string>
+    <string name="no_supported_entities">No Supported Entities</string>
+    <string name="no_supported_entities_summary">Check the documentation for more information</string>
     <string name="no_widgets_summary">When you add a widget to your home screen, you will see it here</string>
     <string name="no_widgets">No Widgets Found</string>
     <string name="none">None</string>

wear/src/main/java/io/homeassistant/companion/android/home/HomePresenter.kt

@@ -1,6 +1,7 @@
 package io.homeassistant.companion.android.home
 
 import io.homeassistant.companion.android.common.data.integration.Entity
+import io.homeassistant.companion.android.common.data.websocket.WebSocketState
 import io.homeassistant.companion.android.common.data.websocket.impl.entities.AreaRegistryResponse
 import io.homeassistant.companion.android.common.data.websocket.impl.entities.AreaRegistryUpdatedEvent
 import io.homeassistant.companion.android.common.data.websocket.impl.entities.DeviceRegistryResponse
@@ -17,9 +18,11 @@ interface HomePresenter {
     fun onViewReady()
     suspend fun onEntityClicked(entityId: String, state: String)
     fun onLogoutClicked()
+    fun onInvalidAuthorization()
     fun onFinish()
 
     suspend fun isConnected(): Boolean
+    fun getWebSocketState(): WebSocketState?
 
     suspend fun getEntities(): List<Entity<*>>?
     suspend fun getEntityUpdates(): Flow<Entity<*>>?

wear/src/main/java/io/homeassistant/companion/android/home/HomePresenterImpl.kt

@@ -8,6 +8,7 @@ import io.homeassistant.companion.android.common.data.integration.DeviceRegistra
 import io.homeassistant.companion.android.common.data.integration.Entity
 import io.homeassistant.companion.android.common.data.integration.IntegrationRepository
 import io.homeassistant.companion.android.common.data.websocket.WebSocketRepository
+import io.homeassistant.companion.android.common.data.websocket.WebSocketState
 import io.homeassistant.companion.android.common.data.websocket.impl.entities.AreaRegistryResponse
 import io.homeassistant.companion.android.common.data.websocket.impl.entities.AreaRegistryUpdatedEvent
 import io.homeassistant.companion.android.common.data.websocket.impl.entities.DeviceRegistryResponse
@@ -101,7 +102,11 @@ class HomePresenterImpl @Inject constructor(
         )
     }
 
-    override fun onLogoutClicked() {
+    override fun onInvalidAuthorization() = finishSession()
+
+    override fun onLogoutClicked() = finishSession()
+
+    private fun finishSession() {
         mainScope.launch {
             authenticationUseCase.revokeSession()
             view.displayOnBoarding()
@@ -132,6 +137,10 @@ class HomePresenterImpl @Inject constructor(
         return integrationUseCase.isRegistered()
     }
 
+    override fun getWebSocketState(): WebSocketState? {
+        return webSocketUseCase.getConnectionState()
+    }
+
     override suspend fun getAreaRegistry(): List<AreaRegistryResponse>? {
         return webSocketUseCase.getAreaRegistry()
     }

wear/src/main/java/io/homeassistant/companion/android/home/MainViewModel.kt

@@ -1,6 +1,7 @@
 package io.homeassistant.companion.android.home
 
 import android.app.Application
+import android.util.Log
 import androidx.compose.runtime.mutableStateListOf
 import androidx.compose.runtime.mutableStateMapOf
 import androidx.compose.runtime.mutableStateOf
@@ -10,6 +11,7 @@ import androidx.lifecycle.viewModelScope
 import dagger.hilt.android.lifecycle.HiltViewModel
 import io.homeassistant.companion.android.HomeAssistantApplication
 import io.homeassistant.companion.android.common.data.integration.Entity
+import io.homeassistant.companion.android.common.data.websocket.WebSocketState
 import io.homeassistant.companion.android.common.data.websocket.impl.entities.AreaRegistryResponse
 import io.homeassistant.companion.android.common.data.websocket.impl.entities.DeviceRegistryResponse
 import io.homeassistant.companion.android.common.data.websocket.impl.entities.EntityRegistryResponse
@@ -27,6 +29,14 @@ import javax.inject.Inject
 @HiltViewModel
 class MainViewModel @Inject constructor(application: Application) : AndroidViewModel(application) {
 
+    companion object {
+        const val TAG = "MainViewModel"
+    }
+
+    enum class LoadingState {
+        LOADING, READY, ERROR
+    }
+
     private lateinit var homePresenter: HomePresenter
     val app = getApplication<HomeAssistantApplication>()
     private val favoritesDao = AppDatabase.getInstance(app.applicationContext).favoritesDao()
@@ -38,6 +48,7 @@ class MainViewModel @Inject constructor(application: Application) : AndroidViewM
     // TODO: This is bad, do this instead: https://stackoverflow.com/questions/46283981/android-viewmodel-additional-arguments
     fun init(homePresenter: HomePresenter) {
         this.homePresenter = homePresenter
+        loadSettings()
         loadEntities()
         getFavorites()
         getSensors()
@@ -68,6 +79,8 @@ class MainViewModel @Inject constructor(application: Application) : AndroidViewM
     var entityListFilter: (Entity<*>) -> Boolean = { true }
 
     // settings
+    var loadingState = mutableStateOf(LoadingState.LOADING)
+        private set
     var isHapticEnabled = mutableStateOf(false)
         private set
     var isToastEnabled = mutableStateOf(false)
@@ -90,7 +103,7 @@ class MainViewModel @Inject constructor(application: Application) : AndroidViewM
 
     var sensors = mutableStateListOf<Sensor>()
 
-    private fun loadEntities() {
+    private fun loadSettings() {
         viewModelScope.launch {
             if (!homePresenter.isConnected()) {
                 return@launch
@@ -101,49 +114,76 @@ class MainViewModel @Inject constructor(application: Application) : AndroidViewM
             isShowShortcutTextEnabled.value = homePresenter.getShowShortcutText()
             templateTileContent.value = homePresenter.getTemplateTileContent()
             templateTileRefreshInterval.value = homePresenter.getTemplateTileRefreshInterval()
+        }
+    }
 
-            homePresenter.getAreaRegistry()?.let {
-                areaRegistry = it
-                areas.addAll(it)
+    fun loadEntities() {
+        viewModelScope.launch {
+            if (!homePresenter.isConnected()) {
+                return@launch
             }
-            deviceRegistry = homePresenter.getDeviceRegistry()
-            entityRegistry = homePresenter.getEntityRegistry()
-            homePresenter.getEntities()?.forEach {
-                if (supportedDomains().contains(it.entityId.split(".")[0])) {
-                    entities[it.entityId] = it
+            try {
+                // Load initial state
+                loadingState.value = LoadingState.LOADING
+                homePresenter.getAreaRegistry()?.let {
+                    areaRegistry = it
+                    areas.addAll(it)
                 }
-            }
-            updateEntityDomains()
-
-            viewModelScope.launch {
-                homePresenter.getEntityUpdates()?.collect {
+                deviceRegistry = homePresenter.getDeviceRegistry()
+                entityRegistry = homePresenter.getEntityRegistry()
+                homePresenter.getEntities()?.forEach {
                     if (supportedDomains().contains(it.entityId.split(".")[0])) {
                         entities[it.entityId] = it
-                        updateEntityDomains()
                     }
                 }
-            }
-            viewModelScope.launch {
-                homePresenter.getAreaRegistryUpdates()?.collect {
-                    areaRegistry = homePresenter.getAreaRegistry()
-                    areas.clear()
-                    areaRegistry?.let {
-                        areas.addAll(it)
+                updateEntityDomains()
+
+                // Finished initial load, update state
+                val webSocketState = homePresenter.getWebSocketState()
+                if (webSocketState == WebSocketState.CLOSED_AUTH) {
+                    homePresenter.onInvalidAuthorization()
+                    return@launch
+                }
+                loadingState.value = if (webSocketState == WebSocketState.ACTIVE) {
+                    LoadingState.READY
+                } else {
+                    LoadingState.ERROR
+                }
+
+                // Listen for updates
+                viewModelScope.launch {
+                    homePresenter.getEntityUpdates()?.collect {
+                        if (supportedDomains().contains(it.entityId.split(".")[0])) {
+                            entities[it.entityId] = it
+                            updateEntityDomains()
+                        }
                     }
-                    updateEntityDomains()
                 }
-            }
-            viewModelScope.launch {
-                homePresenter.getDeviceRegistryUpdates()?.collect {
-                    deviceRegistry = homePresenter.getDeviceRegistry()
-                    updateEntityDomains()
+                viewModelScope.launch {
+                    homePresenter.getAreaRegistryUpdates()?.collect {
+                        areaRegistry = homePresenter.getAreaRegistry()
+                        areas.clear()
+                        areaRegistry?.let {
+                            areas.addAll(it)
+                        }
+                        updateEntityDomains()
+                    }
                 }
-            }
-            viewModelScope.launch {
-                homePresenter.getEntityRegistryUpdates()?.collect {
-                    entityRegistry = homePresenter.getEntityRegistry()
-                    updateEntityDomains()
+                viewModelScope.launch {
+                    homePresenter.getDeviceRegistryUpdates()?.collect {
+                        deviceRegistry = homePresenter.getDeviceRegistry()
+                        updateEntityDomains()
+                    }
+                }
+                viewModelScope.launch {
+                    homePresenter.getEntityRegistryUpdates()?.collect {
+                        entityRegistry = homePresenter.getEntityRegistry()
+                        updateEntityDomains()
+                    }
                 }
+            } catch (e: Exception) {
+                Log.e(TAG, "Exception while loading entities", e)
+                loadingState.value = LoadingState.ERROR
             }
         }
     }

wear/src/main/java/io/homeassistant/companion/android/home/views/HomeView.kt

@@ -56,7 +56,9 @@ fun LoadHomePage(
     val context = LocalContext.current
 
     WearAppTheme {
-        if (mainViewModel.entities.isNullOrEmpty() && mainViewModel.favoriteEntityIds.isNullOrEmpty()) {
+        if (mainViewModel.loadingState.value == MainViewModel.LoadingState.LOADING &&
+            mainViewModel.favoriteEntityIds.isNullOrEmpty()
+        ) {
             Column {
                 ListHeader(id = commonR.string.loading)
                 Chip(
@@ -83,6 +85,7 @@ fun LoadHomePage(
                         mainViewModel,
                         mainViewModel.favoriteEntityIds,
                         { id, state -> mainViewModel.toggleEntity(id, state) },
+                        mainViewModel::loadEntities,
                         { swipeDismissableNavController.navigate(SCREEN_SETTINGS) },
                         { lists, order, filter ->
                             mainViewModel.entityLists.clear()