wactorz_interfaces/

ws.rs

//! WebSocket routes for the Wactorz server.
//!
//! Two routes are mounted under the same axum `Router`:
//!
//! - `/ws`   — Python-compatible aggregated-state bridge.
//!   Compatible with `monitor.html` (and any client expecting
//!   `full_snapshot` / `patch` / `delete_agent` JSON messages).
//!
//! - `/mqtt` — Transparent WebSocket proxy to the Mosquitto broker's WS
//!   listener (configurable host/port, default `localhost:9001`).
//!   Compatible with `mqtt.js` / `frontend/dist/index.html`.
//!
//! Together these two routes ensure **any combination** of
//! `python|rust` backend × `monitor.html|frontend/dist/index.html` frontend
//! works without any client-side changes.
//!
//! ## `/ws` message protocol  (mirrors `monitor_server.py`)
//!
//! **Server → browser** on connect:
//! ```json
//! { "type": "full_snapshot", "state": { "agents": [...], "nodes": [...], ... } }
//! ```
//! **Server → browser** on MQTT event:
//! ```json
//! { "type": "patch", "event": { ... }, "state": { ... } }
//! ```
//! **Server → browser** after delete command:
//! ```json
//! { "type": "delete_agent", "agent_id": "...", "state": { ... } }
//! ```
//! **Browser → server** (commands):
//! ```json
//! { "type": "command", "command": "pause|stop|resume|delete", "agent_id": "..." }
//! ```

use std::collections::HashMap;
use std::sync::Arc;
use std::time::{SystemTime, UNIX_EPOCH};

use axum::{
    Router,
    extract::{
        State,
        ws::{Message, WebSocket, WebSocketUpgrade},
    },
    http::HeaderMap,
    response::IntoResponse,
    routing::get,
};
use futures_util::{SinkExt, StreamExt};
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};
use tokio::sync::{Mutex, broadcast, mpsc};

use wactorz_core::{ActorSystem, Message as ActorMessage};
use wactorz_mqtt::MqttClient;

const AGENT_STALE_SECS: f64 = 90.0;
const TERMINAL_AGENT_GRACE_SECS: f64 = 15.0;

// ── Internal MQTT envelope (Rust MQTT loop → WS state aggregator) ─────────────

/// Raw MQTT message forwarded from the broker event loop.
/// Consumed by [`WsBridge::spawn_monitor_task`]; not sent to browser clients.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WsEnvelope {
    pub topic: String,
    pub payload: Value,
}

// ── Monitor state ─────────────────────────────────────────────────────────────

fn now_secs() -> f64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_secs_f64()
}

/// Mirrors the in-memory state maintained by Python's `monitor_server.py`.
#[derive(Debug, Default)]
pub struct MonitorState {
    agents: HashMap<String, Value>,
    nodes: HashMap<String, Value>,
    alerts: Vec<Value>,
    log_feed: Vec<Value>,
    system_health: Value,
}

impl MonitorState {
    fn prune_stale(&mut self) {
        let now = now_secs();
        self.agents.retain(|_, agent| {
            let last_update = agent
                .get("last_update")
                .and_then(|v| v.as_f64())
                .unwrap_or(now);
            let state = agent.get("state").and_then(|v| v.as_str()).unwrap_or("");
            let max_age = match state {
                "stopped" | "failed" => TERMINAL_AGENT_GRACE_SECS,
                _ => AGENT_STALE_SECS,
            };
            now - last_update <= max_age
        });
    }

    /// Serialisable snapshot sent to browser clients.
    pub fn snapshot(&mut self) -> Value {
        self.prune_stale();
        let agents: Vec<Value> = self.agents.values().cloned().collect();
        let nodes: Vec<Value> = self.nodes.values().cloned().collect();
        let total_cost: f64 = self
            .agents
            .values()
            .filter_map(|a| a.get("cost_usd").and_then(|v| v.as_f64()))
            .sum();
        let alert_end = self.alerts.len().min(10);
        let log_end = self.log_feed.len().min(20);
        json!({
            "agents":          agents,
            "nodes":           nodes,
            "alerts":          &self.alerts[..alert_end],
            "log_feed":        &self.log_feed[..log_end],
            "system_health":   self.system_health,
            "total_cost_usd":  (total_cost * 1_000_000.0).round() / 1_000_000.0,
        })
    }

