huskies/server/src/service/gateway/registration.rs

//! Gateway agent registration — pure logic for managing build agents.
//!
//! Contains `JoinedAgent` and functions that validate and manipulate agent
//! state in memory. All persistence (disk I/O) lives in `io.rs`.

use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;

use super::config::ProjectEntry;

/// A build agent that has registered with this gateway.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JoinedAgent {
    /// Unique ID assigned by the gateway on registration.
    pub id: String,
    /// Human-readable label provided by the agent (e.g. `build-agent-abc123`).
    pub label: String,
    /// The agent's CRDT-sync WebSocket address (e.g. `ws://host:3001/crdt-sync`).
    pub address: String,
    /// Unix timestamp when the agent registered.
    pub registered_at: f64,
    /// Unix timestamp of the last heartbeat from this agent.
    #[serde(default)]
    pub last_seen: f64,
    /// Project this agent is assigned to, if any.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub assigned_project: Option<String>,
}

/// Create a new `JoinedAgent` from registration data.
pub fn create_agent(id: String, label: String, address: String, now: f64) -> JoinedAgent {
    JoinedAgent {
        id,
        label,
        address,
        registered_at: now,
        last_seen: now,
        assigned_project: None,
    }
}

/// Remove an agent by ID from the list. Returns `true` if found and removed.
pub fn remove_agent(agents: &mut Vec<JoinedAgent>, id: &str) -> bool {
    let before = agents.len();
    agents.retain(|a| a.id != id);
    agents.len() < before
}

/// Assign (or unassign) an agent to a project.
///
/// Returns the updated agent on success, or an error if the agent or project
/// is not found.
pub fn assign_agent(
    agents: &mut [JoinedAgent],
    id: &str,
    project: Option<String>,
    projects: &BTreeMap<String, ProjectEntry>,
) -> Result<JoinedAgent, super::Error> {
    // Validate project exists if assigning.
    if let Some(ref p) = project
        && !projects.contains_key(p.as_str())
    {
        return Err(super::Error::ProjectNotFound(format!(
            "unknown project '{p}'"
        )));
    }

    match agents.iter_mut().find(|a| a.id == id) {
        None => Err(super::Error::InvalidAgent(format!("agent not found: {id}"))),
        Some(a) => {
            a.assigned_project = project;
            Ok(a.clone())
        }
    }
}

/// Update an agent's last-seen timestamp. Returns `true` if the agent was found.
pub fn heartbeat(agents: &mut [JoinedAgent], id: &str, now: f64) -> bool {
    match agents.iter_mut().find(|a| a.id == id) {
        None => false,
        Some(a) => {
            a.last_seen = now;
            true
        }
    }
}

// ── Tests ────────────────────────────────────────────────────────────────────

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

    #[test]
    fn create_agent_sets_fields() {
        let agent = create_agent("id-1".into(), "lbl".into(), "ws://a".into(), 100.0);
        assert_eq!(agent.id, "id-1");
        assert_eq!(agent.label, "lbl");
        assert_eq!(agent.address, "ws://a");
        assert_eq!(agent.registered_at, 100.0);
        assert_eq!(agent.last_seen, 100.0);
        assert!(agent.assigned_project.is_none());
    }

    #[test]
    fn remove_agent_by_id() {
        let mut agents = vec![
            create_agent("a".into(), "A".into(), "ws://a".into(), 0.0),
            create_agent("b".into(), "B".into(), "ws://b".into(), 0.0),
        ];
        assert!(remove_agent(&mut agents, "a"));
        assert_eq!(agents.len(), 1);
        assert_eq!(agents[0].id, "b");
    }

    #[test]
    fn remove_agent_missing_returns_false() {
        let mut agents = vec![];
        assert!(!remove_agent(&mut agents, "x"));
    }

    #[test]
    fn assign_agent_to_valid_project() {
        let mut projects = BTreeMap::new();
        projects.insert(
            "proj".into(),
            ProjectEntry {
                url: "http://p".into(),
            },
        );
        let mut agents = vec![create_agent("a".into(), "A".into(), "ws://a".into(), 0.0)];
        let result = assign_agent(&mut agents, "a", Some("proj".into()), &projects);
        assert!(result.is_ok());
        assert_eq!(result.unwrap().assigned_project, Some("proj".into()));
    }

    #[test]
    fn assign_agent_to_unknown_project_fails() {
        let projects = BTreeMap::new();
        let mut agents = vec![create_agent("a".into(), "A".into(), "ws://a".into(), 0.0)];
        let result = assign_agent(&mut agents, "a", Some("nope".into()), &projects);
        assert!(result.is_err());
    }

    #[test]
    fn assign_agent_unknown_id_fails() {
        let projects = BTreeMap::new();
        let mut agents: Vec<JoinedAgent> = vec![];
        let result = assign_agent(&mut agents, "x", None, &projects);
        assert!(result.is_err());
    }

    #[test]
    fn heartbeat_updates_last_seen() {
        let mut agents = vec![create_agent("a".into(), "A".into(), "ws://a".into(), 0.0)];
        assert!(heartbeat(&mut agents, "a", 999.0));
        assert_eq!(agents[0].last_seen, 999.0);
    }

    #[test]
    fn heartbeat_unknown_id_returns_false() {
        let mut agents: Vec<JoinedAgent> = vec![];
        assert!(!heartbeat(&mut agents, "x", 1.0));
    }
}