clawmates/crates/cm-db/src/repo/nodes.rs

//! Fleet node registry: user-connected local-hardware hosts that run the
//! clawmates-node daemon. The daemon authenticates its control channel with the
//! node `token`, then upserts its host-health snapshot on every heartbeat.

use cm_domain::{NodeId, WorkspaceId};
use sqlx::{PgPool, Row};
use time::OffsetDateTime;

use crate::DbError;

/// Latest host-health snapshot for a node.
#[derive(Debug, Clone)]
pub struct NodeHealth {
    pub cpu_pct: f64,
    pub mem_total: i64,
    pub mem_used: i64,
    pub mem_pressure: f64,
    pub swap_used: i64,
    pub disk_total: i64,
    pub disk_free: i64,
    pub load1: f64,
    pub load5: f64,
    pub load15: f64,
    pub container_count: i32,
}

/// A registered node plus its latest health (if it has reported one).
#[derive(Debug, Clone)]
pub struct NodeRow {
    pub id: NodeId,
    pub name: String,
    pub hostname: Option<String>,
    pub local_ip: Option<String>,
    pub status: String,
    pub agent_version: Option<String>,
    pub tailscale_ip: Option<String>,
    pub last_seen: Option<OffsetDateTime>,
    pub created_at: OffsetDateTime,
    pub health: Option<NodeHealth>,
    /// Latest Beszel metrics scalars (the rich bits beyond basic health).
    pub gpu_pct: Option<f64>,
    pub temp_max: Option<f64>,
}

/// Register a new (pending) node with its control-channel token.
pub async fn create(
    pool: &PgPool,
    workspace_id: WorkspaceId,
    name: &str,
    token: &str,
) -> Result<NodeId, DbError> {
    let id = NodeId::new();
    sqlx::query("INSERT INTO nodes (id, workspace_id, name, status, token) VALUES ($1, $2, $3, 'pending', $4)")
        .bind(id.as_uuid())
        .bind(workspace_id.as_uuid())
        .bind(name)
        .bind(token)
        .execute(pool)
        .await?;
    Ok(id)
}

/// Resolve a control-channel token to its node + workspace (daemon auth).
pub async fn auth(pool: &PgPool, token: &str) -> Result<Option<(NodeId, WorkspaceId)>, DbError> {
    let row = sqlx::query("SELECT id, workspace_id FROM nodes WHERE token = $1")
        .bind(token)
        .fetch_optional(pool)
        .await?;
    Ok(row.map(|r| {
        (
            NodeId::from(r.get::<uuid::Uuid, _>("id")),
            WorkspaceId::from(r.get::<uuid::Uuid, _>("workspace_id")),
        )
    }))
}

const SELECT_WITH_HEALTH: &str = "SELECT n.id, n.name, n.hostname, n.local_ip, n.status, n.agent_version, n.tailscale_ip, n.last_seen, n.created_at,
        h.node_id AS health_node, h.cpu_pct, h.mem_total, h.mem_used, h.mem_pressure, h.swap_used,
        h.disk_total, h.disk_free, h.load1, h.load5, h.load15, h.container_count,
        m.gpu_pct AS m_gpu_pct, m.temp_max AS m_temp_max
   FROM nodes n LEFT JOIN node_health h ON h.node_id = n.id
                LEFT JOIN node_metrics m ON m.node_id = n.id";

/// List a workspace's nodes (oldest first) with their latest health.
pub async fn list(pool: &PgPool, workspace_id: WorkspaceId) -> Result<Vec<NodeRow>, DbError> {
    let rows = sqlx::query(&format!(
        "{SELECT_WITH_HEALTH} WHERE n.workspace_id = $1 ORDER BY n.created_at"
    ))
    .bind(workspace_id.as_uuid())
    .fetch_all(pool)
    .await?;
    Ok(rows.into_iter().map(map_node).collect())
}

/// Fetch one node (workspace-scoped) with its latest health.
pub async fn get(
    pool: &PgPool,
    id: NodeId,
    workspace_id: WorkspaceId,
) -> Result<Option<NodeRow>, DbError> {
    let row = sqlx::query(&format!(
        "{SELECT_WITH_HEALTH} WHERE n.id = $1 AND n.workspace_id = $2"
    ))
    .bind(id.as_uuid())
    .bind(workspace_id.as_uuid())
    .fetch_optional(pool)
    .await?;
    Ok(row.map(map_node))
}

