huskies/server/src/chat/transport/matrix/sled_upgrade.rs

//! `upgrade [<project>|all]` gateway chat command — streaming sled binary upgrade.
//!
//! Usage (gateway mode only):
//! - `{bot} upgrade <project>` — upgrade the named sled's binary in-container.
//! - `{bot} upgrade all` — upgrade every registered sled in sequence.
//! - `{bot} upgrade` — list registered projects (shows what can be targeted).
//!
//! The binary comes from the gateway's own artifact store
//! (`~/.huskies/artifacts/`, published by the `release` command) — sleds never
//! download from anywhere but their gateway.  Agents running in a sled are
//! killed by the restart; the pipeline's retry machinery picks the work up
//! again, same as any other agent death.
//!
//! The gateway orchestrates each upgrade in four phases, streaming a marker to
//! the chat room at each step:
//! 1. `[1/4] downloading`    — POSTs to `{sled_url}/api/upgrade`; sled starts download.
//! 2. `[2/4] swapping binary` — gateway received 202; sled atomically renamed the binary.
//! 3. `[3/4] restarting sled` — sled exits cleanly; Docker restarts it with the new binary.
//! 4. `[4/4] reconnected to gateway` — sled's `/health` probe is responding again.
//!
//! After reconnection the gateway polls `/api/version` and verifies the sled's
//! reported git hash matches the published artifact's hash (when a `.hash`
//! sidecar file exists).
//!
//! Concurrent `upgrade` invocations are serialised via a global async mutex so
//! that two simultaneous upgrades cannot interleave their phase markers or race
//! on the sled restart.

use crate::service::gateway::config::ProjectEntry;
use std::collections::BTreeMap;
use std::future::Future;
use std::sync::{Arc, OnceLock};
use std::time::Duration;
use tokio::sync::{Mutex, RwLock};

// ── Serial lock ────────────────────────────────────────────────────────────────

static UPGRADE_LOCK: OnceLock<Mutex<()>> = OnceLock::new();

fn upgrade_lock() -> &'static Mutex<()> {
    UPGRADE_LOCK.get_or_init(|| Mutex::new(()))
}

// ── Command parsing ────────────────────────────────────────────────────────────

/// A parsed `upgrade` command.
#[derive(Debug, PartialEq)]
pub enum UpgradeCommand {
    /// `upgrade <project>` — upgrade the named sled.
    Upgrade {
        /// The project/sled name to upgrade.
        project: String,
    },
    /// `upgrade all` — upgrade every registered sled in sequence.
    UpgradeAll,
    /// `upgrade` with no argument — list available projects.
    ListProjects,
}

/// Parse an `upgrade [<project>]` command from a raw message body.
///
/// Strips the bot mention prefix and checks whether the first word is `upgrade`.
/// Returns `None` when the message is not an upgrade command.
pub fn extract_upgrade_command(
    message: &str,
    bot_name: &str,
    bot_user_id: &str,
) -> Option<UpgradeCommand> {
    let stripped = crate::chat::util::strip_bot_mention(message, bot_name, bot_user_id);
    let trimmed = stripped
        .trim()
        .trim_start_matches(|c: char| !c.is_alphanumeric());

    let (cmd, rest) = match trimmed.split_once(char::is_whitespace) {
        Some((c, r)) => (c, r.trim()),
        None => (trimmed, ""),
    };

    if !cmd.eq_ignore_ascii_case("upgrade") {
        return None;
    }

    if rest.is_empty() {
        return Some(UpgradeCommand::ListProjects);
    }
    let target = rest.split_whitespace().next().unwrap_or(rest);
    if target.eq_ignore_ascii_case("all") {
        Some(UpgradeCommand::UpgradeAll)
    } else {
        Some(UpgradeCommand::Upgrade {
            project: target.to_string(),
        })
    }
}

// ── Handlers ───────────────────────────────────────────────────────────────────

