From 3f22a729b5248ecc4fca95b7a3215c9f90563b02 Mon Sep 17 00:00:00 2001
From: Ryan Lewis <accounts+github@ryandlewis.dev>
Date: Wed, 11 Sep 2024 18:00:53 -0500
Subject: [PATCH] Fix admin action corner cases

---
 .../components/AdminButtons/CancelButton.vue  |  1 +
 .../AdminButtons/LockUnlockButton.vue         |  9 +++++
 .../AdminButtons/PauseResumeButton.vue        | 12 ++++++-
 .../components/AdminButtons/ResetButton.vue   |  3 +-
 .../AdminButtons/SafetyStopButton.vue         |  3 +-
 .../AdminButtons/ShutdownButton.vue           |  3 +-
 src/ui/src/components/Dashboard.vue           | 12 +++----
 src/wei/core/admin.py                         | 34 ++++++++++++++-----
 src/wei/core/state_manager.py                 | 16 +++++++++
 9 files changed, 74 insertions(+), 19 deletions(-)

diff --git a/src/ui/src/components/AdminButtons/CancelButton.vue b/src/ui/src/components/AdminButtons/CancelButton.vue
index 91606d03..a25cb468 100644
--- a/src/ui/src/components/AdminButtons/CancelButton.vue
+++ b/src/ui/src/components/AdminButtons/CancelButton.vue
@@ -28,6 +28,7 @@
         main_url: string;
         module?: string;
         module_status?: string;
+        wc_state?: any;
     }>();
 
     const cancel_url = ref()
diff --git a/src/ui/src/components/AdminButtons/LockUnlockButton.vue b/src/ui/src/components/AdminButtons/LockUnlockButton.vue
index e630cd83..534f981d 100644
--- a/src/ui/src/components/AdminButtons/LockUnlockButton.vue
+++ b/src/ui/src/components/AdminButtons/LockUnlockButton.vue
@@ -27,6 +27,7 @@
     const props = defineProps<{
         main_url: string;
         module?: string;
+        wc_state?: any;
         module_status?: string;
     }>();
 
@@ -53,8 +54,16 @@
         if (props.module_status == 'LOCKED') {
             isLocked.value = true
         } else {
+            if (props.wc_state) {
+                if (props.wc_state.locked) {
+                    isLocked.value = true
+                } else {
+                    isLocked.value = false
+                }
+            } else {
             isLocked.value = false
         }
+    }
     })
     }
 
diff --git a/src/ui/src/components/AdminButtons/PauseResumeButton.vue b/src/ui/src/components/AdminButtons/PauseResumeButton.vue
index 374968fd..ac17529c 100644
--- a/src/ui/src/components/AdminButtons/PauseResumeButton.vue
+++ b/src/ui/src/components/AdminButtons/PauseResumeButton.vue
@@ -24,12 +24,13 @@
   </template>
 
 <script lang="ts" setup>
-    import {defineProps, ref, watchEffect} from 'vue';
+    import { defineProps, ref, watchEffect } from 'vue';
 
     const props = defineProps<{
         main_url: string;
         module?: string;
         module_status?: string;
+        wc_state?: any;
     }>();
 
     const pause_url = ref()
@@ -68,6 +69,15 @@
     })
     }
     else {
+        if (props.wc_state) {
+            if (props.wc_state.paused) {
+                isPaused.value = true
+            } else {
+                isPaused.value = false
+            }
+        } else {
+            isPaused.value = false
+        }
         allowButton.value = true
     }
 
diff --git a/src/ui/src/components/AdminButtons/ResetButton.vue b/src/ui/src/components/AdminButtons/ResetButton.vue
index 2fdb2ba3..074cfd47 100644
--- a/src/ui/src/components/AdminButtons/ResetButton.vue
+++ b/src/ui/src/components/AdminButtons/ResetButton.vue
@@ -22,12 +22,13 @@
 </template>
 
 <script lang="ts" setup>
-    import {defineProps, ref, watchEffect} from 'vue';
+    import { defineProps, ref, watchEffect } from 'vue';
 
     const props = defineProps<{
         main_url: string;
         module?: string;
         module_status?: string;
+        wc_state?: any;
     }>();
 
     const reset_url = ref()
diff --git a/src/ui/src/components/AdminButtons/SafetyStopButton.vue b/src/ui/src/components/AdminButtons/SafetyStopButton.vue
index 1278fb7d..bd04d3b1 100644
--- a/src/ui/src/components/AdminButtons/SafetyStopButton.vue
+++ b/src/ui/src/components/AdminButtons/SafetyStopButton.vue
@@ -21,12 +21,13 @@
 </template>
 
 <script lang="ts" setup>
