mecha10_cli/services/

rl_bundle.rs

//! RL project bundle packaging + upload (LAB-2768)
//!
//! `mecha10 build rl` packages a project's `configs/rl/` (job/task definitions - `jobs/*.yml`,
//! `tasks/*/task.json`) and `configs/simulation/{models,environments}` (the assets those jobs
//! reference) into a single tar.gz and uploads it, scoped by the project's registered
//! `robot_id`, via the endpoint LAB-2767 adds to mecha10-api. `mecha10 run rl` (on the trainer
//! host) downloads and unpacks this same bundle before training starts - see "RL: Vision and
//! Phased Roadmap" doc.
//!
//! This bundle is deliberately separate from `mecha10 config push`'s `configs/` archive
//! (`mecha10_robot_config::build_configs_archive`): that one syncs the *whole* `configs/`
//! directory to the robot-config KV service for the robot/launcher to consume at runtime, while
//! this one is a purpose-built blob (Minio-backed via LAB-2767, not the KV store) for exactly
//! what a remote/Kestra-dispatched training run needs - replacing the old `simulation-assets/*`
//! KV convention `mecha10 sim push`/`pull` used for this purpose (see LAB-2767's context: that
//! mechanism is real, working code, but nothing downstream ever reads what it pushes today).
//!
//! Only `models`/`environments` ride along from `configs/simulation/` - not `images` (the flat
//! billboard-image pool `environment.json`/`task.json` objects reference by path elsewhere,
//! already synced to the robot separately via `mecha10 config push`) - matching this ticket's
//! (LAB-2768) explicit requirement.

use anyhow::{Context, Result};
use flate2::read::GzDecoder;
use flate2::write::GzEncoder;
use flate2::Compression;
use std::path::Path;
use tar::{Archive, Builder};

/// `configs/simulation/` subdirectories folded into the RL bundle - deliberately excludes
/// `images` (the flat billboard-image pool synced separately via `mecha10 config push`).
const SIMULATION_ASSET_DIRS: [&str; 2] = ["models", "environments"];

/// Build the RL project bundle archive (tar.gz) for `project_root`, or `Ok(None)` if the project
/// has neither a `configs/rl/` directory nor any `configs/simulation/{models,environments}`
/// content to package - callers should skip the upload step entirely in that case rather than
/// upload an empty archive.
///
/// Archive layout mirrors the source tree exactly (`configs/rl/...`,
/// `configs/simulation/models/...`, `configs/simulation/environments/...`) so unpacking it on the
/// trainer host (LAB-2767's download side, `mecha10 run rl`) drops files back into the same
/// relative paths a local project uses - the same "archive root mirrors project layout"
/// convention `mecha10_robot_config::build_configs_archive` and
/// `scripts/release/build/package-templates.sh` both already use.
pub fn build_rl_bundle_archive(project_root: &Path) -> Result<Option<Vec<u8>>> {
    let rl_dir = project_root.join("configs").join("rl");
    let simulation_dir = project_root.join("configs").join("simulation");

    let asset_dirs: Vec<_> = SIMULATION_ASSET_DIRS
        .iter()
        .map(|name| (*name, simulation_dir.join(name)))
        .filter(|(_, dir)| dir.exists())
        .collect();

    if !rl_dir.exists() && asset_dirs.is_empty() {
        return Ok(None);
    }

    let mut bytes = Vec::new();
    let encoder = GzEncoder::new(&mut bytes, Compression::default());
    let mut builder = Builder::new(encoder);

    if rl_dir.exists() {
        builder
            .append_dir_all("configs/rl", &rl_dir)
            .context("Failed to archive configs/rl")?;
    }

    for (name, dir) in asset_dirs {
        builder
            .append_dir_all(format!("configs/simulation/{name}"), &dir)
            .with_context(|| format!("Failed to archive configs/simulation/{name}"))?;
    }

    let encoder = builder.into_inner().context("Failed to finish RL bundle archive")?;
    encoder.finish().context("Failed to finish RL bundle archive")?;

    Ok(Some(bytes))
}

