launcher/src/registry.rs

//! Client for the public recipe registry (sm121recipes).
//!
//! Publish: tar.gz the flat bundle produced by `sm121 publish` and POST it.
//! Install: fetch the recipe archive and rebuild the launcher layout under
//! the current directory — `recipe.toml` (+docs) go to `./recipes/<name>/`,
//! image-overlay files (Dockerfile, patches, versions.env, …) to
//! `./docker-images/<name>/`. The registry itself stores the flat bundle.

use std::path::Path;

use anyhow::{Context, Result, bail};
use tracing::info;

pub const DEFAULT_URL: &str = "https://sm121recipes.onlinedb.org";

/// Files that belong in the launcher's `recipes/<name>/` dir: the recipe
/// itself plus the human-facing docs the bundle generates. **Everything else
/// in the bundle is image build context** (Dockerfile, `*.patch`, patcher
/// `*.py`, `versions.env`, `UPSTREAM.md`, `patches/**`, …) and goes to
/// `./docker-images/<name>/`. An allowlist keeps future build-context files
/// from silently mislanding in the recipe dir.
pub(crate) fn is_image_file(entry: &str) -> bool {
    !matches!(entry, "recipe.toml" | "README.md" | "run.sh" | "install.sh")
}

fn client() -> Result<reqwest::Client> {
    reqwest::Client::builder()
        .timeout(std::time::Duration::from_secs(120))
        .build()
        .context("http client")
}

/// Tar.gz the (flat) bundle dir; entry paths are relative to it.
pub(crate) fn tar_gz_dir(dir: &Path) -> Result<Vec<u8>> {
    let mut enc = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default());
    {
        let mut ar = tar::Builder::new(&mut enc);
        for entry in walkdir::WalkDir::new(dir) {
            let e = entry?;
            if e.file_type().is_file() {
                let rel = e.path().strip_prefix(dir)?;
                ar.append_path_with_name(e.path(), rel)?;
            }
        }
        ar.finish()?;
    }
    Ok(enc.finish()?)
}

pub async fn upload(bundle: &Path, name: &str, registry: &str, token: Option<&str>) -> Result<()> {
    let token = token
        .map(str::to_string)
        .or_else(|| std::env::var("SM121RECIPES_TOKEN").ok())
        .with_context(|| "publishing to a registry needs --token <t> or $SM121RECIPES_TOKEN")?;
    let url = format!("{}/api/recipes/{}", registry.trim_end_matches('/'), name);
    let body = tar_gz_dir(bundle)?;
    let resp = client()?
        .post(&url)
        .bearer_auth(token)
        .header(reqwest::header::CONTENT_TYPE, "application/gzip")
        .body(body)
        .send()
        .await
        .context("registry unreachable")?;
    let status = resp.status();
    let text = resp.text().await.unwrap_or_default();
    if !status.is_success() {
        bail!("registry rejected {name}: {status} {text}");
    }
    info!("🌐 published {name} → {url}");
    Ok(())
}

pub async fn handle(cmd: crate::cli::RegistryCommands) -> Result<()> {
    use crate::cli::RegistryCommands;
    match cmd {
        RegistryCommands::List { registry } => list(registry.as_deref()).await,
        RegistryCommands::Show { name, registry } => show(&name, registry.as_deref()).await,
        RegistryCommands::Install {
            name,
            registry,
            force,
        } => install(&name, registry.as_deref(), force).await,
    }
}

pub async fn list(registry: Option<&str>) -> Result<()> {
    let base = registry.unwrap_or(DEFAULT_URL).trim_end_matches('/');
    let items: Vec<serde_json::Value> = client()?
        .get(format!("{base}/api/recipes"))
        .send()
        .await?
        .json()
        .await?;
    if items.is_empty() {
        println!("registry is empty");
        return Ok(());
    }
    for r in items {
        println!(
            "{:<32} {:<8} {:<40} {}",
            r["name"].as_str().unwrap_or("?"),
            r["engine"].as_str().unwrap_or("?"),
            r["model"].as_str().unwrap_or("?"),
            r["title"].as_str().unwrap_or(""),
        );
    }
    Ok(())
}