    fn update_agent(&mut self, agent_id: &str, key: &str, data: Value) {
        let short = &agent_id[..agent_id.len().min(8)];
        let entry = self.agents.entry(agent_id.to_string()).or_insert_with(|| {
            json!({
                "agent_id":   agent_id,
                "name":       short,
                "first_seen": now_secs(),
            })
        });
        if let Some(obj) = entry.as_object_mut() {
            obj.insert(key.to_string(), data);
            obj.insert("last_update".to_string(), json!(now_secs()));
        }
    }

    fn add_log(&mut self, entry: Value) {
        self.log_feed.insert(0, entry);
        if self.log_feed.len() > 100 {
            self.log_feed.pop();
        }
    }

    /// Parse one MQTT message and update internal state.
    ///
    /// Returns `Some((event, is_heartbeat))` when something should be
    /// broadcast, or `None` when the topic is not recognised.
    /// `is_heartbeat` suppresses the event from the browser's log feed
    /// (mirrors Python behaviour).
    pub fn parse_topic(&mut self, topic: &str, payload: Value) -> Option<(Value, bool)> {
        let parts: Vec<&str> = topic.split('/').collect();

        // ── system/# ────────────────────────────────────────────────────────
        if parts[0] == "system" && parts.len() >= 2 {
            match parts[1] {
                "health" => {
                    self.system_health = payload.clone();
                }
                "alerts" => {
                    self.alerts.insert(0, payload.clone());
                    if self.alerts.len() > 50 {
                        self.alerts.pop();
                    }
                }
                _ => {}
            }
            return Some((
                json!({
                    "type":    "system",
                    "subtype": parts[1],
                    "data":    payload,
                }),
                false,
            ));
        }

        // ── agents/{id}/{metric} ─────────────────────────────────────────────
        if parts[0] == "agents" && parts.len() >= 3 {
            let agent_id = parts[1];
            let metric = parts[2];

            match metric {
                "status" => {
                    self.update_agent(agent_id, "status", payload.clone());
                    if let Some(obj) = payload.as_object()
                        && let Some(entry) = self.agents.get_mut(agent_id)
                        && let Some(e) = entry.as_object_mut()
                    {
                        if let Some(n) = obj.get("name") {
                            e.insert("name".into(), n.clone());
                        }
                        if let Some(s) = obj.get("state") {
                            e.insert("state".into(), s.clone());
                        }
                    }
                    self.add_log(json!({
                        "type":      "status",
                        "agent_id":  agent_id,
                        "status":    payload,
                        "timestamp": now_secs(),
                    }));
                }
                "heartbeat" => {
                    self.update_agent(agent_id, "heartbeat", payload.clone());
                    if let Some(obj) = payload.as_object() {
                        let short = &agent_id[..agent_id.len().min(8)];
                        let name = obj.get("name").and_then(|v| v.as_str()).unwrap_or(short);
                        if let Some(entry) = self.agents.get_mut(agent_id)
                            && let Some(e) = entry.as_object_mut()
                        {
                            e.insert("name".into(), json!(name));
                            for k in &["cpu", "state"] {
                                if let Some(v) = obj.get(*k) {
                                    e.insert(k.to_string(), v.clone());
                                }
                            }
                            if let Some(v) = obj.get("memory_mb") {
                                e.insert("mem".into(), v.clone());
                            }
                            if let Some(v) = obj.get("task") {
                                e.insert("task".into(), v.clone());
                            }
                        }
                    }
                    // heartbeat → broadcast state update but suppress from log_feed
                    return Some((
                        json!({
                            "type":     "agent",
                            "agent_id": agent_id,
                            "metric":   metric,
                            "data":     payload,
                        }),
                        true,
                    ));
                }
                "metrics" => {
                    self.update_agent(agent_id, "metrics", payload.clone());
                    if let Some(obj) = payload.as_object()
                        && let Some(entry) = self.agents.get_mut(agent_id)
                        && let Some(e) = entry.as_object_mut()
                    {
                        for k in &[
                            "messages_processed",
                            "cost_usd",
                            "input_tokens",
                            "output_tokens",
                        ] {
                            if let Some(v) = obj.get(*k) {
                                e.insert(k.to_string(), v.clone());
                            }
                        }
                    }
                }
                "logs" => {
                    let mut log = json!({
                        "type":      "log",
                        "agent_id":  agent_id,
                        "timestamp": now_secs(),
                    });
                    if let (Some(src), Some(dst)) = (payload.as_object(), log.as_object_mut()) {
                        for (k, v) in src {
                            dst.entry(k.clone()).or_insert(v.clone());
                        }
                    }
                    self.add_log(log);
                }
                "spawned" => {
                    let mut log = json!({
                        "type":      "spawned",
                        "agent_id":  agent_id,
                        "timestamp": now_secs(),
                    });
                    if let (Some(src), Some(dst)) = (payload.as_object(), log.as_object_mut()) {
                        for (k, v) in src {
                            dst.entry(k.clone()).or_insert(v.clone());
                        }
                    }
                    self.add_log(log);
                }
                "completed" => {
                    self.update_agent(agent_id, "last_completed", payload.clone());
                    self.add_log(json!({
                        "type":      "completed",
                        "agent_id":  agent_id,
                        "timestamp": now_secs(),
                    }));
                }
                "alert" => {
                    let short = &agent_id[..agent_id.len().min(8)];
                    let known_name = self
                        .agents
                        .get(agent_id)
                        .and_then(|a| a.get("name"))
                        .and_then(|v| v.as_str())
                        .unwrap_or(short)
                        .to_string();
                    let enriched = if let Some(obj) = payload.as_object() {
                        let mut e = obj.clone();
                        e.insert("agent_id".into(), json!(agent_id));
                        e.entry("name".to_string())
                            .or_insert_with(|| json!(&known_name));
                        Value::Object(e)
                    } else {
                        json!({ "agent_id": agent_id })
                    };
                    let severity = enriched
                        .get("severity")
                        .and_then(|v| v.as_str())
                        .unwrap_or("warning")
                        .to_string();
                    let name = enriched
                        .get("name")
                        .and_then(|v| v.as_str())
                        .unwrap_or(&known_name)
                        .to_string();
                    self.alerts.insert(0, enriched);
                    if self.alerts.len() > 50 {
                        self.alerts.pop();
                    }
                    self.add_log(json!({
                        "type":      "alert",
                        "agent_id":  agent_id,
                        "name":      name,
                        "message":   format!("{name} unresponsive ({severity})"),
                        "timestamp": now_secs(),
                    }));
                }
                _ => {}
            }
            return Some((
                json!({
                    "type":     "agent",
                    "agent_id": agent_id,
                    "metric":   metric,
                    "data":     payload,
                }),
                false,
            ));
        }

        // ── nodes/{name}/heartbeat ───────────────────────────────────────────
        if parts[0] == "nodes" && parts.len() >= 3 && parts[2] == "heartbeat" {
            let node_name = parts[1];
            if let Some(obj) = payload.as_object() {
                self.nodes.insert(
                    node_name.to_string(),
                    json!({
                        "node":      node_name,
                        "agents":    obj.get("agents").cloned().unwrap_or(json!([])),
                        "last_seen": now_secs(),
                        "online":    true,
                        "node_id":   obj.get("node_id").cloned().unwrap_or(json!("")),
                    }),
                );
            }
            return Some((
                json!({
                    "type":      "node",
                    "node_name": node_name,
                    "data":      payload,
                }),
                false,
            ));
        }

        None
    }
}

