add task query functionallity to worker_ws

1 files changed, 6 insertions, 9 deletions

diff --git a/src/webui_ws.rs b/src/webui_ws.rs
index 6b5821e..57393df 100644
--- a/src/webui_ws.rs
+++ b/src/webui_ws.rs
@@ -1,4 +1,7 @@
-use crate::{State, webui, worker_ws::WorkerID};
+use crate::{
+    State, webui,
+    worker_ws::{TaskState, WorkerID},
+};
 use axum::{
     extract::{
         State as S, WebSocketUpgrade,
@@ -37,20 +40,13 @@ pub enum WebuiEvent {
     },
 }
 
-#[derive(Debug, Serialize, Clone, Copy, PartialEq)]
-#[serde(rename_all = "snake_case")]
-pub enum TaskState {
-    Queue,
-    Loading,
-    Complete,
-}
-
 pub(crate) async fn webui_websocket(
     ws: WebSocketUpgrade,
     S(state): S<Arc<RwLock<State>>>,
 ) -> impl IntoResponse {
     ws.on_upgrade(|ws| webui_websocket_inner(ws, state))
 }
+
 async fn webui_websocket_inner(mut ws: WebSocket, state: Arc<RwLock<State>>) {
     let mut stream = state.read().await.webui_broadcast.subscribe();
     while let Ok(ev) = stream.recv().await {
@@ -63,6 +59,7 @@ async fn webui_websocket_inner(mut ws: WebSocket, state: Arc<RwLock<State>>) {
         }
     }
 }
+
 impl State {
     pub fn send_webui_task_removal(&self, key: &str) {
         if self.webui_broadcast.receiver_count() > 0 {