/// List available projects when `upgrade` is invoked without an argument.
///
/// Returns a Markdown string enumerating the registered project names so the
/// user knows which targets are valid for `upgrade <project>`.
pub async fn handle_upgrade_list_projects(
    projects_store: &Arc<RwLock<BTreeMap<String, ProjectEntry>>>,
) -> String {
    let projects = projects_store.read().await;
    if projects.is_empty() {
        return "No projects are currently registered with the gateway.".to_string();
    }
    let names: Vec<&String> = projects.keys().collect();
    let list = names
        .iter()
        .map(|n| format!("- `{n}`"))
        .collect::<Vec<_>>()
        .join("\n");
    format!("Registered projects (use `upgrade <project>` to upgrade one):\n{list}")
}

/// Resolve the artifact source URL and expected git hash for an upgrade.
///
/// The URL points at the gateway's own artifact endpoint via
/// `host.docker.internal` (resolvable from inside sled containers).  The
/// expected hash comes from the `.hash` sidecar written by `release`, when
/// present.  `HUSKIES_GATEWAY_BINARY_URL` overrides the URL (no hash check).
///
/// Returns `Err` with a user-facing message when no artifact has been
/// published yet.
fn resolve_artifact_source(gateway_port: Option<u16>) -> Result<(String, Option<String>), String> {
    if let Ok(url) = std::env::var("HUSKIES_GATEWAY_BINARY_URL") {
        return Ok((url, None));
    }

    let artifact_name = crate::http::SLED_ARTIFACT_NAME;
    let artifact_path = crate::http::artifacts_dir().join(artifact_name);
    if !artifact_path.exists() {
        return Err(format!(
            "No published artifact at `{}`. Run `release` first to build and publish one.",
            artifact_path.display()
        ));
    }

    let expected_hash =
        std::fs::read_to_string(crate::http::artifacts_dir().join(format!("{artifact_name}.hash")))
            .ok()
            .map(|s| s.trim().to_string())
            .filter(|s| !s.is_empty());

    let url = format!(
        "http://host.docker.internal:{}/api/artifacts/{artifact_name}",
        gateway_port.unwrap_or(3000)
    );
    Ok((url, expected_hash))
}

/// Upgrade every registered sled in sequence, streaming per-sled phase markers.
///
/// Returns a summary listing the outcome for each sled.
pub async fn handle_upgrade_all<F, Fut>(
    projects_store: &Arc<RwLock<BTreeMap<String, ProjectEntry>>>,
    gateway_port: Option<u16>,
    send_phase: F,
) -> String
where
    F: Fn(String) -> Fut,
    Fut: Future<Output = ()>,
{
    let names: Vec<String> = {
        let projects = projects_store.read().await;
        projects.keys().cloned().collect()
    };
    if names.is_empty() {
        return "No projects are currently registered with the gateway.".to_string();
    }

    // Fail fast before touching any sled if there is nothing to distribute.
    if let Err(e) = resolve_artifact_source(gateway_port) {
        return e;
    }

    let mut results: Vec<String> = Vec::with_capacity(names.len());
    for name in &names {
        let outcome = handle_sled_upgrade(name, projects_store, gateway_port, |msg| {
            send_phase(format!("**{name}** {msg}"))
        })
        .await;
        results.push(format!("- {name}: {outcome}"));
    }

    format!(
        "Upgrade sweep over {} sled(s) complete:\n{}",
        names.len(),
        results.join("\n")
    )
}

/// Upgrade a named sled by streaming phase markers to the chat room.
///
/// Acquires the global upgrade lock to serialise concurrent invocations.  Each
/// phase is announced by calling `send_phase` before the corresponding work
/// begins.  On any failure, an error message is returned and the previous
/// binary remains active on the sled.
///
/// Agents running in the sled are killed by the restart; the pipeline's
/// retry machinery re-queues their work.
pub async fn handle_sled_upgrade<F, Fut>(
    project: &str,
    projects_store: &Arc<RwLock<BTreeMap<String, ProjectEntry>>>,
    gateway_port: Option<u16>,
    send_phase: F,
) -> String
where
    F: Fn(String) -> Fut,
    Fut: Future<Output = ()>,
{
    // ── Look up project URL ──────────────────────────────────────────────────
    let sled_url = {
        let projects = projects_store.read().await;
        match projects.get(project).and_then(|e| e.url.clone()) {
            Some(u) => u,
            None => {
                let available: Vec<&String> = projects.keys().collect();
                return format!(
                    "Project `{project}` not found. Registered projects: {}",
                    available
                        .iter()
                        .map(|s| s.as_str())
                        .collect::<Vec<_>>()
                        .join(", ")
                );
            }
        }
    };

    // ── Resolve binary source ────────────────────────────────────────────────
    let (source_url, expected_hash) = match resolve_artifact_source(gateway_port) {
        Ok(v) => v,
        Err(e) => return e,
    };

    run_sled_upgrade(project, &sled_url, &source_url, expected_hash, send_phase).await
}