// ── Shared bridge state ───────────────────────────────────────────────────────

#[derive(Clone)]
pub struct BridgeState {
    /// MQTT → WS broadcast (raw envelopes, consumed by monitor task).
    pub mqtt_tx: broadcast::Sender<WsEnvelope>,
    /// Aggregated monitor state shared across all `/ws` connections.
    pub monitor: Arc<Mutex<MonitorState>>,
    /// Broadcast channel: serialised JSON patches to all `/ws` clients.
    pub monitor_tx: broadcast::Sender<String>,
    /// MQTT client for publishing commands received from the browser.
    pub mqtt_client: Arc<MqttClient>,
    /// Live actor registry used for direct browser -> actor chat routing.
    pub system: ActorSystem,
    /// Mosquitto WebSocket host (for `/mqtt` proxy).
    pub mqtt_host: String,
    /// Mosquitto WebSocket port (for `/mqtt` proxy, default 9001).
    pub mqtt_ws_port: u16,
}

// ── WsBridge ──────────────────────────────────────────────────────────────────

pub struct WsBridge {
    state: BridgeState,
}

impl WsBridge {
    pub fn new(
        mqtt_tx: broadcast::Sender<WsEnvelope>,
        mqtt_client: Arc<MqttClient>,
        system: ActorSystem,
        mqtt_host: String,
        mqtt_ws_port: u16,
    ) -> Self {
        let (monitor_tx, _) = broadcast::channel::<String>(256);
        Self {
            state: BridgeState {
                mqtt_tx,
                monitor: Arc::new(Mutex::new(MonitorState::default())),
                monitor_tx,
                mqtt_client,
                system,
                mqtt_host,
                mqtt_ws_port,
            },
        }
    }

