Initial implementation: relay server + common protocol + client stub

crates/server/src/ws_handler.rs

@@ -0,0 +1,99 @@
+use axum::{
+    extract::{State, WebSocketUpgrade},
+    response::IntoResponse,
+};
+use axum::extract::ws::{Message, WebSocket};
+use futures_util::{SinkExt, StreamExt};
+use tokio::sync::mpsc;
+use uuid::Uuid;
+use tracing::{debug, error, info, warn};
+
+use helios_common::protocol::{ClientMessage, ServerMessage};
+use crate::{AppState, session::Session};
+
+pub async fn ws_upgrade(
+    ws: WebSocketUpgrade,
+    State(state): State<AppState>,
+) -> impl IntoResponse {
+    ws.on_upgrade(move |socket| handle_socket(socket, state))
+}
+
+async fn handle_socket(socket: WebSocket, state: AppState) {
+    let session_id = Uuid::new_v4();
+    let (cmd_tx, mut cmd_rx) = mpsc::channel::<ServerMessage>(64);
+
+    // Register session (label filled in on Hello)
+    let session = Session {
+        id: session_id,
+        label: None,
+        cmd_tx,
+    };
+    state.sessions.insert(session);
+    info!("Client connected: session={session_id}");
+
+    let (mut ws_tx, mut ws_rx) = socket.split();
+
+    // Spawn task: forward server commands → WS
+    let send_task = tokio::spawn(async move {
+        while let Some(msg) = cmd_rx.recv().await {
+            match serde_json::to_string(&msg) {
+                Ok(json) => {
+                    if let Err(e) = ws_tx.send(Message::Text(json.into())).await {
+                        error!("WS send error for session={session_id}: {e}");
+                        break;
+                    }
+                }
+                Err(e) => {
+                    error!("Serialization error for session={session_id}: {e}");
+                }
+            }
+        }
+    });
+
+    // Main loop: receive client messages
+    while let Some(result) = ws_rx.next().await {
+        match result {
+            Ok(Message::Text(text)) => {
+                match serde_json::from_str::<ClientMessage>(&text) {
+                    Ok(msg) => handle_client_message(session_id, msg, &state).await,
+                    Err(e) => {
+                        warn!("Invalid JSON from session={session_id}: {e} | raw={text}");
+                    }
+                }
+            }
+            Ok(Message::Close(_)) => {
+                info!("Client disconnected gracefully: session={session_id}");
+                break;
+            }
+            Ok(Message::Ping(_)) | Ok(Message::Pong(_)) | Ok(Message::Binary(_)) => {}
+            Err(e) => {
+                error!("WS receive error for session={session_id}: {e}");
+                break;
+            }
+        }
+    }
+
+    send_task.abort();
+    state.sessions.remove(&session_id);
+    info!("Session cleaned up: session={session_id}");
+}
+
+async fn handle_client_message(session_id: Uuid, msg: ClientMessage, state: &AppState) {
+    match &msg {
+        ClientMessage::Hello { label } => {
+            if let Some(lbl) = label {
+                state.sessions.set_label(&session_id, lbl.clone());
+            }
+            debug!("Hello from session={session_id}, label={label:?}");
+        }
+        ClientMessage::ScreenshotResponse { request_id, .. }
+        | ClientMessage::ExecResponse { request_id, .. }
+        | ClientMessage::Ack { request_id }
+        | ClientMessage::Error { request_id, .. } => {
+            let rid = *request_id;
+            if !state.sessions.resolve_pending(rid, msg) {
+                warn!("No pending request for request_id={rid} (session={session_id})");
+            }
+        }
+    }
+}