pub async fn show(name: &str, registry: Option<&str>) -> Result<()> {
    let base = registry.unwrap_or(DEFAULT_URL).trim_end_matches('/');
    let v: serde_json::Value = client()?
        .get(format!("{base}/api/recipes/{name}"))
        .send()
        .await?
        .json()
        .await?;
    print!(
        "{}",
        v["recipe_toml"].as_str().unwrap_or("(no recipe_toml)")
    );
    if let Some(files) = v["files"].as_array() {
        info!(
            "files: {}",
            files
                .iter()
                .filter_map(|f| f.as_str())
                .collect::<Vec<_>>()
                .join(", ")
        );
    }
    Ok(())
}

pub async fn install(name: &str, registry: Option<&str>, force: bool) -> Result<()> {
    validate_name(name)?;
    let base = registry.unwrap_or(DEFAULT_URL).trim_end_matches('/');
    let resp = client()?
        .get(format!("{base}/recipes/{name}/archive.tar.gz"))
        .send()
        .await
        .context("registry unreachable")?;
    if !resp.status().is_success() {
        bail!(
            "recipe '{name}' not found at {base} (HTTP {})",
            resp.status()
        );
    }
    let bytes = resp.bytes().await?;

    let rdir = Path::new("recipes").join(name);
    let idir = Path::new("docker-images").join(name);
    if rdir.exists() && !force {
        bail!("{} exists — overwrite with --force", rdir.display());
    }
    std::fs::create_dir_all(&rdir)?;

    let dec = flate2::read::GzDecoder::new(&bytes[..]);
    let mut ar = tar::Archive::new(dec);
    let mut n = 0usize;
    for entry in ar.entries()? {
        let mut entry = entry?;
        let path = entry.path()?.to_path_buf();
        let Some(rel) = path.to_str().map(str::to_string).filter(|s| !s.is_empty()) else {
            continue;
        };
        if !entry.header().entry_type().is_file() {
            continue;
        }
        if rel.contains("..") || rel.starts_with('/') {
            bail!("unsafe path in registry archive: {rel}");
        }
        let dst = if is_image_file(&rel) { &idir } else { &rdir };
        let out = dst.join(&rel);
        if let Some(p) = out.parent() {
            std::fs::create_dir_all(p)?;
        }
        let mut f = std::fs::File::create(&out)?;
        std::io::copy(&mut entry, &mut f)?;
        n += 1;
    }
    if n == 0 {
        bail!("archive for '{name}' contained no files");
    }
    if idir.exists() {
        info!(
            "📦 {name}: recipe → {}, image overlay → {}",
            rdir.display(),
            idir.display()
        );
    } else {
        info!("📦 {name}: recipe → {}", rdir.display());
    }
    Ok(())
}

fn validate_name(name: &str) -> Result<()> {
    if name.is_empty()
        || !name
            .chars()
            .all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || matches!(c, '-' | '_' | '.'))
    {
        bail!("invalid recipe name '{name}'");
    }
    Ok(())
}

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

    #[test]
    fn flat_bundle_splits_into_launcher_layout() {
        for img in [
            "Dockerfile",
            "modelopt.patch",
            "patches/a.patch",
            "patch_mla.py",
            "versions.env",
            "UPSTREAM.md",
        ] {
            assert!(is_image_file(img), "{img} should be an image file");
        }
        for rec in ["recipe.toml", "README.md", "run.sh", "install.sh"] {
            assert!(!is_image_file(rec), "{rec} should be a recipe file");
        }
    }

    #[test]
    fn names_are_slugs() {
        assert!(validate_name("glm53-flash-nvfp4-lab").is_ok());
        assert!(validate_name("../x").is_err());
        assert!(validate_name("Upper").is_err());
        assert!(validate_name("").is_err());
    }

    #[test]
    fn tar_gz_roundtrips_bundle() {
        let dir = std::env::temp_dir().join(format!("sm121-cli-{}", std::process::id()));
        std::fs::create_dir_all(dir.join("patches")).unwrap();
        std::fs::write(dir.join("recipe.toml"), b"engine = \"vllm\"").unwrap();
        std::fs::write(dir.join("patches/a.patch"), b"diff").unwrap();
        let bytes = tar_gz_dir(&dir).unwrap();
        let mut ar = tar::Archive::new(flate2::read::GzDecoder::new(&bytes[..]));
        let mut names: Vec<String> = ar
            .entries()
            .unwrap()
            .map(|e| e.unwrap().path().unwrap().to_string_lossy().into_owned())
            .collect();
        names.sort();
        assert_eq!(names, ["patches/a.patch", "recipe.toml"]);
        let _ = std::fs::remove_dir_all(&dir);
    }
}