/// Run the four-phase upgrade against a sled whose source URL is already
/// resolved.  Split from [`handle_sled_upgrade`] so tests can drive the wire
/// behaviour without a published artifact on the host.
async fn run_sled_upgrade<F, Fut>(
    project: &str,
    sled_url: &str,
    source_url: &str,
    expected_hash: Option<String>,
    send_phase: F,
) -> String
where
    F: Fn(String) -> Fut,
    Fut: Future<Output = ()>,
{
    let container_name = format!("huskies-{project}");

    // ── Acquire serial lock ──────────────────────────────────────────────────
    let _lock = upgrade_lock().lock().await;

    let client = reqwest::Client::builder()
        .timeout(Duration::from_secs(30))
        .build()
        .unwrap_or_default();

    // ── Phase 1: downloading ─────────────────────────────────────────────────
    send_phase("[1/4] downloading\u{2026}".to_string()).await;

    let upgrade_url = format!("{}/api/upgrade", sled_url.trim_end_matches('/'));
    let body = serde_json::json!({ "source_url": source_url });

    let resp = match client.post(&upgrade_url).json(&body).send().await {
        Ok(r) => r,
        Err(e) => {
            return format!(
                "Upgrade failed at **[1/4] downloading**: could not reach sled at `{upgrade_url}`.\n\
                 Error: {e}\n\n\
                 The previous version remains active."
            );
        }
    };

    if !resp.status().is_success() && resp.status().as_u16() != 202 {
        let status = resp.status();
        let body_text = resp.text().await.unwrap_or_default();
        return format!(
            "Upgrade failed at **[1/4] downloading**: sled returned HTTP {status}.\n\
             Response: {body_text}\n\n\
             The previous version remains active."
        );
    }

    // ── Phase 2: swapping binary ─────────────────────────────────────────────
    // The sled accepted the request (202) and is downloading + atomically
    // replacing the binary in the background.
    send_phase("[2/4] swapping binary\u{2026}".to_string()).await;

    // ── Phase 3: restarting sled ─────────────────────────────────────────────
    // The sled will re-exec momentarily; announce before the health loop.
    send_phase("[3/4] restarting sled\u{2026}".to_string()).await;

    // ── Wait for sled to come back up ────────────────────────────────────────
    let health_url = format!("{}/health", sled_url.trim_end_matches('/'));
    // Give the sled a few seconds to start the download + re-exec before polling.
    tokio::time::sleep(Duration::from_secs(3)).await;

    let reconnected = wait_for_health(&client, &health_url, 120).await;
    if !reconnected {
        return format!(
            "Upgrade failed at **[4/4] reconnected to gateway**: sled at `{sled_url}` did not \
             come back online within 120 seconds after the upgrade was triggered.\n\n\
             Check the container logs: `docker logs huskies-{project}`"
        );
    }

    // ── Phase 4: reconnected ─────────────────────────────────────────────────
    send_phase("[4/4] reconnected to gateway".to_string()).await;

    // ── Verify convergence ───────────────────────────────────────────────────
    match fetch_sled_version(&client, sled_url).await {
        Some((version, git_hash)) => match expected_hash {
            Some(expected) if git_hash == expected => {
                format!("upgraded to v{version} ({git_hash}) — matches published artifact")
            }
            Some(expected) => format!(
                "**upgrade did not converge**: sled reports {git_hash}, published artifact \
                 is {expected}. The sled is healthy but running the wrong binary — check \
                 `docker logs {container_name}`."
            ),
            None => format!("upgraded to v{version} ({git_hash})"),
        },
        None => format!(
            "upgraded and healthy, but `/api/version` is unavailable — the sled is \
             probably still on a pre-version-endpoint binary. Check `docker logs \
             {container_name}`."
        ),
    }
}

// ── Helpers ───────────────────────────────────────────────────────────────────

