pavel/bot
1
0
Fork 0
Code Issues Pull requests Projects Releases 22 Packages Wiki Activity Actions

websockets

Browse source
This commit is contained in:
pavel 2026-02-11 16:54:44 +01:00
commit f0b53b67ce
9 changed files with 268 additions and 28 deletions
Download patch file Download diff file Expand all files Collapse all files

83
Cargo.lock generated
View file

@ -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"

2
Cargo.toml
View file

@ -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"

89
frontend/src/main.js
View file

@ -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();
}

1
frontend/vite.config.js
View file

@ -6,6 +6,7 @@ export default defineConfig({
'/api': {
target: 'http://localhost:3000',
changeOrigin: true,
ws: true,
}
}
}

27
src/domain/tasks.rs
View file

@ -22,7 +22,7 @@ pub struct UpdateTaskRequest {
pub cron: Option<String>,
}
#[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<TaskRunResponse>,
}
#[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<chrono::FixedOffset>,
}
#[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<TaskResponse> {

24
src/scheduler.rs
View file

@ -13,12 +13,14 @@ pub struct Scheduler {
db: DatabaseConnection,
tasks_to_jobs: DashMap<Uuid, Uuid>,
config: Arc<crate::config::Config>,
pub tx: tokio::sync::broadcast::Sender<crate::server::notifications::WsEvent>,
}
impl Scheduler {
pub async fn new(
db: DatabaseConnection,
config: Arc<crate::config::Config>,
tx: tokio::sync::broadcast::Sender<crate::server::notifications::WsEvent>,
) -> AppResult<Self> {
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<crate::config::Config>,
tx: tokio::sync::broadcast::Sender<crate::server::notifications::WsEvent>,
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(())

8
src/server/mod.rs
View file

@ -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,

46
src/server/notifications.rs Normal file
View file

@ -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;
}
}
}

16
src/server/tasks.rs
View file

@ -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(