websockets

src/server/mod.rs

@@ -1,5 +1,6 @@
 pub mod auth;
 pub mod chat;
+pub mod notifications;
 pub mod tasks;
 
 use axum::{
@@ -24,6 +25,7 @@ pub struct AppState {
     pub config: Arc<crate::config::Config>,
     pub verifier: Arc<crate::domain::auth::JwksVerifier>,
     pub authenticator: Arc<crate::domain::auth::Authenticator>,
+    pub tx: tokio::sync::broadcast::Sender<crate::server::notifications::WsEvent>,
 }
 
 pub async fn start(config: crate::config::Config) -> AppResult<()> {
@@ -31,8 +33,10 @@ pub async fn start(config: crate::config::Config) -> AppResult<()> {
 
     let config = Arc::new(config);
 
+    let (tx, _) = tokio::sync::broadcast::channel(100);
+
     let scheduler = Arc::new(
-        Scheduler::new(db.clone(), config.clone())
+        Scheduler::new(db.clone(), config.clone(), tx.clone())
             .await
             .map_err(|e| crate::error::AppError::Internal(e.to_string()))?,
     );
@@ -56,6 +60,7 @@ pub async fn start(config: crate::config::Config) -> AppResult<()> {
         config: config.clone(),
         verifier,
         authenticator,
+        tx,
     });
 
     let app = build_app(state, &config);
@@ -122,6 +127,7 @@ fn build_app(state: Arc<AppState>, config: &crate::config::Config) -> Router {
         .route("/api/auth/refresh", post(auth::auth_refresh))
         .route("/api/auth/logout", post(auth::auth_logout))
         .route("/api/chat", post(chat::chat_handler))
+        .route("/api/ws", get(notifications::ws_handler))
         .layer(cors)
         .layer(tower_http::set_header::SetResponseHeaderLayer::overriding(
             axum::http::header::CONTENT_SECURITY_POLICY,

src/server/notifications.rs

@@ -0,0 +1,46 @@
+use crate::domain::tasks::{RecentRunResponse, TaskResponse};
+use crate::server::AppState;
+use axum::{
+    extract::{
+        State,
+        ws::{Message, WebSocket, WebSocketUpgrade},
+    },
+    response::IntoResponse,
+};
+use serde::Serialize;
+use std::sync::Arc;
+
+#[derive(Serialize, Clone, Debug)]
+#[serde(tag = "type", content = "data")]
+pub enum WsEvent {
+    TaskCreated(TaskResponse),
+    TaskUpdated(TaskResponse),
+    RunStarted(RecentRunResponse),
+    RunFinished(TaskResponse),
+}
+
+pub async fn ws_handler(
+    ws: WebSocketUpgrade,
+    State(state): State<Arc<AppState>>,
+) -> impl IntoResponse {
+    ws.on_upgrade(|socket| handle_socket(socket, state))
+}
+
+async fn handle_socket(mut socket: WebSocket, state: Arc<AppState>) {
+    let mut rx = state.tx.subscribe();
+
+    while let Ok(event) = rx.recv().await {
+        let msg = match serde_json::to_string(&event) {
+            Ok(json) => json,
+            Err(e) => {
+                tracing::error!("Failed to serialize WsEvent: {}", e);
+                continue;
+            }
+        };
+
+        if socket.send(Message::Text(msg)).await.is_err() {
+            // Client disconnected
+            break;
+        }
+    }
+}

src/server/tasks.rs

@@ -84,7 +84,13 @@ pub async fn create_task(
         let _ = state.scheduler.add_task_job(task_id, cron).await;
     }
 
-    tasks::get_task_inner(task_id, &state.db).await.map(Json)
+    let task_response = tasks::get_task_inner(task_id, &state.db).await?;
+    let _ = state
+        .tx
+        .send(crate::server::notifications::WsEvent::TaskCreated(
+            task_response.clone(),
+        ));
+    Ok(Json(task_response))
 }
 
 pub async fn rerun_task(
@@ -150,7 +156,13 @@ pub async fn update_task(
         let _ = state.scheduler.remove_task_job(id).await;
     }
 
-    tasks::get_task_inner(id, &state.db).await.map(Json)
+    let task_response = tasks::get_task_inner(id, &state.db).await?;
+    let _ = state
+        .tx
+        .send(crate::server::notifications::WsEvent::TaskUpdated(
+            task_response.clone(),
+        ));
+    Ok(Json(task_response))
 }
 
 pub async fn get_task(