    /// Spawn a background task that:
    ///
    /// 1. Subscribes to `nodes/#` so remote-node heartbeats reach the bridge.
    /// 2. Consumes raw MQTT envelopes from the broadcast channel.
    /// 3. Updates [`MonitorState`].
    /// 4. Broadcasts Python-compatible JSON patches to all `/ws` clients.
    pub fn spawn_monitor_task(&self) {
        // Subscribe to nodes/# so remote node heartbeats are received.
        // agents/# and system/# are subscribed in main.rs; nodes/# is the
        // bridge's own concern.
        let mqtt_for_sub = Arc::clone(&self.state.mqtt_client);
        tokio::spawn(async move {
            if let Err(e) = mqtt_for_sub.subscribe("nodes/#").await {
                tracing::warn!(
                    "[ws-bridge] nodes/# subscribe failed (broker may not be running): {e}"
                );
            } else {
                tracing::info!("[ws-bridge] subscribed to nodes/#");
            }
        });

        let mut rx = self.state.mqtt_tx.subscribe();
        let monitor = Arc::clone(&self.state.monitor);
        let monitor_tx = self.state.monitor_tx.clone();

        tokio::spawn(async move {
            while let Ok(envelope) = rx.recv().await {
                let msg = {
                    let mut st = monitor.lock().await;
                    match st.parse_topic(&envelope.topic, envelope.payload) {
                        None => continue,
                        Some((event, is_heartbeat)) => {
                            let snap = st.snapshot();
                            let log_event = if is_heartbeat { Value::Null } else { event };
                            serde_json::to_string(&json!({
                                "type":  "patch",
                                "event": log_event,
                                "state": snap,
                            }))
                            .unwrap_or_default()
                        }
                    }
                };
                if !msg.is_empty() {
                    let _ = monitor_tx.send(msg);
                }
            }
        });
    }

    /// Build the axum `Router` with `/ws` and `/mqtt` routes.
    pub fn router(&self) -> Router {
        Router::new()
            .route("/ws", get(ws_handler))
            .route("/mqtt", get(mqtt_proxy_handler))
            .with_state(self.state.clone())
    }
}

// ── /ws handler: Python-compatible aggregated state ───────────────────────────

