Fix environment status not updating

code-asher · code-asher · commit 1d83f50e19d4 · 2024-09-16T13:04:30.000-08:00
It looks like you have to use the per-environment listener.  My guess is
that they do not update an environment if the ID is the same.  I made
the environment comparable, which also lets use use a set instead.
diff --git a/src/main/kotlin/com/coder/gateway/CoderRemoteEnvironment.kt b/src/main/kotlin/com/coder/gateway/CoderRemoteEnvironment.kt
@@ -25,11 +25,19 @@ class CoderRemoteEnvironment(
 ) : AbstractRemoteProviderEnvironment(observablePropertiesFactory) {
     override fun getId(): String = "${workspace.name}.${agent.name}"
     override fun getName(): String = "${workspace.name}.${agent.name}"
-    private val status = WorkspaceAndAgentStatus.from(workspace, agent)
+    private var status = WorkspaceAndAgentStatus.from(workspace, agent)
 
-
-    // Map each state to whether a connection can be attempted.
-    private var state = status.toRemoteEnvironmentState()
+    /**
+     * Update the workspace/agent status to the listeners, if it has changed.
+     */
+    fun update(workspace: Workspace, agent: WorkspaceAgent) {
+        val newStatus = WorkspaceAndAgentStatus.from(workspace, agent)
+        if (newStatus != status) {
+            status = newStatus
+            val state = status.toRemoteEnvironmentState()
+            listenerSet.forEach { it.consume(state) }
+        }
+    }
 
     /**
      * The contents are provided by the SSH view provided by Toolbox, all we
@@ -47,14 +55,26 @@ class CoderRemoteEnvironment(
     override fun setVisible(visibilityState: EnvironmentVisibilityState) {}
 
     /**
-     * Immediately send the state to the listener.
-     *
-     * Currently we consume the entire workspace list and are not updating
-     * individual workspaces, so the state here is static and the listener is
-     * only used once.
+     * Immediately send the state to the listener and store for updates.
      */
     override fun addStateListener(consumer: EnvironmentStateConsumer): Boolean {
-        consumer.consume(state)
+        consumer.consume(status.toRemoteEnvironmentState())
         return super.addStateListener(consumer)
     }
+
+    /**
+     * An environment is equal if it has the same ID.
+     */
+    override fun equals(other: Any?): Boolean {
+        if (other == null) return false
+        if (this === other) return true // Note the triple ===
+        if (other !is CoderRemoteEnvironment) return false
+        if (getId() != other.getId()) return false
+        return true
+    }
+
+    /**
+     * Companion to equals, for sets.
+     */
+    override fun hashCode(): Int = getId().hashCode()
 }
diff --git a/src/main/kotlin/com/coder/gateway/CoderRemoteProvider.kt b/src/main/kotlin/com/coder/gateway/CoderRemoteProvider.kt
@@ -51,7 +51,7 @@ class CoderRemoteProvider(
 
     // Current polling job.
     private var pollJob: Job? = null
-    private var lastEnvironments: List<CoderRemoteEnvironment>? = null
+    private var lastEnvironments: Set<CoderRemoteEnvironment>? = null
 
     // Create our services from the Toolbox ones.
     private val settingsService = CoderSettingsService(settingsStore)
@@ -96,23 +96,27 @@ class CoderRemoteProvider(
                             //       different information?
                             it.name
                         }?.map { agent ->
-                            CoderRemoteEnvironment(client, ws, agent, observablePropertiesFactory)
+                            // If we have an environment already, update that.
+                            val env = CoderRemoteEnvironment(client, ws, agent, observablePropertiesFactory)
+                            lastEnvironments?.firstOrNull { it == env }?.let {
+                                it.update(ws, agent)
+                                it
+                            } ?: env
                         } ?: emptyList()
                     }
-                }
+                }.toSet()
 
                 // In case we logged out while running the query.
                 if (!isActive) {
                     return@launch
                 }
 
                 // Reconfigure if a new environment is found.
-                val newEnvironments = environments
-                    .filter { a -> lastEnvironments?.any { b -> a.id == b.id } != true }
-                    .map { it.name }.toSet()
-                if (newEnvironments.isNotEmpty()) {
+                // TODO@JB: Should we use the add/remove listeners instead?
+                val newEnvironments = lastEnvironments?.let { environments.subtract(it) }
+                if (newEnvironments?.isNotEmpty() == true) {
                     logger.info("Found new environment(s), reconfiguring CLI: {}", newEnvironments)
-                    cli.configSsh(newEnvironments)
+                    cli.configSsh(newEnvironments.map { it.name }.toSet())
                 }
 
                 consumer.consumeEnvironments(environments)
