launcher/src/runlock.rs

//! Recipe state bookkeeping:
//!
//! **Run ledger** — `recipes/<name>/.lock` (gitignored, machine state):
//! every `sm121 serve` opens a run (unix-epoch start, cluster), every
//! verified `sm121 stop` closes it with its uptime; a serve over a stale
//! open run (crash skipped `stop`) closes it first. `sm121 report` reads
//! who is up right now, the current uptime and lifetime launch counts.
//!
//! **Archived flag** — `archived = true` in the recipe's own toml
//! (round-trip-safe via `toml_edit`): archived recipes leave the report
//! output until un-archived.

use anyhow::Result;
use serde::{Deserialize, Serialize};
use std::path::{Path, PathBuf};

/// `sm121 runlock` subcommands — offline .lock ledger management. `sm121
/// stop` only closes the ledger after verifying every node of the cluster,
/// so a dead node (or a killed serve) leaves the run open forever; these
/// fix that without touching the nodes.
#[derive(clap::Subcommand)]
pub enum RunlockCommands {
    /// Close the open run of a recipe's .lock ledger (offline)
    Close {
        /// Recipe name (its directory under the recipe sources)
        recipe: String,
    },
    /// List every recipe with an open (running) ledger entry
    List,
}

/// One launch: unix-epoch seconds (UTC), open-ended while running.
#[derive(Serialize, Deserialize, Clone)]
pub(crate) struct RunEntry {
    pub(crate) start: i64,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub(crate) end: Option<i64>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub(crate) uptime_secs: Option<i64>,
    /// cluster the run landed on (None for tmux deploys)
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub(crate) cluster: Option<String>,
}

#[derive(Serialize, Deserialize, Default)]
pub(crate) struct Lock {
    /// every launch ever, newest last; capped to the newest [`HISTORY_CAP`]
    #[serde(default)]
    pub(crate) runs: Vec<RunEntry>,
    /// the serve that is up right now (absent when stopped)
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub(crate) running: Option<RunEntry>,
    /// lifetime launch count — survives history capping
    #[serde(default)]
    pub(crate) total: u64,
}

const HISTORY_CAP: usize = 100;

impl Lock {
    pub(crate) fn load(dir: &Path) -> Self {
        std::fs::read_to_string(dir.join(".lock"))
            .ok()
            .and_then(|s| serde_json::from_str(&s).ok())
            .unwrap_or_default()
    }

    fn save(&self, dir: &Path) -> Result<()> {
        std::fs::create_dir_all(dir)?;
        std::fs::write(dir.join(".lock"), serde_json::to_string_pretty(self)?)?;
        Ok(())
    }
}

fn recipe_dir(recipe: &str, sources: &[String]) -> Result<PathBuf> {
    crate::config::recipe_file(recipe, sources)
        .and_then(|p| p.parent().map(Path::to_path_buf))
        .ok_or_else(|| anyhow::anyhow!("recipe dir not found for '{recipe}'"))
}

/// `sm121 serve` went up: close a stale open run (a crash skipped `stop`),
/// then open the new one. `cluster` is None for tmux deploys.
pub(crate) fn mark_start(recipe: &str, sources: &[String], cluster: Option<&str>) -> Result<()> {
    let dir = recipe_dir(recipe, sources)?;
    let mut l = Lock::load(&dir);
    let t = crate::stamp::now_secs();
    if let Some(_stale) = l.running.take()
        && let Some(last) = l.runs.last_mut()
        && last.end.is_none()
    {
        last.end = Some(t);
        last.uptime_secs = Some((t - last.start).max(0));
    }
    let run = RunEntry {
        start: t,
        end: None,
        uptime_secs: None,
        cluster: cluster.map(String::from),
    };
    l.runs.push(run.clone());
    if l.runs.len() > HISTORY_CAP {
        l.runs.drain(..l.runs.len() - HISTORY_CAP);
    }
    l.running = Some(run);
    l.total += 1;
    l.save(&dir)
}

