clawmates/crates/cm-api/src/routes/nodes.rs

//! Fleet node registry HTTP/SSE/WS routes: pair a new node, list nodes with live
//! health, stream health updates, run a verification command, deregister, and the
//! daemon's outbound control channel.

use std::convert::Infallible;
use std::sync::Arc;
use std::time::Duration;

use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade};
use axum::extract::{Path, Query, State};
use axum::response::sse::{Event, KeepAlive, Sse};
use axum::response::{IntoResponse, Response};
use axum::Json;
use cm_db::repo::nodes;
use cm_domain::NodeId;
use futures::{SinkExt, StreamExt};
use serde::Deserialize;
use serde_json::{json, Value};
use uuid::Uuid;

use crate::fleet::{run_channel, NodeHub};
use crate::{ApiError, AppState, Authed};

/// An unguessable control-channel token (two time-ordered UUIDs).
fn gen_token() -> String {
    format!("{}{}", Uuid::now_v7().simple(), Uuid::now_v7().simple())
}

fn node_json(n: &nodes::NodeRow) -> Value {
    json!({
        "id": n.id,
        "name": n.name,
        "hostname": n.hostname,
        "localIp": n.local_ip,
        "status": n.status,
        "agentVersion": n.agent_version,
        "tailscaleIp": n.tailscale_ip,
        "lastSeen": n.last_seen.map(|t| t.unix_timestamp()),
        "createdAt": n.created_at.unix_timestamp(),
        "gpuPct": n.gpu_pct,
        "tempMax": n.temp_max,
        "health": n.health.as_ref().map(|h| json!({
            "cpuPct": h.cpu_pct,
            "memTotal": h.mem_total,
            "memUsed": h.mem_used,
            "memPressure": h.mem_pressure,
            "swapUsed": h.swap_used,
            "diskTotal": h.disk_total,
            "diskFree": h.disk_free,
            "load1": h.load1,
            "load5": h.load5,
            "load15": h.load15,
            "containerCount": h.container_count,
        })),
    })
}

#[derive(Deserialize)]
pub struct PairReq {
    pub name: Option<String>,
}

/// `POST /api/nodes/pair` — register a pending node + mint its control-channel
/// token. The frontend builds the `curl … | bash` install command (it knows its
/// own origin); we just return the token + node id.
pub async fn pair(
    State(state): State<AppState>,
    Authed(user): Authed,
    Json(req): Json<PairReq>,
) -> Result<Json<Value>, ApiError> {
    let token = gen_token();
    let name = req
        .name
        .filter(|s| !s.trim().is_empty())
        .unwrap_or_else(|| "New node".to_owned());
    let id = nodes::create(&state.pool, user.workspace_id, &name, &token).await?;
    Ok(Json(json!({ "id": id, "token": token })))
}

/// `GET /api/nodes` — list the workspace's nodes with their latest health.
pub async fn list(
    State(state): State<AppState>,
    Authed(user): Authed,
) -> Result<Json<Value>, ApiError> {
    let rows = nodes::list(&state.pool, user.workspace_id).await?;
    Ok(Json(
        json!({ "nodes": rows.iter().map(node_json).collect::<Vec<_>>() }),
    ))
}

/// `GET /api/nodes/live` — SSE stream of the node list + health (2s poll).
pub async fn live(State(state): State<AppState>, Authed(user): Authed) -> impl IntoResponse {
    let pool = state.pool.clone();
    let ws = user.workspace_id;
    let stream = async_stream::stream! {
        loop {
            if let Ok(rows) = nodes::list(&pool, ws).await {
                let arr: Vec<_> = rows.iter().map(node_json).collect();
                let data = serde_json::to_string(&arr).unwrap_or_else(|_| "[]".to_owned());
                yield Ok::<_, Infallible>(Event::default().event("nodes").data(data));
            }
            tokio::time::sleep(Duration::from_secs(2)).await;
        }
    };
    Sse::new(stream).keep_alive(KeepAlive::default())
}

/// `POST /api/nodes/{id}/exec-test` — run a verification command on the node and
/// return its output (wizard step 3: "we can run commands on your behalf").
pub async fn exec_test(
    State(state): State<AppState>,
    Authed(user): Authed,
    Path(id): Path<Uuid>,
) -> Result<Json<Value>, ApiError> {
    let node_id = NodeId::from(id);
    // Scope: the node must belong to the caller's workspace.
    let node = nodes::get(&state.pool, node_id, user.workspace_id)
        .await?
        .ok_or(ApiError::NotFound)?;
    match state.node_hub.verify(node_id).await {
        Ok(out) => Ok(Json(
            json!({ "ok": out.ok, "output": out.output, "node": node.name }),
        )),
        Err(e) => Ok(Json(json!({ "ok": false, "output": e, "node": node.name }))),
    }
}