async fn ws_handler(ws: WebSocketUpgrade, State(state): State<BridgeState>) -> impl IntoResponse {
    ws.on_upgrade(move |socket| handle_ws_socket(socket, state))
}

async fn handle_ws_socket(socket: WebSocket, state: BridgeState) {
    let mut monitor_rx = state.monitor_tx.subscribe();
    let (mut ws_send, mut ws_recv) = socket.split();

    // Send a full state snapshot immediately on connect (mirrors Python behaviour)
    let snap_json = {
        let mut st = state.monitor.lock().await;
        serde_json::to_string(&json!({
            "type":  "full_snapshot",
            "state": st.snapshot(),
        }))
        .unwrap_or_default()
    };
    if ws_send.send(Message::Text(snap_json.into())).await.is_err() {
        return;
    }
    let config_json = serde_json::to_string(&json!({
        "type": "config",
        "chat_mode": "direct_ws",
    }))
    .unwrap_or_default();
    if ws_send
        .send(Message::Text(config_json.into()))
        .await
        .is_err()
    {
        return;
    }

    // Per-client direct-reply channel: slash command responses bypass the
    // broadcast and go only to this specific connection.
    let (reply_tx, mut reply_rx) = mpsc::channel::<String>(32);

    // Send task: merges broadcast patches and per-client direct replies.
    let send_task = tokio::spawn(async move {
        loop {
            tokio::select! {
                Ok(json) = monitor_rx.recv() => {
                    if ws_send.send(Message::Text(json.into())).await.is_err() {
                        break;
                    }
                }
                Some(json) = reply_rx.recv() => {
                    if ws_send.send(Message::Text(json.into())).await.is_err() {
                        break;
                    }
                }
                else => break,
            }
        }
    });

    // Handle inbound messages (commands and slash commands from the browser)
    while let Some(Ok(msg)) = ws_recv.next().await {
        match msg {
            Message::Close(_) => break,
            Message::Text(text) => {
                let trimmed = text.trim();
                if trimmed.starts_with('/') {
                    // Slash command — reply only to this client
                    let reply = handle_slash_command(trimmed, &state).await;
                    let _ = reply_tx.send(reply).await;
                } else {
                    handle_browser_message(trimmed, &state).await;
                }
            }
            _ => {}
        }
    }
    send_task.abort();
}

