Add websocket feature

Author: max-lt

Cargo.lock

@@ -1,11 +1,15 @@
 [package]
 name = "openworkers-logs"
-version = "0.1.3"
+version = "0.1.4"
 edition = "2024"
 
+[features]
+websocket = ["dep:actix-ws"]
+
 [dependencies]
 actix-web = "4.12.1"
 actix-web-lab = "0.24.3"
+actix-ws = { version = "0.3", optional = true }
 tokio = { version = "1.43.0", features = ["full"] }
 sqlx = { version = "0.8.3", features = ["runtime-tokio", "postgres", "uuid", "chrono"] }
 async-nats = "0.45.0"

Cargo.toml

@@ -46,8 +46,8 @@ RUN --mount=type=cache,id=cargo-$TARGETPLATFORM,sharing=locked,target=$CARGO_HOM
     --mount=type=cache,id=cargo-$TARGETPLATFORM,sharing=locked,target=$CARGO_HOME/registry \
     --mount=type=cache,id=cargo-$TARGETPLATFORM,sharing=locked,target=/build/target \
     case "$TARGETPLATFORM" in \
-        "linux/arm64") cargo chef cook --release --target aarch64-unknown-linux-gnu --recipe-path recipe.json ;; \
-        "linux/amd64") cargo chef cook --release --target x86_64-unknown-linux-gnu --recipe-path recipe.json ;; \
+        "linux/arm64") cargo chef cook --release --features websocket --target aarch64-unknown-linux-gnu --recipe-path recipe.json ;; \
+        "linux/amd64") cargo chef cook --release --features websocket --target x86_64-unknown-linux-gnu --recipe-path recipe.json ;; \
         *) echo "Unsupported platform: $TARGETPLATFORM" && exit 1 ;; \
     esac
 
@@ -58,9 +58,9 @@ RUN --mount=type=cache,id=cargo-$TARGETPLATFORM,sharing=locked,target=$CARGO_HOM
     --mount=type=cache,id=cargo-$TARGETPLATFORM,sharing=locked,target=$CARGO_HOME/registry \
     --mount=type=cache,id=cargo-$TARGETPLATFORM,sharing=locked,target=/build/target \
     case "$TARGETPLATFORM" in \
-        "linux/arm64") cargo build --release --target aarch64-unknown-linux-gnu && \
+        "linux/arm64") cargo build --release --features websocket --target aarch64-unknown-linux-gnu && \
                         cp /build/target/aarch64-unknown-linux-gnu/release/openworkers-logs /build/output ;; \
-        "linux/amd64") cargo build --release --target x86_64-unknown-linux-gnu && \
+        "linux/amd64") cargo build --release --features websocket --target x86_64-unknown-linux-gnu && \
                         cp /build/target/x86_64-unknown-linux-gnu/release/openworkers-logs /build/output ;; \
         *) echo "Unsupported platform: $TARGETPLATFORM" && exit 1 ;; \
     esac

Dockerfile.multi

@@ -8,6 +8,8 @@ use futures::StreamExt;
 use uuid::Uuid;
 
 use db::{LogEntry, LogLevel, create_pool, insert_log};
+#[cfg(feature = "websocket")]
+use routes::ws_worker_logs;
 use routes::{AppState, health, stream_worker_logs};
 
 #[actix_web::main]