/// `POST /api/nodes/{id}/sandbox-check` — run a hardened throwaway container on
/// the node to confirm it can host agent workloads.
pub async fn sandbox_check(
    State(state): State<AppState>,
    Authed(user): Authed,
    Path(id): Path<Uuid>,
) -> Result<Json<Value>, ApiError> {
    let node_id = NodeId::from(id);
    nodes::get(&state.pool, node_id, user.workspace_id)
        .await?
        .ok_or(ApiError::NotFound)?;
    match state.node_hub.sandbox_check(node_id).await {
        Ok(out) => Ok(Json(json!({ "ok": out.ok, "output": out.output }))),
        Err(e) => Ok(Json(json!({ "ok": false, "output": e }))),
    }
}

/// `DELETE /api/nodes/{id}` — deregister a node (workspace-scoped).
pub async fn remove(
    State(state): State<AppState>,
    Authed(user): Authed,
    Path(id): Path<Uuid>,
) -> Result<Json<Value>, ApiError> {
    nodes::delete(&state.pool, NodeId::from(id), user.workspace_id).await?;
    Ok(Json(json!({ "ok": true })))
}

#[derive(Deserialize)]
pub struct AgentQuery {
    pub token: String,
}

/// `GET /api/nodes/agent?token=…` — the daemon's outbound control channel. The
/// WS handshake can't carry a bearer header, so the daemon authenticates with
/// its node token in the query string (like the Terminal WS ticket).
pub async fn agent_ws(
    State(state): State<AppState>,
    Query(q): Query<AgentQuery>,
    upgrade: WebSocketUpgrade,
) -> Response {
    let auth = nodes::auth(&state.pool, &q.token).await.ok().flatten();
    let Some((node_id, _workspace_id)) = auth else {
        return ApiError::Unauthorized.into_response();
    };
    let pool = state.pool.clone();
    let hub = state.node_hub.clone();
    upgrade.on_upgrade(move |socket| run_channel(pool, hub, node_id, socket))
}

/// `GET /api/nodes/{id}/tools` — installed dev-tool versions for a node, the
/// latest upstream version, and an update flag (Phase 1: read-only). `glm` is
/// derived from `claude` (it runs Claude Code with a z.ai config).
pub async fn tools(
    State(state): State<AppState>,
    Authed(user): Authed,
    Path(id): Path<Uuid>,
) -> Result<Json<Value>, ApiError> {
    let node_id = NodeId::from(id);
    nodes::get(&state.pool, node_id, user.workspace_id)
        .await?
        .ok_or(ApiError::NotFound)?;
    let installed: std::collections::HashMap<String, String> =
        cm_db::repo::node_tools::list(&state.pool, node_id)
            .await
            .unwrap_or_default()
            .into_iter()
            .collect();
    let latest = cm_db::repo::node_tools::all_latest(&state.pool)
        .await
        .unwrap_or_default();

    // (probe key, display name, ui key) in display order; `glm` mirrors `claude`.
    let order = [
        ("docker", "Docker", "docker"),
        ("rust", "Rust", "rust"),
        ("claude", "Claude Code", "claude"),
        ("kimi-cli", "Kimi", "kimi"),
        ("glm", "GLM (Claude Code)", "glm"),
        ("ollama", "Ollama", "ollama"),
    ];
    let mut out: Vec<Value> = Vec::new();
    for (probe, name, key) in order {
        let src = if probe == "glm" { "claude" } else { probe };
        if let Some(inst) = installed.get(src) {
            let lat = latest.get(src).cloned();
            let update = lat.as_ref().map(|l| l != inst).unwrap_or(false);
            out.push(json!({
                "name": name, "key": key, "installed": inst,
                "latest": lat, "updateAvailable": update,
            }));
        }
    }
    Ok(Json(json!({ "tools": out })))
}

/// `POST /api/nodes/{id}/tools/{tool}/update` — run the daemon's fixed update for a
/// tool (claude/glm/kimi/ollama), then it re-probes so the version refreshes.
pub async fn tool_update(
    State(state): State<AppState>,
    Authed(user): Authed,
    Path((id, tool)): Path<(Uuid, String)>,
) -> Result<Json<Value>, ApiError> {
    let node_id = NodeId::from(id);
    nodes::get(&state.pool, node_id, user.workspace_id)
        .await?
        .ok_or(ApiError::NotFound)?;
    if !["claude", "glm", "kimi", "ollama", "rust"].contains(&tool.as_str()) {
        return Ok(Json(json!({ "ok": false, "output": "tool not updatable" })));
    }
    // ~180s: tool updates (npm/uv/brew/installer) legitimately run long.
    match state
        .node_hub
        .call_timeout(node_id, "tool_update", json!({ "tool": tool }), 180)
        .await
    {
        Ok(out) => Ok(Json(json!({ "ok": out.ok, "output": out.output }))),
        Err(e) => Ok(Json(json!({ "ok": false, "output": e }))),
    }
}