/// Handle a slash command sent by a browser client over `/ws`.
///
/// Mirrors the Python `handle_slash` dispatcher in `monitor_server.py`.
/// Returns a JSON string to send back to that specific client only.
async fn handle_slash_command(text: &str, state: &BridgeState) -> String {
    let parts: Vec<&str> = text.split_whitespace().collect();
    let cmd = parts.first().map(|s| s.to_lowercase()).unwrap_or_default();

    let content = match cmd.as_str() {
        "/help" | "/h" => "Commands:\n\
             \x20 /agents                        list all active agents\n\
             \x20 /nodes                         list remote nodes\n\
             \x20 /help                          show this help\n\n\
             Everything else is forwarded to the main orchestrator."
            .to_string(),

        "/agents" => {
            let st = state.monitor.lock().await;
            if st.agents.is_empty() {
                "No agents running.".to_string()
            } else {
                let mut lines = vec!["Agents:".to_string()];
                let mut names: Vec<&str> = st
                    .agents
                    .values()
                    .filter_map(|a| a.get("name").and_then(|v| v.as_str()))
                    .collect();
                names.sort_unstable();
                for name in names {
                    // Find the full agent entry for this name
                    let entry = st
                        .agents
                        .values()
                        .find(|a| a.get("name").and_then(|v| v.as_str()) == Some(name));
                    let state_str = entry
                        .and_then(|a| a.get("state"))
                        .and_then(|v| v.as_str())
                        .unwrap_or("?");
                    let agent_id = entry
                        .and_then(|a| a.get("agent_id"))
                        .and_then(|v| v.as_str())
                        .unwrap_or("");
                    let id_short = &agent_id[..agent_id.len().min(8)];
                    lines.push(format!("  [{state_str:8}] @{name:<22} {id_short}"));
                }
                lines.join("\n")
            }
        }

        "/nodes" => {
            let st = state.monitor.lock().await;
            let mut lines = vec!["Nodes:".to_string()];
            if st.nodes.is_empty() {
                lines.push("  (no remote nodes)".to_string());
            } else {
                let mut node_names: Vec<&str> = st.nodes.keys().map(|s| s.as_str()).collect();
                node_names.sort_unstable();
                for node_name in node_names {
                    if let Some(nd) = st.nodes.get(node_name) {
                        let online = nd.get("online").and_then(|v| v.as_bool()).unwrap_or(false);
                        let status = if online { "online" } else { "OFFLINE" };
                        let agents: Vec<String> = nd
                            .get("agents")
                            .and_then(|v| v.as_array())
                            .map(|arr| {
                                arr.iter()
                                    .filter_map(|v| v.as_str())
                                    .map(|s| format!("@{s}"))
                                    .collect()
                            })
                            .unwrap_or_default();
                        let agent_list = if agents.is_empty() {
                            "(no agents)".to_string()
                        } else {
                            agents.join(", ")
                        };
                        lines.push(format!("  {node_name:<20} {status:<6}   {agent_list}"));
                    }
                }
            }
            lines.join("\n")
        }

        _ => format!("Unknown command: {cmd}. Type /help for available commands."),
    };

    serde_json::to_string(&json!({
        "type":      "chat",
        "from":      "monitor",
        "content":   content,
        "timestamp": now_secs(),
    }))
    .unwrap_or_default()
}

async fn handle_browser_message(text: &str, state: &BridgeState) {
    let Ok(cmd) = serde_json::from_str::<Value>(text) else {
        return;
    };
    match cmd.get("type").and_then(|v| v.as_str()) {
        Some("command") => handle_browser_command(cmd, state).await,
        Some("chat") => handle_browser_chat(cmd, state).await,
        _ => {}
    }
}

async fn handle_browser_chat(cmd: Value, state: &BridgeState) {
    let Some(content) = cmd.get("content").and_then(|v| v.as_str()) else {
        return;
    };
    let target_name = cmd
        .get("agent_name")
        .and_then(|v| v.as_str())
        .filter(|s| !s.is_empty())
        .unwrap_or("main-actor");

    let Some(entry) = state.system.registry.get_by_name(target_name).await else {
        tracing::warn!("[ws] chat target not found: {target_name}");
        return;
    };
    let msg = ActorMessage::text(
        Some("user".to_string()),
        Some(entry.id.clone()),
        content.to_string(),
    );
    if let Err(err) = state.system.registry.send(&entry.id, msg).await {
        tracing::warn!("[ws] chat route failed for {target_name}: {err}");
    }
}

async fn handle_browser_command(cmd: Value, state: &BridgeState) {
    let Some(command) = cmd.get("command").and_then(|v| v.as_str()) else {
        return;
    };
    let Some(agent_id) = cmd.get("agent_id").and_then(|v| v.as_str()) else {
        return;
    };

    let valid = ["pause", "stop", "resume", "delete"];
    if !valid.contains(&command) {
        tracing::warn!("[ws] Unknown command: {command}");
        return;
    }

    tracing::info!(
        "[ws] {} -> {}",
        command.to_uppercase(),
        &agent_id[..agent_id.len().min(8)]
    );

    // Publish command to MQTT
    let mqtt_payload = json!({
        "command":   command,
        "sender":    "monitor-dashboard",
        "timestamp": now_secs(),
    });
    let topic = format!("agents/{agent_id}/commands");
    if let Err(e) = state.mqtt_client.publish_json(&topic, &mqtt_payload).await {
        tracing::error!("[ws] MQTT publish failed: {e}");
        return;
    }

    // Optimistic state update + broadcast
    let msg = {
        let mut st = state.monitor.lock().await;
        if command == "delete" {
            st.agents.remove(agent_id);
            let snap = st.snapshot();
            serde_json::to_string(&json!({
                "type":     "delete_agent",
                "agent_id": agent_id,
                "state":    snap,
            }))
            .unwrap_or_default()
        } else {
            if let Some(entry) = st.agents.get_mut(agent_id)
                && let Some(e) = entry.as_object_mut()
            {
                let new_state = match command {
                    "stop" => "stopped",
                    "pause" => "paused",
                    "resume" => "running",
                    _ => return,
                };
                e.insert("state".into(), json!(new_state));
            }
            let snap = st.snapshot();
            serde_json::to_string(&json!({
                "type":  "patch",
                "state": snap,
            }))
            .unwrap_or_default()
        }
    };

    if !msg.is_empty() {
        let _ = state.monitor_tx.send(msg);
    }
}