@@ -93,14 +95,19 @@ async fn main() -> std::io::Result<()> {
     log::info!("Starting HTTP server on 0.0.0.0:{}", port);
 
     HttpServer::new(move || {
-        App::new()
+        let app = App::new()
             .app_data(web::Data::new(AppState {
                 pool: pool.clone(),
                 nats_client: nats_client.clone(),
             }))
             .wrap(middleware::Logger::default())
             .service(health)
-            .service(stream_worker_logs)
+            .service(stream_worker_logs);
+
+        #[cfg(feature = "websocket")]
+        let app = app.service(ws_worker_logs);
+
+        app
     })
     .bind(("0.0.0.0", port))?
     .run()

src/main.rs

@@ -1,5 +1,5 @@
-use base64::engine::general_purpose::STANDARD;
 use base64::engine::Engine;
+use base64::engine::general_purpose::STANDARD;
 
 pub async fn nats_connect() -> async_nats::Client {
     let nats_servers = std::env::var("NATS_SERVERS").expect("NATS_SERVERS must be set");

src/nats.rs

@@ -1,10 +1,15 @@
-use actix_web::{get, web, HttpResponse, Responder};
+use actix_web::{HttpResponse, Responder, get, web};
 use actix_web_lab::sse::{self, Sse};
 use futures::StreamExt;
 use sqlx::PgPool;
 use uuid::Uuid;
 
-use crate::db::{get_logs, LogEntry};
+#[cfg(feature = "websocket")]
+use actix_web::HttpRequest;
+#[cfg(feature = "websocket")]
+use actix_ws::Message;
+
+use crate::db::{LogEntry, get_logs};
 
 pub struct AppState {
     pub pool: PgPool,
@@ -113,3 +118,107 @@ pub async fn stream_worker_logs(
 
     Sse::from_stream(stream).with_keep_alive(std::time::Duration::from_secs(5))
 }
+
+#[cfg(feature = "websocket")]
+#[get("/api/v1/workers/{worker_id}/ws-logs")]
+pub async fn ws_worker_logs(
+    req: HttpRequest,
+    body: web::Payload,
+    worker_id: web::Path<Uuid>,
+    data: web::Data<AppState>,
+) -> Result<HttpResponse, actix_web::Error> {
+    let worker_id = *worker_id;
+
+    let (response, mut session, mut msg_stream) = actix_ws::handle(&req, body)?;
+
+    // Fetch last 10 logs from database
+    let historical_logs = match get_logs(&data.pool, worker_id, 10).await {
+        Ok(logs) => logs,
+        Err(e) => {
+            log::error!("Failed to fetch historical logs: {:?}", e);
+            vec![]
+        }
+    };
+
+    // Subscribe to NATS for this specific worker's logs
+    let subject = format!("{}.console.>", worker_id);
+    let nats_sub = match data.nats_client.subscribe(subject).await {
+        Ok(sub) => sub,
+        Err(e) => {
+            log::error!("Failed to subscribe to NATS: {:?}", e);
+            return Ok(response);
+        }
+    };
+
+    // Spawn WebSocket handler task
+    actix_web::rt::spawn(async move {
+        let mut nats_sub = nats_sub;
+
+        // Send historical logs first (oldest first)
+        for log_entry in historical_logs.into_iter().rev() {
+            let json = serde_json::json!({
+                "date": log_entry.date.timestamp_millis(),
+                "level": format!("{:?}", log_entry.level).to_lowercase(),
+                "message": log_entry.message
+            });
+
+            if session.text(json.to_string()).await.is_err() {
+                return;
+            }
+        }
+
+        // Handle incoming messages and NATS subscription concurrently
+        loop {
+            tokio::select! {
+                // Handle incoming WebSocket messages (ping/pong, close)
+                Some(msg) = msg_stream.next() => {
+                    match msg {
+                        Ok(Message::Ping(bytes)) => {
+                            if session.pong(&bytes).await.is_err() {
+                                break;
+                            }
+                        }
+                        Ok(Message::Close(_)) => {
+                            break;
+                        }
+                        Err(_) => {
+                            break;
+                        }
+                        _ => {}
+                    }
+                }
+
+                // Stream logs from NATS
+                Some(nats_msg) = nats_sub.next() => {
+                    let level_str = nats_msg.subject.split('.').nth(2).unwrap_or("info");
+
+                    let message = match String::from_utf8(nats_msg.payload.to_vec()) {
+                        Ok(m) => m,
+                        Err(_) => continue,
+                    };
+
+                    let json = serde_json::json!({
+                        "date": chrono::Utc::now().timestamp_millis(),
+                        "level": level_str,
+                        "message": message
+                    });
+
+                    if session.text(json.to_string()).await.is_err() {
+                        break;
+                    }
+                }
+
+                // Send ping every 30 seconds to keep connection alive
+                _ = tokio::time::sleep(std::time::Duration::from_secs(30)) => {
+                    if session.ping(b"").await.is_err() {
+                        break;
+                    }
+                }
+            }
+        }
+
+        let _ = session.close(None).await;
+    });
+
+    Ok(response)
+}