/// Tarball filename for an RL project bundle - matches the `<name>-v<version>-<suffix>.tar.gz`
/// convention `scripts/release/build/package-templates.sh` (`mecha10-v{VERSION}-templates.tar.gz`)
/// and `handlers::upload::handle_upload` (`{name}-v{version}-{platform}-{arch}.tar.gz`) already
/// use for versioned release archives.
pub fn rl_bundle_tarball_name(project_name: &str, version: &str) -> String {
    format!("{project_name}-v{version}-rl-bundle.tar.gz")
}

/// URL for LAB-2767's robot-scoped RL project bundle upload endpoint. Isolated in its own
/// function - matching `services::publish_status::build_versions_url`'s pattern - so if LAB-2767
/// lands with a different path than assumed here, only this one function (and the tests pinned
/// to it) need updating.
///
/// Assumed shape mirrors the existing robot-scoped, Minio-backed archive endpoint
/// (`RobotConfigClient::upload_configs_archive`'s `PUT /api/robots/{robot_id}/configs/archive`):
/// `PUT {control_plane_url}/api/robots/{robot_id}/rl-bundle`.
fn build_rl_bundle_url(control_plane_url: &str, robot_id: &str) -> String {
    format!(
        "{}/api/robots/{}/rl-bundle",
        control_plane_url.trim_end_matches('/'),
        robot_id
    )
}

/// Upload a pre-built RL project bundle archive for `robot_id`, authenticated the same way every
/// other mecha10 CLI command talks to mecha10-api (`Authorization: Bearer <api_key>` +
/// `X-User-ID`, see `handlers::upload::upload_to_registry`/`services::publish_status`).
pub async fn upload_rl_bundle(
    control_plane_url: &str,
    robot_id: &str,
    archive_bytes: Vec<u8>,
    tarball_name: &str,
    api_key: &str,
    user_id: &str,
) -> Result<()> {
    let url = build_rl_bundle_url(control_plane_url, robot_id);

    let part = reqwest::multipart::Part::bytes(archive_bytes)
        .file_name(tarball_name.to_string())
        .mime_str("application/gzip")
        .context("Invalid mime type")?;
    let form = reqwest::multipart::Form::new().part("file", part);

    let client = reqwest::Client::new();
    let response = client
        .put(&url)
        .header("Authorization", format!("Bearer {}", api_key))
        .header("X-User-ID", user_id)
        .multipart(form)
        .send()
        .await
        .context("Failed to reach mecha10-api")?;

    let status = response.status();
    if !status.is_success() {
        let body = response.text().await.unwrap_or_default();
        anyhow::bail!("RL bundle upload failed with status {}: {}", status, body);
    }

    Ok(())
}

/// Download the RL project bundle previously uploaded for `robot_id` via [`upload_rl_bundle`] -
/// LAB-2769's download side of LAB-2767's endpoint, called by `mecha10 run rl
/// --download-bundle-only` on the trainer host. Authenticated identically to `upload_rl_bundle`
/// (`Authorization: Bearer <api_key>` + `X-User-ID`); returns the raw tar.gz bytes for
/// [`unpack_rl_bundle_archive`] to unpack.
pub async fn download_rl_bundle(
    control_plane_url: &str,
    robot_id: &str,
    api_key: &str,
    user_id: &str,
) -> Result<Vec<u8>> {
    let url = build_rl_bundle_url(control_plane_url, robot_id);

    let client = reqwest::Client::new();
    let response = client
        .get(&url)
        .header("Authorization", format!("Bearer {}", api_key))
        .header("X-User-ID", user_id)
        .send()
        .await
        .context("Failed to reach mecha10-api")?;

    let status = response.status();
    if !status.is_success() {
        let body = response.text().await.unwrap_or_default();
        anyhow::bail!("RL bundle download failed with status {}: {}", status, body);
    }

    let bytes = response
        .bytes()
        .await
        .context("Failed to read RL bundle download response body")?;

    Ok(bytes.to_vec())
}