/// Record a heartbeat: mark the node online + refresh its version/tailscale IP,
/// and upsert its latest host-health snapshot.
#[allow(clippy::too_many_arguments)]
pub async fn heartbeat(
    pool: &PgPool,
    id: NodeId,
    agent_version: Option<&str>,
    tailscale_ip: Option<&str>,
    hostname: Option<&str>,
    local_ip: Option<&str>,
    h: &NodeHealth,
) -> Result<(), DbError> {
    sqlx::query(
        // Preserve a rule-/operator-set `draining` state across heartbeats; a node
        // is only un-drained by an explicit set_status.
        "UPDATE nodes SET status = CASE WHEN status = 'draining' THEN 'draining' ELSE 'online' END,
            last_seen = now(),
            agent_version = COALESCE($2, agent_version),
            tailscale_ip = COALESCE($3, tailscale_ip),
            hostname = COALESCE($4, hostname),
            local_ip = COALESCE($5, local_ip)
         WHERE id = $1",
    )
    .bind(id.as_uuid())
    .bind(agent_version)
    .bind(tailscale_ip)
    .bind(hostname)
    .bind(local_ip)
    .execute(pool)
    .await?;
    sqlx::query(
        "INSERT INTO node_health
            (node_id, captured_at, cpu_pct, mem_total, mem_used, mem_pressure, swap_used,
             disk_total, disk_free, load1, load5, load15, container_count)
         VALUES ($1, now(), $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)
         ON CONFLICT (node_id) DO UPDATE SET
            captured_at = now(), cpu_pct = excluded.cpu_pct, mem_total = excluded.mem_total,
            mem_used = excluded.mem_used, mem_pressure = excluded.mem_pressure,
            swap_used = excluded.swap_used, disk_total = excluded.disk_total,
            disk_free = excluded.disk_free, load1 = excluded.load1, load5 = excluded.load5,
            load15 = excluded.load15, container_count = excluded.container_count",
    )
    .bind(id.as_uuid())
    .bind(h.cpu_pct)
    .bind(h.mem_total)
    .bind(h.mem_used)
    .bind(h.mem_pressure)
    .bind(h.swap_used)
    .bind(h.disk_total)
    .bind(h.disk_free)
    .bind(h.load1)
    .bind(h.load5)
    .bind(h.load15)
    .bind(h.container_count)
    .execute(pool)
    .await?;
    Ok(())
}

/// Set a node's status (e.g. 'offline' when its channel drops, 'draining' on
/// deregister).
pub async fn set_status(pool: &PgPool, id: NodeId, status: &str) -> Result<(), DbError> {
    sqlx::query("UPDATE nodes SET status = $2 WHERE id = $1")
        .bind(id.as_uuid())
        .bind(status)
        .execute(pool)
        .await?;
    Ok(())
}

/// Record what a node reports it can host, for placement predicates.
///
/// Replaces rather than merges: the node sends its complete view on every
/// report, so a capability it has *stopped* having (firecracker uninstalled,
/// `/dev/kvm` gone after a reboot into a non-virt kernel) must disappear here
/// too. Merging would let a stale `true` survive forever.
pub async fn set_capabilities(
    pool: &PgPool,
    id: NodeId,
    capabilities: &serde_json::Value,
) -> Result<(), DbError> {
    sqlx::query("UPDATE nodes SET capabilities = $2 WHERE id = $1")
        .bind(id.as_uuid())
        .bind(capabilities)
        .execute(pool)
        .await?;
    Ok(())
}

/// Online nodes that report every one of `required` as `true`.
///
/// The predicate side of placement. Nothing is assumed: a node that has never
/// reported has `capabilities = '{}'`, which fails every requirement — an
/// unqueried node and an incapable node are treated identically, because
/// scheduling work onto a node whose abilities are unknown is how you get a
/// mission that cannot start and does not say why.
pub async fn online_with_capabilities(
    pool: &PgPool,
    workspace_id: uuid::Uuid,
    required: &[&str],
) -> Result<Vec<NodeId>, DbError> {
    let needed: serde_json::Value = required
        .iter()
        .map(|k| ((*k).to_string(), serde_json::Value::Bool(true)))
        .collect::<serde_json::Map<_, _>>()
        .into();
    let rows: Vec<(uuid::Uuid,)> = sqlx::query_as(
        "SELECT id FROM nodes
          WHERE workspace_id = $1 AND status = 'online' AND capabilities @> $2
          ORDER BY last_seen DESC NULLS LAST",
    )
    .bind(workspace_id)
    .bind(&needed)
    .fetch_all(pool)
    .await?;
    Ok(rows.into_iter().map(|(id,)| NodeId::from(id)).collect())
}