// ── /mqtt handler: transparent proxy to Mosquitto WS ─────────────────────────
//
// The browser's mqtt.js speaks the MQTT binary protocol over WebSocket.
// We forward every frame verbatim to/from Mosquitto's WS listener (port 9001
// by default, or whatever --mqtt-ws-port is set to).
//
// Supports the "mqtt" subprotocol header so mqtt.js is satisfied.

async fn mqtt_proxy_handler(
    ws: WebSocketUpgrade,
    headers: HeaderMap,
    State(state): State<BridgeState>,
) -> impl IntoResponse {
    // Echo back whichever MQTT sub-protocol the client announced
    let proto = headers
        .get("sec-websocket-protocol")
        .and_then(|v| v.to_str().ok())
        .map(|s| s.to_string());

    let ws = ws.protocols(["mqtt", "mqttv3.1"]);
    ws.on_upgrade(move |socket| proxy_to_mosquitto(socket, state, proto))
}

async fn proxy_to_mosquitto(socket: WebSocket, state: BridgeState, proto: Option<String>) {
    use tokio_tungstenite::connect_async;
    use tokio_tungstenite::tungstenite::Message as TMsg;
    use tokio_tungstenite::tungstenite::client::IntoClientRequest;

    let upstream_url = format!("ws://{}:{}/", state.mqtt_host, state.mqtt_ws_port);

    // Build a proper client handshake request, then add the MQTT sub-protocol.
    let request = {
        let mut request = match upstream_url.as_str().into_client_request() {
            Ok(r) => r,
            Err(e) => {
                tracing::warn!("[mqtt-proxy] bad upstream request: {e}");
                return;
            }
        };
        let p = proto.as_deref().unwrap_or("mqtt");
        if let Ok(value) = p.parse() {
            request
                .headers_mut()
                .insert("Sec-WebSocket-Protocol", value);
        }
        request
    };

    let upstream = match connect_async(request).await {
        Ok((stream, _)) => stream,
        Err(e) => {
            tracing::warn!(
                "[mqtt-proxy] upstream connect failed ({}): {e}",
                upstream_url
            );
            return;
        }
    };

    let (mut up_send, mut up_recv) = upstream.split();
    let (mut cl_send, mut cl_recv) = socket.split();

    // upstream → client
    let up_to_cl = tokio::spawn(async move {
        while let Some(Ok(msg)) = up_recv.next().await {
            let out = match msg {
                TMsg::Binary(b) => Message::Binary(b),
                TMsg::Text(t) => Message::Text(t.as_str().into()),
                TMsg::Close(_) => break,
                _ => continue,
            };
            if cl_send.send(out).await.is_err() {
                break;
            }
        }
    });

    // client → upstream
    while let Some(Ok(msg)) = cl_recv.next().await {
        let fwd = match msg {
            Message::Binary(b) => TMsg::Binary(b),
            Message::Text(t) => TMsg::Text(t.as_str().into()),
            Message::Close(_) => break,
            _ => continue,
        };
        if up_send.send(fwd).await.is_err() {
            break;
        }
    }

    up_to_cl.abort();
}