Implement custom environment state

This way, we can show our status labels.
coder · code-asher · Sep 5, 2024 · Sep 12, 2024 · Sep 13, 2024 · Sep 13, 2024
commit 00d7321721dcf1f22c09c243913374f5a85b1fee
diff --git a/src/main/kotlin/com/coder/gateway/CoderRemoteEnvironment.kt b/src/main/kotlin/com/coder/gateway/CoderRemoteEnvironment.kt
@@ -1,16 +1,14 @@
 package com.coder.gateway
 
+import com.coder.gateway.models.WorkspaceAndAgentStatus
 import com.coder.gateway.sdk.CoderRestClient
 import com.coder.gateway.sdk.v2.models.Workspace
 import com.coder.gateway.sdk.v2.models.WorkspaceAgent
-import com.coder.gateway.sdk.v2.models.WorkspaceAgentStatus
-import com.coder.gateway.sdk.v2.models.WorkspaceStatus
 import com.coder.gateway.views.EnvironmentView
 import com.jetbrains.toolbox.gateway.AbstractRemoteProviderEnvironment
 import com.jetbrains.toolbox.gateway.EnvironmentVisibilityState
 import com.jetbrains.toolbox.gateway.environments.EnvironmentContentsView
 import com.jetbrains.toolbox.gateway.states.EnvironmentStateConsumer
-import com.jetbrains.toolbox.gateway.states.StandardRemoteEnvironmentState
 import com.jetbrains.toolbox.gateway.ui.ObservablePropertiesFactory
 import java.util.concurrent.CompletableFuture
 
@@ -27,28 +25,11 @@ 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)
 
-    // Active (and unhealthy) here indicate that the workspace is in a state
-    // where a connection can be attempted, not that the workspace is up and
-    // running.  Once a connection is actually initiated, the CLI will then
-    // start the workspace if it is off.
-    private var state = when (workspace.latestBuild.status) {
-        WorkspaceStatus.PENDING -> StandardRemoteEnvironmentState.Active
-        WorkspaceStatus.STARTING -> StandardRemoteEnvironmentState.Active
-        WorkspaceStatus.RUNNING -> when (agent.status) {
-            WorkspaceAgentStatus.CONNECTED -> StandardRemoteEnvironmentState.Active
-            WorkspaceAgentStatus.DISCONNECTED -> StandardRemoteEnvironmentState.Unreachable
-            WorkspaceAgentStatus.TIMEOUT -> StandardRemoteEnvironmentState.Unhealthy
-            WorkspaceAgentStatus.CONNECTING -> StandardRemoteEnvironmentState.Active
-        }
-        WorkspaceStatus.STOPPING -> StandardRemoteEnvironmentState.Initializing
-        WorkspaceStatus.STOPPED -> StandardRemoteEnvironmentState.Active
-        WorkspaceStatus.FAILED -> StandardRemoteEnvironmentState.Unhealthy
-        WorkspaceStatus.CANCELING -> StandardRemoteEnvironmentState.Initializing
-        WorkspaceStatus.CANCELED -> StandardRemoteEnvironmentState.Active
-        WorkspaceStatus.DELETING -> StandardRemoteEnvironmentState.Deleting
-        WorkspaceStatus.DELETED -> StandardRemoteEnvironmentState.Deleted
-    }
+
+    // Map each state to whether a connection can be attempted.
+    private var state = status.toRemoteEnvironmentState()

     /**
      * The contents are provided by the SSH view provided by Toolbox, all we

diff --git a/src/main/kotlin/com/coder/gateway/models/WorkspaceAndAgentStatus.kt b/src/main/kotlin/com/coder/gateway/models/WorkspaceAndAgentStatus.kt
@@ -5,6 +5,8 @@ import com.coder.gateway.sdk.v2.models.WorkspaceAgent
 import com.coder.gateway.sdk.v2.models.WorkspaceAgentLifecycleState
 import com.coder.gateway.sdk.v2.models.WorkspaceAgentStatus
 import com.coder.gateway.sdk.v2.models.WorkspaceStatus
+import com.jetbrains.toolbox.gateway.states.Color
+import com.jetbrains.toolbox.gateway.states.CustomRemoteEnvironmentState
 
 /**
  * WorkspaceAndAgentStatus represents the combined status of a single agent and
@@ -46,6 +48,30 @@ enum class WorkspaceAndAgentStatus(val label: String, val description: String) {
     READY("Ready", "The agent is ready to accept connections."),
     ;
 
+    /**
+     * Return the environment state for Toolbox, which tells it the label, color
+     * and whether the environment is reachable.
+     *
+     * We mark all ready and pending states as reachable since if the workspace
+     * is pending the cli will wait for it anyway.
+     *
+     * Additionally, terminal states like stopped are also marked as reachable,
+     * since the cli will start them.
+     */
+    fun toRemoteEnvironmentState(): CustomRemoteEnvironmentState {
+        // Use comments; no named arguments for non-Kotlin functions.
+        // TODO@JB: Is there a set of default colors we could use?
+        return CustomRemoteEnvironmentState(
+            label,
+            Color(200, 200, 200, 200), // darkThemeColor
+            Color(104, 112, 128, 255), // lightThemeColor
+            Color(224, 224, 240, 26), // darkThemeBackgroundColor
+            Color(224, 224, 245, 250), // lightThemeBackgroundColor
+            ready() || pending() || canStart(), // reachable
+            null, // iconId
+        )
+    }
+
     /**
      * Return true if the agent is in a connectable state.
      */
@@ -68,6 +94,12 @@ enum class WorkspaceAndAgentStatus(val label: String, val description: String) {
             .contains(this)
     }
 
+    /**
+     * Return true if the workspace can be started.
+     */
+    fun canStart(): Boolean = listOf(STOPPED, FAILED, CANCELED)
+        .contains(this)
+
     // We want to check that the workspace is `running`, the agent is
     // `connected`, and the agent lifecycle state is `ready` to ensure the best
     // possible scenario for attempting a connection.