-    import {defineProps, ref, watchEffect} from 'vue';
+    import { defineProps, ref } from 'vue';
 
     const props = defineProps<{
         main_url: string;
         module?: string;
         module_status?: string;
+        wc_state?: any;
     }>();
 
     const safetyStop_url = ref()
diff --git a/src/ui/src/components/AdminButtons/ShutdownButton.vue b/src/ui/src/components/AdminButtons/ShutdownButton.vue
index c4561523..81ddce04 100644
--- a/src/ui/src/components/AdminButtons/ShutdownButton.vue
+++ b/src/ui/src/components/AdminButtons/ShutdownButton.vue
@@ -22,12 +22,13 @@
 </template>
 
 <script lang="ts" setup>
-    import {defineProps, ref, watchEffect} from 'vue';
+    import { defineProps, ref, watchEffect } from 'vue';
 
     const props = defineProps<{
         main_url: string;
         module?: string;
         module_status?: string;
+        wc_state?: any;
     }>();
 
     const shutdown_url = ref()
diff --git a/src/ui/src/components/Dashboard.vue b/src/ui/src/components/Dashboard.vue
index 77aebaac..cf38aac7 100644
--- a/src/ui/src/components/Dashboard.vue
+++ b/src/ui/src/components/Dashboard.vue
@@ -21,12 +21,12 @@
               <v-card-title class="text-center">
                 <h2>{{ wc_info.name }}</h2>
                 <div class="d-flex justify-center">
-                  <PauseResumeButton :main_url="main_url" class="ml-2"/>
-                  <CancelButton :main_url="main_url" class="ml-2" />
-                  <ResetButton :main_url="main_url" class="ml-2" />
-                  <LockUnlockButton :main_url="main_url" class="ml-2"/>
-                  <ShutdownButton :main_url="main_url" class="ml-2" />
-                  <SafetyStopButton :main_url="main_url" class="ml-2"/>
+                  <PauseResumeButton :main_url="main_url" :wc_state=wc_state class="ml-2"/>
+                  <CancelButton :main_url="main_url" :wc_state=wc_state class="ml-2" />
+                  <ResetButton :main_url="main_url" :wc_state=wc_state class="ml-2" />
+                  <LockUnlockButton :main_url="main_url" :wc_state=wc_state class="ml-2"/>
+                  <ShutdownButton :main_url="main_url" :wc_state=wc_state class="ml-2" />
+                  <SafetyStopButton :main_url="main_url" :wc_state=wc_state class="ml-2"/>
                 </div>
               </v-card-title>
             <v-card-text>
diff --git a/src/wei/core/admin.py b/src/wei/core/admin.py
index f85bfb15..7b9de1f6 100644
--- a/src/wei/core/admin.py
+++ b/src/wei/core/admin.py
@@ -1,12 +1,14 @@
 """Handles admin commands and related logic."""
 
 from wei.types.interface_types import InterfaceMap
-from wei.types.module_types import AdminCommands, Module
+from wei.types.module_types import AdminCommands, Module, ModuleStatus
+from wei.utils import threaded_task
 
 
+@threaded_task
 def send_safety_stop(module: Module) -> None:
     """Safety stops a module"""
-    if AdminCommands.SAFETY_STOP in module.about.admin_commands:
+    if check_can_send_admin_command(module, AdminCommands.SAFETY_STOP):
         InterfaceMap.interfaces[module.interface].send_admin_command(
             module, AdminCommands.SAFETY_STOP
         )
@@ -16,9 +18,10 @@ def send_safety_stop(module: Module) -> None:
         send_pause(module)
 
 
+@threaded_task
 def send_reset(module: Module) -> None:
     """Resets a module"""
-    if AdminCommands.RESET in module.about.admin_commands:
+    if check_can_send_admin_command(module, AdminCommands.RESET):
         InterfaceMap.interfaces[module.interface].send_admin_command(
             module, AdminCommands.RESET
         )
@@ -27,9 +30,10 @@ def send_reset(module: Module) -> None:
         print(f"Module {module.name} does not support resetting.")
 
 
+@threaded_task
 def send_pause(module: Module) -> None:
     """Pauses a module"""