/// Poll `GET {health_url}` every 3 seconds until it returns 200 or `timeout_secs` elapses.
///
/// Returns `true` when the probe succeeds, `false` on timeout.
async fn wait_for_health(client: &reqwest::Client, health_url: &str, timeout_secs: u64) -> bool {
    let deadline = std::time::Instant::now() + Duration::from_secs(timeout_secs);
    let poll = Duration::from_secs(3);
    loop {
        match client.get(health_url).send().await {
            Ok(r) if r.status().is_success() => return true,
            _ => {}
        }
        if std::time::Instant::now() >= deadline {
            return false;
        }
        tokio::time::sleep(poll).await;
    }
}

/// Fetch `(version, git_hash)` from the sled's `/api/version` endpoint.
///
/// Returns `None` when the endpoint is unreachable or malformed — e.g. a sled
/// still running a binary that predates the endpoint.
async fn fetch_sled_version(client: &reqwest::Client, sled_url: &str) -> Option<(String, String)> {
    let url = format!("{}/api/version", sled_url.trim_end_matches('/'));
    let val: serde_json::Value = client.get(&url).send().await.ok()?.json().await.ok()?;
    let version = val.get("version").and_then(|v| v.as_str())?.to_string();
    let git_hash = val.get("git_hash").and_then(|v| v.as_str())?.to_string();
    Some((version, git_hash))
}

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

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

    // ── extract_upgrade_command ───────────────────────────────────────────────

    #[test]
    fn extract_upgrade_with_project() {
        let cmd = extract_upgrade_command("Timmy upgrade huskies-server", "Timmy", "@timmy:home");
        assert_eq!(
            cmd,
            Some(UpgradeCommand::Upgrade {
                project: "huskies-server".to_string()
            })
        );
    }

    #[test]
    fn extract_upgrade_no_arg_is_list() {
        let cmd = extract_upgrade_command("Timmy upgrade", "Timmy", "@timmy:home");
        assert_eq!(cmd, Some(UpgradeCommand::ListProjects));
    }

    #[test]
    fn extract_upgrade_with_full_user_id() {
        let cmd = extract_upgrade_command("@timmy:home upgrade myapp", "Timmy", "@timmy:home");
        assert_eq!(
            cmd,
            Some(UpgradeCommand::Upgrade {
                project: "myapp".to_string()
            })
        );
    }

    #[test]
    fn extract_non_upgrade_returns_none() {
        let cmd = extract_upgrade_command("Timmy status", "Timmy", "@timmy:home");
        assert!(cmd.is_none());
    }

    #[test]
    fn extract_upgrade_case_insensitive() {
        let cmd = extract_upgrade_command("Timmy UPGRADE alpha", "Timmy", "@timmy:home");
        assert_eq!(
            cmd,
            Some(UpgradeCommand::Upgrade {
                project: "alpha".to_string()
            })
        );
    }

    // ── handle_upgrade_list_projects ─────────────────────────────────────────

    #[tokio::test]
    async fn list_projects_empty_store() {
        let store: Arc<RwLock<BTreeMap<String, ProjectEntry>>> =
            Arc::new(RwLock::new(BTreeMap::new()));
        let msg = handle_upgrade_list_projects(&store).await;
        assert!(
            msg.contains("No projects"),
            "empty store should say no projects: {msg}"
        );
    }

    #[tokio::test]
    async fn list_projects_shows_names() {
        use std::collections::BTreeMap;
        let mut map = BTreeMap::new();
        map.insert(
            "alpha".to_string(),
            ProjectEntry {
                url: Some("http://localhost:3001".into()),
                auth_token: None,
                ssh_port: None,
                host_path: None,
            },
        );
        map.insert(
            "beta".to_string(),
            ProjectEntry {
                url: Some("http://localhost:3002".into()),
                auth_token: None,
                ssh_port: None,
                host_path: None,
            },
        );
        let store = Arc::new(RwLock::new(map));
        let msg = handle_upgrade_list_projects(&store).await;
        assert!(msg.contains("alpha"), "should list alpha: {msg}");
        assert!(msg.contains("beta"), "should list beta: {msg}");
    }

    // ── handle_sled_upgrade validation ───────────────────────────────────────

    #[tokio::test]
    async fn upgrade_unknown_project_returns_error() {
        let store: Arc<RwLock<BTreeMap<String, ProjectEntry>>> =
            Arc::new(RwLock::new(BTreeMap::new()));
        let phases: std::sync::Mutex<Vec<String>> = std::sync::Mutex::new(vec![]);
        let result = handle_sled_upgrade("nonexistent", &store, Some(3000), |msg| {
            phases.lock().unwrap().push(msg);
            async {}
        })
        .await;
        assert!(
            result.contains("not found"),
            "should say not found: {result}"
        );
        // No phase markers should have been emitted before the validation error.
        assert!(
            phases.lock().unwrap().is_empty(),
            "no phases should be emitted for unknown project"
        );
    }

    #[tokio::test]
    async fn upgrade_project_with_no_url_fails_gracefully() {
        let mut map = BTreeMap::new();
        map.insert(
            "myapp".to_string(),
            ProjectEntry {
                url: None,
                auth_token: None,
                ssh_port: None,
                host_path: None,
            },
        );
        let store = Arc::new(RwLock::new(map));
        let result = handle_sled_upgrade("myapp", &store, Some(3000), |_msg| async {}).await;
        assert!(
            result.contains("not found"),
            "project with no URL should say not found: {result}"
        );
    }

    #[test]
    fn extract_upgrade_all() {
        let cmd = extract_upgrade_command("Timmy upgrade all", "Timmy", "@timmy:home");
        assert_eq!(cmd, Some(UpgradeCommand::UpgradeAll));
        let cmd = extract_upgrade_command("@timmy upgrade ALL", "Timmy", "@timmy:home");
        assert_eq!(cmd, Some(UpgradeCommand::UpgradeAll));
    }

    #[tokio::test]
    async fn upgrade_all_empty_store_reports_no_projects() {
        let store: Arc<RwLock<BTreeMap<String, ProjectEntry>>> =
            Arc::new(RwLock::new(BTreeMap::new()));
        let msg = handle_upgrade_all(&store, Some(3000), |_msg| async {}).await;
        assert!(
            msg.contains("No projects"),
            "empty store should say no projects: {msg}"
        );
    }

    #[tokio::test]
    async fn upgrade_unreachable_sled_reports_failure() {
        let phases: std::sync::Mutex<Vec<String>> = std::sync::Mutex::new(vec![]);
        let result = run_sled_upgrade(
            "myapp",
            "http://127.0.0.1:1", // port 1 is never listening
            "http://127.0.0.1:1/api/artifacts/huskies-linux-arm64",
            None,
            |msg| {
                phases.lock().unwrap().push(msg);
                async {}
            },
        )
        .await;
        // Phase 1 marker must have been sent before the failed request.
        let sent = phases.lock().unwrap().clone();
        assert!(
            sent.iter().any(|m| m.contains("[1/4]")),
            "phase 1 marker must be sent: {sent:?}"
        );
        assert!(
            result.contains("downloading") || result.contains("reach"),
            "error should mention the failure: {result}"
        );
        assert!(
            result.contains("previous version"),
            "error should confirm old version is active: {result}"
        );
    }

    // ── wait_for_health ───────────────────────────────────────────────────────

    #[tokio::test]
    async fn wait_for_health_immediate_success() {
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let port = listener.local_addr().unwrap().port();

        let handle = tokio::spawn(async move {
            if let Ok((mut stream, _)) = listener.accept().await {
                use tokio::io::AsyncWriteExt;
                let mut buf = [0u8; 4096];
                let _ = tokio::io::AsyncReadExt::read(&mut stream, &mut buf).await;
                let _ = stream
                    .write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\nok")
                    .await;
            }
        });

        let client = reqwest::Client::new();
        let url = format!("http://127.0.0.1:{port}/health");
        let ok = wait_for_health(&client, &url, 5).await;
        assert!(ok, "should return true when health probe succeeds");
        handle.abort();
    }

    #[tokio::test]
    async fn wait_for_health_timeout() {
        let client = reqwest::Client::builder()
            .timeout(Duration::from_millis(100))
            .build()
            .unwrap();
        // Nothing listening on port 1.
        let ok = wait_for_health(&client, "http://127.0.0.1:1/health", 1).await;
        assert!(!ok, "should return false when health probe never succeeds");
    }
}