/// `sm121 stop` verified the containers down: close the open run with its
/// uptime. A stop without a tracked run (foreign serve, already-stopped)
/// is a no-op.
pub(crate) fn mark_stop(recipe: &str, sources: &[String]) -> Result<()> {
    let dir = recipe_dir(recipe, sources)?;
    let mut l = Lock::load(&dir);
    if l.running.take().is_none() {
        return Ok(());
    }
    let t = crate::stamp::now_secs();
    if let Some(last) = l.runs.last_mut()
        && last.end.is_none()
    {
        last.end = Some(t);
        last.uptime_secs = Some((t - last.start).max(0));
    }
    l.save(&dir)
}

/// Every recipe with an open (running) ledger entry: (name, start, cluster).
/// Same dir precedence as collect: first dir wins on name collisions.
pub(crate) fn open_runs(config: &crate::config::Config) -> Vec<(String, i64, Option<String>)> {
    let mut out = Vec::new();
    let mut seen = std::collections::HashSet::new();
    for dir in crate::config::recipe_dirs(config) {
        let Ok(entries) = std::fs::read_dir(&dir) else {
            continue;
        };
        for entry in entries.flatten() {
            let path = entry.path();
            let Some(name) = entry.file_name().to_str().map(String::from) else {
                continue;
            };
            if !seen.insert(name.clone()) || !path.join("recipe.toml").is_file() {
                continue;
            }
            if let Some(run) = Lock::load(&path).running {
                out.push((name, run.start, run.cluster));
            }
        }
    }
    out.sort_by(|a, b| a.0.cmp(&b.0));
    out
}

/// `sm121 runlock` — offline ledger management: no node verification
/// (unlike `sm121 stop`, which refuses to close the ledger when it cannot
/// reach every node of the cluster). Used after a node died mid-run or the
/// serve process was killed.
pub(crate) fn run(config: &crate::config::Config, cmd: &RunlockCommands) -> Result<()> {
    match cmd {
        RunlockCommands::Close { recipe } => {
            let open = open_runs(config)
                .into_iter()
                .find(|(n, _, _)| n == recipe)
                .map(|(_, start, cluster)| (start, cluster));
            match open {
                None => println!("no open run for '{recipe}' (nothing to close)"),
                Some((start, cluster)) => {
                    mark_stop(recipe, &config.recipe_sources)?;
                    println!(
                        "closed open run of '{recipe}' (started {} UTC, cluster {})",
                        crate::stamp::fmt_ts(start),
                        cluster.as_deref().unwrap_or("tmux"),
                    );
                }
            }
        }
        RunlockCommands::List => {
            let open = open_runs(config);
            if open.is_empty() {
                println!("no open runs (all ledgers closed)");
                return Ok(());
            }
            for (name, start, cluster) in open {
                println!(
                    "{name}\tstarted {} UTC\tcluster {}",
                    crate::stamp::fmt_ts(start),
                    cluster.as_deref().unwrap_or("tmux"),
                );
            }
        }
    }
    Ok(())
}

