feat: add Google OAuth, remote worker system, and file browser
- Google OAuth login with JWT session cookies, per-user project isolation - Remote worker registration via WebSocket, execute_on_worker/list_workers agent tools - File browser UI in workflow view, file upload/download API - Deploy script switched to local build, added tori.euphon.cloud ingress
This commit is contained in:
104
src/ws_worker.rs
Normal file
104
src/ws_worker.rs
Normal file
@@ -0,0 +1,104 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::{
|
||||
extract::{State, WebSocketUpgrade, ws::{Message, WebSocket}},
|
||||
response::Response,
|
||||
routing::get,
|
||||
Router,
|
||||
};
|
||||
use futures::{SinkExt, StreamExt};
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::worker::{WorkerInfo, WorkerManager, WorkerResult};
|
||||
|
||||
pub fn router(mgr: Arc<WorkerManager>) -> Router {
|
||||
Router::new()
|
||||
.route("/", get(ws_handler))
|
||||
.with_state(mgr)
|
||||
}
|
||||
|
||||
async fn ws_handler(
|
||||
ws: WebSocketUpgrade,
|
||||
State(mgr): State<Arc<WorkerManager>>,
|
||||
) -> Response {
|
||||
ws.on_upgrade(move |socket| handle_worker_socket(socket, mgr))
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
#[serde(tag = "type")]
|
||||
enum WorkerMessage {
|
||||
#[serde(rename = "register")]
|
||||
Register { info: WorkerInfo },
|
||||
#[serde(rename = "result")]
|
||||
Result(WorkerResult),
|
||||
}
|
||||
|
||||
async fn handle_worker_socket(socket: WebSocket, mgr: Arc<WorkerManager>) {
|
||||
let (mut sender, mut receiver) = socket.split();
|
||||
|
||||
// First message must be registration
|
||||
let (name, mut job_rx) = loop {
|
||||
match receiver.next().await {
|
||||
Some(Ok(Message::Text(text))) => {
|
||||
match serde_json::from_str::<WorkerMessage>(&text) {
|
||||
Ok(WorkerMessage::Register { info }) => {
|
||||
let name = info.name.clone();
|
||||
let rx = mgr.register(name.clone(), info).await;
|
||||
// Ack
|
||||
let ack = serde_json::json!({ "type": "registered", "name": &name });
|
||||
let _ = sender.send(Message::Text(ack.to_string().into())).await;
|
||||
break (name, rx);
|
||||
}
|
||||
_ => {
|
||||
let _ = sender.send(Message::Text(
|
||||
r#"{"type":"error","message":"First message must be register"}"#.into(),
|
||||
)).await;
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
Some(Ok(Message::Close(_))) | None => return,
|
||||
_ => continue,
|
||||
}
|
||||
};
|
||||
|
||||
// Main loop: forward jobs to worker, receive results
|
||||
let name_clone = name.clone();
|
||||
let mgr_clone = mgr.clone();
|
||||
|
||||
// Task: send jobs from job_rx to the WebSocket
|
||||
let send_task = tokio::spawn(async move {
|
||||
while let Some(req) = job_rx.recv().await {
|
||||
let msg = serde_json::json!({
|
||||
"type": "execute",
|
||||
"job_id": req.job_id,
|
||||
"script": req.script,
|
||||
});
|
||||
if sender.send(Message::Text(msg.to_string().into())).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
// Task: receive results from the WebSocket
|
||||
let recv_task = tokio::spawn(async move {
|
||||
while let Some(Ok(msg)) = receiver.next().await {
|
||||
match msg {
|
||||
Message::Text(text) => {
|
||||
if let Ok(WorkerMessage::Result(result)) = serde_json::from_str(&text) {
|
||||
mgr_clone.report_result(result).await;
|
||||
}
|
||||
}
|
||||
Message::Close(_) => break,
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
tokio::select! {
|
||||
_ = send_task => {},
|
||||
_ = recv_task => {},
|
||||
}
|
||||
|
||||
mgr.unregister(&name_clone).await;
|
||||
}
|
||||
Reference in New Issue
Block a user