/// Unpack an RL project bundle archive (as built by [`build_rl_bundle_archive`] and returned by
/// [`download_rl_bundle`]) into `dest_root` - entries land at `dest_root/configs/rl/...` and/or
/// `dest_root/configs/simulation/{models,environments}/...`, mirroring the project-relative
/// layout the archive was built from (same "archive root mirrors project layout" convention as
/// `mecha10_robot_config::extract_configs_archive`). Overwrites files already present at the
/// destination; does not delete anything at the destination that's absent from the archive.
pub fn unpack_rl_bundle_archive(archive_bytes: &[u8], dest_root: &Path) -> Result<()> {
    std::fs::create_dir_all(dest_root)
        .with_context(|| format!("Failed to create destination directory {}", dest_root.display()))?;

    let decoder = GzDecoder::new(archive_bytes);
    let mut archive = Archive::new(decoder);
    archive
        .unpack(dest_root)
        .context("Failed to unpack RL project bundle")?;

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use flate2::read::GzDecoder;
    use std::collections::BTreeSet;
    use tar::Archive;
    use tempfile::tempdir;

    fn archive_entries(bytes: &[u8]) -> BTreeSet<String> {
        let decoder = GzDecoder::new(bytes);
        let mut archive = Archive::new(decoder);
        archive
            .entries()
            .unwrap()
            .map(|e| e.unwrap().path().unwrap().to_string_lossy().to_string())
            .collect()
    }

    #[test]
    fn build_rl_bundle_archive_returns_none_when_nothing_to_package() {
        let dir = tempdir().unwrap();
        assert!(build_rl_bundle_archive(dir.path()).unwrap().is_none());
    }

    #[test]
    fn build_rl_bundle_archive_includes_configs_rl_recursively() {
        let dir = tempdir().unwrap();
        std::fs::create_dir_all(dir.path().join("configs/rl/jobs")).unwrap();
        std::fs::write(dir.path().join("configs/rl/jobs/train.yml"), "flow_type: train").unwrap();
        std::fs::create_dir_all(dir.path().join("configs/rl/tasks/task1")).unwrap();
        std::fs::write(dir.path().join("configs/rl/tasks/task1/task.json"), "{}").unwrap();

        let bytes = build_rl_bundle_archive(dir.path()).unwrap().unwrap();
        let entries = archive_entries(&bytes);

        assert!(entries.contains("configs/rl/jobs/train.yml"));
        assert!(entries.contains("configs/rl/tasks/task1/task.json"));
    }

    #[test]
    fn build_rl_bundle_archive_includes_models_and_environments_only() {
        let dir = tempdir().unwrap();
        std::fs::create_dir_all(dir.path().join("configs/simulation/models/rover")).unwrap();
        std::fs::write(dir.path().join("configs/simulation/models/rover/model.json"), "{}").unwrap();
        std::fs::create_dir_all(dir.path().join("configs/simulation/environments/room")).unwrap();
        std::fs::write(
            dir.path().join("configs/simulation/environments/room/environment.json"),
            "{}",
        )
        .unwrap();
        std::fs::create_dir_all(dir.path().join("configs/simulation/images")).unwrap();
        std::fs::write(dir.path().join("configs/simulation/images/aiko.jpg"), "not-a-real-jpg").unwrap();

        let bytes = build_rl_bundle_archive(dir.path()).unwrap().unwrap();
        let entries = archive_entries(&bytes);

        assert!(entries.contains("configs/simulation/models/rover/model.json"));
        assert!(entries.contains("configs/simulation/environments/room/environment.json"));
        assert!(
            !entries.iter().any(|e| e.starts_with("configs/simulation/images")),
            "images/ must be excluded from the RL bundle: {entries:?}"
        );
    }

    #[test]
    fn build_rl_bundle_archive_packages_rl_configs_even_without_simulation_assets() {
        let dir = tempdir().unwrap();
        std::fs::create_dir_all(dir.path().join("configs/rl/jobs")).unwrap();
        std::fs::write(dir.path().join("configs/rl/jobs/train.yml"), "flow_type: train").unwrap();

        let bytes = build_rl_bundle_archive(dir.path()).unwrap().unwrap();
        let entries = archive_entries(&bytes);
        assert!(entries.contains("configs/rl/jobs/train.yml"));
        assert!(
            !entries.iter().any(|e| e.starts_with("configs/simulation")),
            "no simulation assets exist in this project, none should be archived: {entries:?}"
        );
    }

    #[test]
    fn build_rl_bundle_archive_packages_simulation_assets_even_without_configs_rl() {
        let dir = tempdir().unwrap();
        std::fs::create_dir_all(dir.path().join("configs/simulation/models")).unwrap();
        std::fs::write(dir.path().join("configs/simulation/models/m.json"), "{}").unwrap();

        let bytes = build_rl_bundle_archive(dir.path()).unwrap().unwrap();
        let entries = archive_entries(&bytes);
        assert!(entries.contains("configs/simulation/models/m.json"));
        assert!(
            !entries.iter().any(|e| e.starts_with("configs/rl")),
            "no configs/rl exists in this project, none should be archived: {entries:?}"
        );
    }

    #[test]
    fn rl_bundle_tarball_name_matches_versioned_archive_convention() {
        assert_eq!(
            rl_bundle_tarball_name("my-robot", "1.2.3"),
            "my-robot-v1.2.3-rl-bundle.tar.gz"
        );
    }

    #[test]
    fn build_rl_bundle_url_scopes_by_robot_id_and_strips_trailing_slash() {
        assert_eq!(
            build_rl_bundle_url("https://api.example.com/", "robot-1"),
            "https://api.example.com/api/robots/robot-1/rl-bundle"
        );
    }

    #[tokio::test]
    async fn upload_rl_bundle_errors_when_server_unreachable() {
        // Point at a URL nothing is listening on - upload_rl_bundle must surface a real error,
        // not panic or silently succeed.
        let result = upload_rl_bundle(
            "http://127.0.0.1:9",
            "robot-1",
            b"fake archive".to_vec(),
            "robot-1-v1.0.0-rl-bundle.tar.gz",
            "test-api-key",
            "user-1",
        )
        .await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn download_rl_bundle_errors_when_server_unreachable() {
        // Mirrors upload_rl_bundle_errors_when_server_unreachable - download_rl_bundle must
        // surface a real error, not panic or silently return an empty archive.
        let result = download_rl_bundle("http://127.0.0.1:9", "robot-1", "test-api-key", "user-1").await;
        assert!(result.is_err());
    }

    #[test]
    fn unpack_rl_bundle_archive_round_trips_build_rl_bundle_archive() {
        let src = tempdir().unwrap();
        std::fs::create_dir_all(src.path().join("configs/rl/jobs")).unwrap();
        std::fs::write(src.path().join("configs/rl/jobs/train.yml"), "flow_type: train").unwrap();
        std::fs::create_dir_all(src.path().join("configs/simulation/models/rover")).unwrap();
        std::fs::write(src.path().join("configs/simulation/models/rover/model.json"), "{}").unwrap();

        let bytes = build_rl_bundle_archive(src.path()).unwrap().unwrap();

        let dest = tempdir().unwrap();
        unpack_rl_bundle_archive(&bytes, dest.path()).unwrap();

        assert_eq!(
            std::fs::read_to_string(dest.path().join("configs/rl/jobs/train.yml")).unwrap(),
            "flow_type: train"
        );
        assert_eq!(
            std::fs::read_to_string(dest.path().join("configs/simulation/models/rover/model.json")).unwrap(),
            "{}"
        );
    }

    #[test]
    fn unpack_rl_bundle_archive_creates_missing_destination_directory() {
        let src = tempdir().unwrap();
        std::fs::create_dir_all(src.path().join("configs/rl/jobs")).unwrap();
        std::fs::write(src.path().join("configs/rl/jobs/train.yml"), "flow_type: train").unwrap();
        let bytes = build_rl_bundle_archive(src.path()).unwrap().unwrap();

        let dest_parent = tempdir().unwrap();
        let dest = dest_parent.path().join("does-not-exist-yet");
        assert!(!dest.exists());

        unpack_rl_bundle_archive(&bytes, &dest).unwrap();

        assert!(dest.join("configs/rl/jobs/train.yml").exists());
    }

    #[test]
    fn unpack_rl_bundle_archive_errors_on_invalid_archive_bytes() {
        let dest = tempdir().unwrap();
        let result = unpack_rl_bundle_archive(b"not a gzip archive", dest.path());
        assert!(result.is_err());
    }
}