/// Online nodes that can host a microVM **and** hold the image `backend` names.
///
/// KVM alone is the wrong predicate. The first real microVM mission was placed
/// on a node reporting `microvm: true` that did not have `rootfs-claude.ext4`;
/// it failed by name rather than booting the wrong image, but whether a mission
/// ran came down to which capable node was listed first.
///
/// `backend = None` means the node's default image, which reports itself as
/// `"default"` — so the requirement is never vacuous. A node running an older
/// daemon has no `rootfs` key at all and matches nothing, which is the same
/// treatment an unqueried node gets for every other capability: unknown is not
/// permission.
///
/// Capability only — this says a node COULD run the backend, not that it has room.
/// Capacity is `cm_api::vm_placement`'s job.
/// microVM phases already pinned to a node but whose VM may not exist yet.
///
/// The other half of "how much is this node committed to". `vm_list` reports
/// BOOTED VMs; between `phase_runner` choosing a node and the guest answering,
/// there is a window of seconds in which a phase is a real 8 GiB claim that no
/// node can report. Two missions launched together both survey a node as empty
/// and both land on it.
///
/// Returns (node, phase_id, iteration) so the caller can build the same
/// deterministic vm id the executor uses and union the two sets by identity
/// rather than adding them — a phase whose VM HAS booted must count once, not
/// twice.
pub async fn pinned_microvm_phases(
    pool: &PgPool,
    workspace_id: uuid::Uuid,
) -> Result<Vec<(NodeId, uuid::Uuid, i32)>, DbError> {
    let rows: Vec<(uuid::Uuid, uuid::Uuid, i32)> = sqlx::query_as(
        "SELECT m.target_node_id, p.id, p.iteration
           FROM mission_phases p
           JOIN missions m ON m.id = p.mission_id
          WHERE m.workspace_id = $1
            AND m.runtime_kind = 'microvm'
            AND m.status = 'running'
            AND m.target_node_id IS NOT NULL
            -- `pending` counts: it is about to become a VM. `completed`/`failed`
            -- do not: their VM is destroyed on every exit path of
            -- `run_phase_in_vm`, so counting them would shrink the fleet by the
            -- number of missions it has ever run.
            AND p.status IN ('pending', 'running')",
    )
    .bind(workspace_id)
    .fetch_all(pool)
    .await?;
    Ok(rows
        .into_iter()
        .map(|(n, p, i)| (NodeId::from(n), p, i))
        .collect())
}

pub async fn online_for_backend(
    pool: &PgPool,
    workspace_id: uuid::Uuid,
    backend: Option<&str>,
) -> Result<Vec<NodeId>, DbError> {
    online_for_backends(pool, workspace_id, &[backend_key(backend).to_string()]).await
}

/// Nodes that can run EVERY one of these backends.
///
/// A composed mission runs its whole graph on one node, and the graph's nodes
/// may each name their own backend — an independent verifier on another
/// provider is the entire point of the roster. Asking only for the mission's
/// backend placed such a mission on a node with `claude` and no
/// `canary-claude`, and the run died at the second graph node with
/// `no rootfs for backend "canary-claude" on this node`. The full harness
/// caught it; nothing before it had a reason to.
pub async fn online_for_backends(
    pool: &PgPool,
    workspace_id: uuid::Uuid,
    backends: &[String],
) -> Result<Vec<NodeId>, DbError> {
    // `@>` on the array asks "does this node's list contain ALL of these" —
    // containment, not intersection, which is exactly the question here and the
    // whole reason the node reports an array rather than a count.
    let rows: Vec<(uuid::Uuid,)> = sqlx::query_as(
        "SELECT id FROM nodes
          WHERE workspace_id = $1 AND status = 'online'
            AND capabilities @> '{\"microvm\": true}'::jsonb
            AND capabilities -> 'rootfs' @> $2::jsonb
          -- Deterministic, NOT `last_seen DESC`. Ranking now happens in
          -- `cm_api::vm_placement`, against real capacity. Ordering by last_seen
          -- and taking `.first()` was the placement algorithm until now: among
          -- healthy nodes all heartbeating every 5s, that is arbitrary — it sent
          -- concurrent missions to whichever node's packet landed most recently,
          -- with no regard for what was already running there.
          ORDER BY id",
    )
    .bind(workspace_id)
    .bind(serde_json::Value::Array(
        backends
            .iter()
            .map(|b| serde_json::Value::String(b.clone()))
            .collect(),
    ))
    .fetch_all(pool)
    .await?;
    Ok(rows.into_iter().map(|(id,)| NodeId::from(id)).collect())
}