/// `POST /api/nodes/{id}/terminal/ticket` — mint a single-use terminal ticket
/// (the browser WS handshake can't carry a bearer header).
pub async fn terminal_ticket(
    State(state): State<AppState>,
    Authed(user): Authed,
    Path(id): Path<Uuid>,
) -> Result<Json<Value>, ApiError> {
    let node_id = NodeId::from(id);
    nodes::get(&state.pool, node_id, user.workspace_id)
        .await?
        .ok_or(ApiError::NotFound)?;
    if !state.node_hub.is_online(node_id).await {
        return Ok(Json(json!({ "error": "node is offline" })));
    }
    Ok(Json(
        json!({ "ticket": state.node_hub.mint_ticket(node_id).await }),
    ))
}

/// `GET /api/nodes/{id}/terminal/ws?ticket=…` — bridge a browser xterm to a host
/// shell on the node (PTY proxied over the daemon's control channel).
pub async fn terminal_ws(
    State(state): State<AppState>,
    Path(id): Path<Uuid>,
    Query(q): Query<AgentQuery>,
    upgrade: WebSocketUpgrade,
) -> Response {
    let node_id = NodeId::from(id);
    match state.node_hub.redeem_ticket(&q.token).await {
        Some(t) if t == node_id => {}
        _ => return ApiError::Unauthorized.into_response(),
    }
    let hub = state.node_hub.clone();
    upgrade.on_upgrade(move |socket| bridge_terminal(hub, node_id, socket))
}

/// Browser→server control frames over the terminal WS: a resize, the
/// `fallback` request to open the WS-relay PTY, or WebRTC signaling (offer/ICE).
#[derive(Deserialize)]
struct TermCtrl {
    #[serde(rename = "type")]
    kind: String,
    cols: Option<u16>,
    rows: Option<u16>,
    sdp: Option<String>,
    candidate: Option<String>,
    sdp_mid: Option<String>,
    sdp_mline_index: Option<u16>,
}

/// The terminal WS is BOTH the WebRTC signaling channel and the fallback data
/// path. The browser tries a direct DataChannel first (offer/ICE relayed here);
/// if that fails it sends `{type:"fallback"}` and we open the WS-relay PTY.
async fn bridge_terminal(hub: Arc<NodeHub>, node_id: NodeId, socket: WebSocket) {
    let Some((sid, mut pty_rx, mut sig_rx)) = hub.open_session(node_id).await else {
        return;
    };
    let (mut ws_tx, mut ws_rx) = socket.split();
    let to_browser = async {
        loop {
            tokio::select! {
                bytes = pty_rx.recv() => match bytes {
                    Some(b) => {
                        if ws_tx.send(Message::Binary(b.into())).await.is_err() { break; }
                    }
                    None => break,
                },
                text = sig_rx.recv() => match text {
                    Some(t) => {
                        if ws_tx.send(Message::Text(t.into())).await.is_err() { break; }
                    }
                    None => break,
                },
            }
        }
    };
    let to_node = async {
        while let Some(Ok(msg)) = ws_rx.next().await {
            match msg {
                // Binary = keystrokes over the WS fallback path only (the direct
                // DataChannel carries its own input).
                Message::Binary(b) => hub.terminal_input(node_id, sid, b.as_ref()).await,
                Message::Text(t) => {
                    let Ok(c) = serde_json::from_str::<TermCtrl>(t.as_str()) else {
                        continue;
                    };
                    let cols = c.cols.unwrap_or(80);
                    let rows = c.rows.unwrap_or(24);
                    match c.kind.as_str() {
                        "resize" => hub.terminal_resize(node_id, sid, cols, rows).await,
                        // Host shell (no container) — the Infra node terminal.
                        "fallback" => hub.open_pty(node_id, sid, cols, rows, None, None).await,
                        "webrtc_offer" => {
                            hub.webrtc_offer(
                                node_id,
                                sid,
                                c.sdp.as_deref().unwrap_or(""),
                                None,
                                None,
                            )
                            .await
                        }
                        "webrtc_ice" => {
                            hub.webrtc_ice(
                                node_id,
                                sid,
                                c.candidate.as_deref().unwrap_or(""),
                                c.sdp_mid.as_deref(),
                                c.sdp_mline_index,
                            )
                            .await
                        }
                        "webrtc_close" => hub.webrtc_close(node_id, sid).await,
                        _ => {}
                    }
                }
                Message::Close(_) => break,
                _ => {}
            }
        }
    };
    tokio::select! {
        _ = to_browser => {},
        _ = to_node => {},
    }
    hub.terminal_close(node_id, sid).await;
}