/// Flip the recipe's `archived` flag in its own toml (comment-preserving).
pub(crate) fn set_archived(recipe: &str, sources: &[String], archived: bool) -> Result<()> {
    let file = crate::config::recipe_file(recipe, sources)
        .ok_or_else(|| anyhow::anyhow!("recipe file not found for '{recipe}'"))?;
    let text = std::fs::read_to_string(&file)?;
    let mut doc: toml_edit::DocumentMut = text.parse()?;
    doc["archived"] = toml_edit::value(archived);
    std::fs::write(&file, doc.to_string())?;
    Ok(())
}

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

    fn dir(tag: &str) -> PathBuf {
        let d = std::env::temp_dir().join(format!("sm121-runlock-{}-{tag}", std::process::id()));
        std::fs::create_dir_all(d.join("r-lock")).unwrap();
        std::fs::write(
            d.join("r-lock/recipe.toml"),
            "engine = \"vllm\"\nname = \"T\"\n",
        )
        .unwrap();
        d
    }

    fn sources(d: &Path) -> Vec<String> {
        vec![d.to_string_lossy().to_string()]
    }

    #[test]
    fn start_stop_roundtrip_and_counts() {
        let d = dir("rt");
        let src = sources(&d);
        mark_start("r-lock", &src, Some("c1")).unwrap();
        let l = Lock::load(&d.join("r-lock"));
        assert_eq!(l.total, 1);
        assert!(l.running.is_some());
        assert_eq!(l.runs.len(), 1);
        assert_eq!(l.runs[0].cluster.as_deref(), Some("c1"));
        mark_stop("r-lock", &src).unwrap();
        let l = Lock::load(&d.join("r-lock"));
        assert!(l.running.is_none());
        assert_eq!(l.runs[0].end, Some(l.runs[0].start), "instant stop");
        assert_eq!(l.runs[0].uptime_secs, Some(0));
        // stop without a tracked run: no-op
        mark_stop("r-lock", &src).unwrap();
        assert_eq!(Lock::load(&d.join("r-lock")).total, 1);
        std::fs::remove_dir_all(&d).ok();
    }

    #[test]
    fn start_over_stale_run_closes_it() {
        let d = dir("stale");
        let src = sources(&d);
        mark_start("r-lock", &src, None).unwrap();
        // rewind the open run's start so the stale-close uptime is positive
        let p = d.join("r-lock/.lock");
        let mut l = Lock::load(&d.join("r-lock"));
        l.runs[0].start -= 100;
        l.running.as_mut().unwrap().start -= 100;
        std::fs::write(&p, serde_json::to_string(&l).unwrap()).unwrap();
        // a new serve without a stop in between: the stale run closes
        mark_start("r-lock", &src, Some("c2")).unwrap();
        let l = Lock::load(&d.join("r-lock"));
        assert_eq!(l.total, 2, "two launches");
        assert_eq!(l.runs.len(), 2);
        assert_eq!(l.runs[0].uptime_secs, Some(100), "stale closed");
        assert!(l.runs[1].end.is_none(), "new run open");
        assert_eq!(l.running.as_ref().unwrap().cluster.as_deref(), Some("c2"));
        std::fs::remove_dir_all(&d).ok();
    }

    #[test]
    fn history_is_capped_but_total_grows() {
        let d = dir("cap");
        let src = sources(&d);
        for _ in 0..HISTORY_CAP + 5 {
            mark_start("r-lock", &src, None).unwrap();
            mark_stop("r-lock", &src).unwrap();
        }
        let l = Lock::load(&d.join("r-lock"));
        assert_eq!(l.total, HISTORY_CAP as u64 + 5);
        assert_eq!(l.runs.len(), HISTORY_CAP);
        assert!(l.running.is_none());
        std::fs::remove_dir_all(&d).ok();
    }

    #[test]
    fn open_runs_lists_only_open_ledgers() {
        let d = dir("open");
        let src = sources(&d);
        mark_start("r-lock", &src, Some("c1")).unwrap();
        std::fs::create_dir_all(d.join("other")).unwrap();
        std::fs::write(
            d.join("other/recipe.toml"),
            "engine = \"vllm\"\nname = \"O\"\n",
        )
        .unwrap();
        mark_start("other", &src, Some("c1")).unwrap();
        mark_stop("other", &src).unwrap();
        let cfg: crate::config::Config =
            toml::from_str(&format!("recipe_sources = [\"{}\"]\n", d.display())).unwrap();
        let open = open_runs(&cfg);
        let mine = open.iter().find(|(n, _, _)| n == "r-lock");
        assert!(mine.is_some(), "open run not listed: {open:?}");
        assert_eq!(mine.unwrap().2.as_deref(), Some("c1"));
        assert!(
            open.iter().all(|(n, _, _)| n != "other"),
            "closed run must not be listed: {open:?}",
        );
        std::fs::remove_dir_all(&d).ok();
    }

    #[test]
    fn archived_flag_flips_round_trip() {
        let d = dir("arch");
        let src = sources(&d);
        // recipe with comments — toml_edit must preserve them
        let toml = d.join("r-lock/recipe.toml");
        std::fs::write(
            &toml,
            "# keep me\nengine = \"vllm\" # inline\nname = \"T\"\n",
        )
        .unwrap();
        set_archived("r-lock", &src, true).unwrap();
        let after = std::fs::read_to_string(&toml).unwrap();
        assert!(after.contains("archived = true"), "{after}");
        assert!(after.contains("# keep me"), "comment lost: {after}");
        assert!(
            after.contains("engine = \"vllm\" # inline"),
            "inline lost: {after}"
        );
        set_archived("r-lock", &src, false).unwrap();
        assert!(
            std::fs::read_to_string(&toml)
                .unwrap()
                .contains("archived = false")
        );
        std::fs::remove_dir_all(&d).ok();
    }
}