/// The name a backend reports itself as in a node's `rootfs` list.
///
/// Must agree with `clawmates-node::microvm::rootfs_for`, which resolves the same
/// three spellings to the default image. If these two drift, placement promises
/// an image the booter cannot find — or refuses one it has.
pub fn backend_key(backend: Option<&str>) -> &str {
    match backend {
        None | Some("") | Some("default") => "default",
        Some(b) => b,
    }
}

/// Mark online nodes whose last heartbeat is older than `secs` as offline.
pub async fn mark_stale_offline(pool: &PgPool, secs: i64) -> Result<(), DbError> {
    sqlx::query(
        "UPDATE nodes SET status = 'offline'
         WHERE status = 'online'
           AND (last_seen IS NULL OR last_seen < now() - ($1 * interval '1 second'))",
    )
    .bind(secs)
    .execute(pool)
    .await?;
    Ok(())
}

/// Remove a node from the registry (workspace-scoped).
pub async fn delete(pool: &PgPool, id: NodeId, workspace_id: WorkspaceId) -> Result<(), DbError> {
    sqlx::query("DELETE FROM nodes WHERE id = $1 AND workspace_id = $2")
        .bind(id.as_uuid())
        .bind(workspace_id.as_uuid())
        .execute(pool)
        .await?;
    Ok(())
}

/// A node's current status (for metrics-aware placement: skip 'draining').
pub async fn status_of(pool: &PgPool, id: NodeId) -> Result<Option<String>, DbError> {
    Ok(
        sqlx::query_scalar::<_, String>("SELECT status FROM nodes WHERE id = $1")
            .bind(id.as_uuid())
            .fetch_optional(pool)
            .await?,
    )
}

fn map_node(r: sqlx::postgres::PgRow) -> NodeRow {
    let health = r
        .get::<Option<uuid::Uuid>, _>("health_node")
        .map(|_| NodeHealth {
            cpu_pct: r.get("cpu_pct"),
            mem_total: r.get("mem_total"),
            mem_used: r.get("mem_used"),
            mem_pressure: r.get("mem_pressure"),
            swap_used: r.get("swap_used"),
            disk_total: r.get("disk_total"),
            disk_free: r.get("disk_free"),
            load1: r.get("load1"),
            load5: r.get("load5"),
            load15: r.get("load15"),
            container_count: r.get("container_count"),
        });
    NodeRow {
        id: NodeId::from(r.get::<uuid::Uuid, _>("id")),
        name: r.get("name"),
        hostname: r.get("hostname"),
        local_ip: r.get("local_ip"),
        status: r.get("status"),
        agent_version: r.get("agent_version"),
        tailscale_ip: r.get("tailscale_ip"),
        last_seen: r.get("last_seen"),
        created_at: r.get("created_at"),
        health,
        gpu_pct: r.get("m_gpu_pct"),
        temp_max: r.get("m_temp_max"),
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /// The three spellings that mean "the node's default image" must all resolve
    /// to the name the node actually advertises for it. A mismatch here makes
    /// placement reject every node for an ordinary mission with no backend set.
    #[test]
    fn the_default_backend_has_one_name() {
        for spelling in [None, Some(""), Some("default")] {
            assert_eq!(backend_key(spelling), "default", "{spelling:?}");
        }
    }

    /// And a named backend is passed through verbatim — it is matched against the
    /// node's list, which is built from the filenames on its disk.
    #[test]
    fn a_named_backend_is_not_rewritten() {
        assert_eq!(backend_key(Some("claude")), "claude");
        assert_eq!(backend_key(Some("agent-terminal")), "agent-terminal");
    }
}

/// Remember a node's idle memory footprint.
///
/// Only ever called with a reading taken while the node had ZERO phase VMs
/// committed — that is the one moment the number is honestly observable.
/// Writing it at any other time would record the VMs as part of the host.
pub async fn set_mem_baseline(pool: &PgPool, node_id: NodeId, mib: i64) -> Result<(), DbError> {
    sqlx::query("UPDATE nodes SET mem_baseline_mib = $2 WHERE id = $1")
        .bind(node_id.as_uuid())
        .bind(mib)
        .execute(pool)
        .await?;
    Ok(())
}