-    if AdminCommands.PAUSE in module.about.admin_commands:
+    if check_can_send_admin_command(module, AdminCommands.PAUSE):
         InterfaceMap.interfaces[module.interface].send_admin_command(
             module, AdminCommands.PAUSE
         )
@@ -39,9 +43,10 @@ def send_pause(module: Module) -> None:
         send_cancel(module)
 
 
+@threaded_task
 def send_resume(module: Module) -> None:
     """Resumes a module"""
-    if AdminCommands.RESUME in module.about.admin_commands:
+    if check_can_send_admin_command(module, AdminCommands.RESUME):
         InterfaceMap.interfaces[module.interface].send_admin_command(
             module, AdminCommands.RESUME
         )
@@ -50,9 +55,10 @@ def send_resume(module: Module) -> None:
         print(f"Module {module.name} does not support resuming.")
 
 
+@threaded_task
 def send_cancel(module: Module) -> None:
     """Cancels a module"""
-    if AdminCommands.CANCEL in module.about.admin_commands:
+    if check_can_send_admin_command(module, AdminCommands.CANCEL):
         InterfaceMap.interfaces[module.interface].send_admin_command(
             module, AdminCommands.CANCEL
         )
@@ -61,9 +67,10 @@ def send_cancel(module: Module) -> None:
         print(f"Module {module.name} does not support canceling.")
 
 
+@threaded_task
 def send_shutdown(module: Module) -> None:
     """Shuts down a module"""
-    if AdminCommands.SHUTDOWN in module.about.admin_commands:
+    if check_can_send_admin_command(module, AdminCommands.SHUTDOWN):
         InterfaceMap.interfaces[module.interface].send_admin_command(
             module, AdminCommands.SHUTDOWN
         )
@@ -72,9 +79,10 @@ def send_shutdown(module: Module) -> None:
         print(f"Module {module.name} does not support shutting down.")
 
 
+@threaded_task
 def send_lock(module: Module) -> None:
     """Locks a module"""
-    if AdminCommands.LOCK in module.about.admin_commands:
+    if check_can_send_admin_command(module, AdminCommands.LOCK):
         InterfaceMap.interfaces[module.interface].send_admin_command(
             module, AdminCommands.LOCK
         )
@@ -83,12 +91,20 @@ def send_lock(module: Module) -> None:
         print(f"Module {module.name} does not support locking.")
 
 
+@threaded_task
 def send_unlock(module: Module) -> None:
     """Unlocks a module"""
-    if AdminCommands.UNLOCK in module.about.admin_commands:
+    if check_can_send_admin_command(module, AdminCommands.UNLOCK):
         InterfaceMap.interfaces[module.interface].send_admin_command(
             module, AdminCommands.UNLOCK
         )
         print(f"Module {module.name} has been unlocked.")
     else:
         print(f"Module {module.name} does not support unlocking.")
+
+
+def check_can_send_admin_command(module: Module, command: AdminCommands) -> bool:
+    """Checks if a module can send an admin command"""
+    return not module.state.status == ModuleStatus.UNKNOWN and (
+        module.about is None or command in module.about.admin_commands
+    )
diff --git a/src/wei/core/state_manager.py b/src/wei/core/state_manager.py
index 70b60280..93f42af7 100644
--- a/src/wei/core/state_manager.py
+++ b/src/wei/core/state_manager.py
@@ -111,6 +111,18 @@ def paused(self, value: bool) -> None:
             f"{self._workcell_prefix}:paused", "true" if value else "false"
         )
 
+    @property
+    def locked(self) -> bool:
+        """Get the lock state of the workcell"""
+        return self._redis_client.get(f"{self._workcell_prefix}:locked") == "true"
+
+    @locked.setter
+    def locked(self, value: bool) -> None:
+        """Set the lock state of the workcell"""
+        self._redis_client.set(
+            f"{self._workcell_prefix}:locked", "true" if value else "false"
+        )
+
     @property
     def shutdown(self) -> bool:
         """Get the shutdown state of the workcell"""
@@ -167,6 +179,9 @@ def get_state(self) -> Dict[str, Dict[Any, Any]]:
             "modules": self._modules.to_dict(),
             "workflows": self._workflow_runs.to_dict(),
             "workcell": self._workcell.to_dict(),
+            "paused": self.paused,
+            "locked": self.locked,
+            "shutdown": self.shutdown,
         }
 
     @property
@@ -207,6 +222,7 @@ def clear_state(
         self._workcell.clear()
         self.state_change_marker = "0"
         self.paused = False
+        self.locked = False
         self.shutdown = False
         self.mark_state_changed()