Improve Wear OS error handling on home, check websocket connection st…

…ate (#2261)

* Websocket: make connected Job Deferred to get errors

 - Switch the job used for monitoring whether a connection has been established to use Deferred, which will allow us to get the exception that caused the job to be cancelled (and also makes the code a little easier)

* Websocket: add variable for state

 - Add a variable to the websocket repository to monitor the state of the connection. This will allow the app to respond to handled exceptions.

* More Websocket race condition fixes

 - If multiple outgoing messages are pending before the connection is opened, if we receive an invalid auth message or onFailure callback this will be processed last. The connected job is still 'completed', so it shouldn't return true but return if it was successful. This will ensure that any messages after the first one that triggered the failure, will be ignored until the job is reset.
 - Set the connection state to authenticating before sending the message, to make sure that we don't accidentally override a very quick response

* Handle error states in HomeView

 - Add a state for the loading progress to the view model in order to be able to update the home view with errors if necessary
 - If loading the main data failed because of authentication, reset the session like in the phone app and send the user back to onboarding
 - Fix the app getting stuck on a loading screen, even if everything loaded correctly, if there are no supported entity domains in the response

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()