diff --git a/Cargo.lock b/Cargo.lock index 73d8ffb..8030dcb 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -161,6 +161,7 @@ checksum = "edca88bc138befd0323b20752846e6587272d3b03b0343c8ea28a6f819e6e71f" dependencies = [ "async-trait", "axum-core", + "base64", "bytes", "futures-util", "http", @@ -179,8 +180,10 @@ dependencies = [ "serde_json", "serde_path_to_error", "serde_urlencoded", + "sha1", "sync_wrapper", "tokio", + "tokio-tungstenite", "tower", "tower-layer", "tower-service", @@ -336,7 +339,7 @@ dependencies = [ "sea-orm-migration", "serde", "serde_json", - "thiserror", + "thiserror 2.0.18", "tokio", "tokio-cron-scheduler", "tower-http 0.5.2", @@ -667,6 +670,12 @@ dependencies = [ "parking_lot_core", ] +[[package]] +name = "data-encoding" +version = "2.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7a1e2f27636f116493b8b860f5546edb47c8d8f8ea73e1d2a20be88e28d1fea" + [[package]] name = "der" version = "0.7.10" @@ -2505,7 +2514,7 @@ dependencies = [ "serde_json", "sqlx", "strum 0.26.3", - "thiserror", + "thiserror 2.0.18", "time", "tracing", "url", @@ -2602,7 +2611,7 @@ dependencies = [ "proc-macro2", "quote", "syn 2.0.114", - "thiserror", + "thiserror 2.0.18", ] [[package]] @@ -2814,7 +2823,7 @@ checksum = "297f631f50729c8c99b84667867963997ec0b50f32b2a7dbcab828ef0541e8bb" dependencies = [ "num-bigint", "num-traits", - "thiserror", + "thiserror 2.0.18", "time", ] @@ -2912,7 +2921,7 @@ dependencies = [ "serde_json", "sha2", "smallvec", - "thiserror", + "thiserror 2.0.18", "time", "tokio", "tokio-stream", @@ -3000,7 +3009,7 @@ dependencies = [ "smallvec", "sqlx-core", "stringprep", - "thiserror", + "thiserror 2.0.18", "time", "tracing", "uuid", @@ -3043,7 +3052,7 @@ dependencies = [ "smallvec", "sqlx-core", "stringprep", - "thiserror", + "thiserror 2.0.18", "time", "tracing", "uuid", @@ -3070,7 +3079,7 @@ dependencies = [ "serde", "serde_urlencoded", "sqlx-core", - "thiserror", + "thiserror 2.0.18", "time", "tracing", "url", @@ -3221,13 +3230,33 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "thiserror" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +dependencies = [ + "thiserror-impl 1.0.69", +] + [[package]] name = "thiserror" version = "2.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" dependencies = [ - "thiserror-impl", + "thiserror-impl 2.0.18", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.114", ] [[package]] @@ -3381,6 +3410,18 @@ dependencies = [ "tokio", ] +[[package]] +name = "tokio-tungstenite" +version = "0.24.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "edc5f74e248dc973e0dbb7b74c7e0d6fcc301c694ff50049504004ef4d0cdcd9" +dependencies = [ + "futures-util", + "log", + "tokio", + "tungstenite", +] + [[package]] name = "tokio-util" version = "0.7.18" @@ -3554,6 +3595,24 @@ version = "0.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b" +[[package]] +name = "tungstenite" +version = "0.24.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "18e5b8366ee7a95b16d32197d0b2604b43a0be89dc5fac9f8e96ccafbaedda8a" +dependencies = [ + "byteorder", + "bytes", + "data-encoding", + "http", + "httparse", + "log", + "rand", + "sha1", + "thiserror 1.0.69", + "utf-8", +] + [[package]] name = "typenum" version = "1.19.0" @@ -3611,6 +3670,12 @@ dependencies = [ "serde", ] +[[package]] +name = "utf-8" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09cc8ee72d2a9becf2f2febe0205bbed8fc6615b7cb429ad062dc7b7ddd036a9" + [[package]] name = "utf8_iter" version = "1.0.4" diff --git a/Cargo.toml b/Cargo.toml index 6d3bcb2..e7a13dc 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -12,7 +12,7 @@ tokio = { version = "1", features = ["full"] } reqwest = { version = "0.12", features = ["json"] } serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" -axum = "0.7" +axum = { version = "0.7", features = ["ws"] } tower-http = { version = "0.5", features = ["cors", "set-header", "limit"] } sea-orm = { version = "1.1", features = ["sqlx-postgres", "runtime-tokio-rustls", "macros"] } sea-orm-migration = "1.1" diff --git a/frontend/src/main.js b/frontend/src/main.js index 76062aa..a3569ae 100644 --- a/frontend/src/main.js +++ b/frontend/src/main.js @@ -587,17 +587,86 @@ async function checkSession() { return false; } -let isPolling = false; -async function startAutoRefresh() { - setInterval(async () => { - if (isPolling) return; - isPolling = true; +let socket = null; +let reconnectDelay = 1000; + +function connectWebSocket() { + const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:'; + const wsUrl = `${protocol}//${window.location.host}/api/ws`; + + console.log('Connecting to WebSocket:', wsUrl); + socket = new WebSocket(wsUrl); + + socket.onopen = () => { + console.log('WebSocket connected'); + reconnectDelay = 1000; + // Initial fetch to sync state + fetchTasks(); + }; + + socket.onmessage = (event) => { try { - await fetchTasks(); - } finally { - isPolling = false; + const { type, data } = JSON.parse(event.data); + console.log('WebSocket event:', type, data); + + switch (type) { + case 'TaskCreated': + state.tasks.unshift(data); + renderApp(); + showToast('New task created', 'success'); + break; + case 'TaskUpdated': + case 'RunFinished': + const index = state.tasks.findIndex(t => t.id === data.id); + if (index !== -1) { + const wasSelected = state.selectedTaskId === data.id; + state.tasks[index] = data; + if (wasSelected) { + // Update selected run if we were following latest + const wasFollowingLatest = state.selectedRunId === state.tasks[index].runs?.[1]?.id || !state.selectedRunId; + if (wasFollowingLatest && data.runs && data.runs.length > 0) { + state.selectedRunId = data.runs[0].id; + } + } + } else { + state.tasks.unshift(data); + } + renderApp(); + if (type === 'RunFinished') { + showToast(`Task run completed: ${data.goal}`, 'info'); + } + break; + case 'RunStarted': + const taskIndex = state.tasks.findIndex(t => t.id === data.task_id); + if (taskIndex !== -1) { + // We don't have the full task update here, but we can update status + // For simplicity, we just trigger a fetch or wait for RunFinished + // But let's at least show it's running in the UI if selected + if (state.tasks[taskIndex].runs) { + // Prepend a dummy run or just fetch + fetchTasks(); + } + } + showToast(`Task started: ${data.goal}`, 'info'); + break; + } + } catch (e) { + console.error('Error handling WebSocket message:', e); } - }, 3000); + }; + + socket.onclose = () => { + console.log('WebSocket disconnected. Reconnecting...'); + setTimeout(() => { + reconnectDelay = Math.min(reconnectDelay * 1.5, 30000); + connectWebSocket(); + }, reconnectDelay); + }; + + socket.onerror = (error) => { + console.error('WebSocket error:', error); + socket.close(); + }; } async function showLogin() { @@ -661,7 +730,7 @@ async function initializeApp() { appEl.classList.remove('hidden'); loginOverlay.classList.add('hidden'); await fetchTasks(); - startAutoRefresh(); + connectWebSocket(); } else { showLogin(); } diff --git a/frontend/vite.config.js b/frontend/vite.config.js index 7ce721b..a71c6f4 100644 --- a/frontend/vite.config.js +++ b/frontend/vite.config.js @@ -6,6 +6,7 @@ export default defineConfig({ '/api': { target: 'http://localhost:3000', changeOrigin: true, + ws: true, } } } diff --git a/src/domain/tasks.rs b/src/domain/tasks.rs index 3856fb2..6e7962d 100644 --- a/src/domain/tasks.rs +++ b/src/domain/tasks.rs @@ -22,7 +22,7 @@ pub struct UpdateTaskRequest { pub cron: Option, } -#[derive(Serialize)] +#[derive(Serialize, Clone, Debug)] pub struct TaskResponse { pub id: Uuid, pub goal: String, @@ -31,7 +31,7 @@ pub struct TaskResponse { pub runs: Vec, } -#[derive(Serialize)] +#[derive(Serialize, Clone, Debug)] pub struct TaskRunResponse { pub id: Uuid, pub status: String, @@ -40,7 +40,7 @@ pub struct TaskRunResponse { pub created_at: chrono::DateTime, } -#[derive(Serialize)] +#[derive(Serialize, Clone, Debug)] pub struct RecentRunResponse { pub id: Uuid, pub task_id: Uuid, @@ -75,6 +75,18 @@ pub async fn execute_agent_run( .await .map_err(crate::error::AppError::Database)?; + let _ = _scheduler + .tx + .send(crate::server::notifications::WsEvent::RunStarted( + RecentRunResponse { + id: run_id, + task_id, + goal: goal.clone(), + status: "running".to_string(), + created_at: Utc::now().into(), + }, + )); + let mut agent = Agent::new( db.clone(), config.zen_api_key.clone(), @@ -113,7 +125,14 @@ pub async fn execute_agent_run( .await .map_err(crate::error::AppError::Database)?; - get_task_inner(task_id, db).await + let task_response = get_task_inner(task_id, db).await?; + let _ = _scheduler + .tx + .send(crate::server::notifications::WsEvent::RunFinished( + task_response.clone(), + )); + + Ok(task_response) } pub async fn get_task_inner(id: Uuid, db: &DatabaseConnection) -> AppResult { diff --git a/src/scheduler.rs b/src/scheduler.rs index 1bbbb95..9aad91b 100644 --- a/src/scheduler.rs +++ b/src/scheduler.rs @@ -13,12 +13,14 @@ pub struct Scheduler { db: DatabaseConnection, tasks_to_jobs: DashMap, config: Arc, + pub tx: tokio::sync::broadcast::Sender, } impl Scheduler { pub async fn new( db: DatabaseConnection, config: Arc, + tx: tokio::sync::broadcast::Sender, ) -> AppResult { let scheduler = JobScheduler::new() .await @@ -32,6 +34,7 @@ impl Scheduler { db, tasks_to_jobs: DashMap::new(), config, + tx, }) } @@ -43,12 +46,14 @@ impl Scheduler { let db = self.db.clone(); let config = self.config.clone(); + let tx = self.tx.clone(); let job = Job::new_async(cron_expr, move |_uuid, _l| { let db = db.clone(); let config = config.clone(); + let tx = tx.clone(); Box::pin(async move { - if let Err(e) = Self::run_task(db, config, task_id).await { + if let Err(e) = Self::run_task(db, config, tx, task_id).await { tracing::error!("Error in scheduled task {}: {}", task_id, e); } }) @@ -81,6 +86,7 @@ impl Scheduler { async fn run_task( db: DatabaseConnection, config: Arc, + tx: tokio::sync::broadcast::Sender, task_id: Uuid, ) -> AppResult<()> { let task = Task::find_by_id(task_id) @@ -105,6 +111,16 @@ impl Scheduler { use sea_orm::ActiveModelTrait; run.insert(&db).await.map_err(AppError::Database)?; + let _ = tx.send(crate::server::notifications::WsEvent::RunStarted( + crate::domain::tasks::RecentRunResponse { + id: run_id, + task_id, + goal: task.goal.clone(), + status: "running".to_string(), + created_at: chrono::Utc::now().into(), + }, + )); + // Start agent in background let mut agent = Agent::new( db.clone(), @@ -143,6 +159,12 @@ impl Scheduler { e ); } + + if let Ok(task_response) = crate::domain::tasks::get_task_inner(task_id, &db).await { + let _ = tx.send(crate::server::notifications::WsEvent::RunFinished( + task_response, + )); + } }); Ok(()) diff --git a/src/server/mod.rs b/src/server/mod.rs index b40635f..6979d6e 100644 --- a/src/server/mod.rs +++ b/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, pub verifier: Arc, pub authenticator: Arc, + pub tx: tokio::sync::broadcast::Sender, } 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, 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, diff --git a/src/server/notifications.rs b/src/server/notifications.rs new file mode 100644 index 0000000..dcc62a3 --- /dev/null +++ b/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>, +) -> impl IntoResponse { + ws.on_upgrade(|socket| handle_socket(socket, state)) +} + +async fn handle_socket(mut socket: WebSocket, state: Arc) { + 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; + } + } +} diff --git a/src/server/tasks.rs b/src/server/tasks.rs index 9782b9c..5c5da2b 100644 --- a/src/server/tasks.rs +++ b/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(