mecha10_cli/services/

rl_bundle.rs

//! RL project bundle download (LAB-2769)
//!
//! `mecha10 run rl --download-bundle-only` (on the trainer host) downloads and unpacks a
//! robot-scoped RL project bundle - a tar.gz of a project's `configs/rl/` (job/task
//! definitions) and `configs/simulation/{models,environments}` (the assets those jobs
//! reference) - before training starts. See "RL: Vision and Phased Roadmap" doc.
//!
//! This bundle is deliberately separate from `mecha10 config push` (`mecha10_robot_config`'s
//! per-key config sync): that one syncs `configs/` file-by-file 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.

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

/// 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: `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
    )
}

/// Download the RL project bundle previously uploaded for `robot_id` -
/// 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. 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::write::GzEncoder;
    use flate2::Compression;
    use tempfile::tempdir;

    /// Builds a minimal tar.gz archive containing a single `configs/rl/jobs/train.yml` entry -
    /// a stand-in for a real RL project bundle, since bundle construction itself no longer lives
    /// in this module (see module doc).
    fn sample_bundle_archive() -> Vec<u8> {
        let mut bytes = Vec::new();
        let encoder = GzEncoder::new(&mut bytes, Compression::default());
        let mut builder = tar::Builder::new(encoder);

        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();
        builder
            .append_dir_all("configs/rl", src.path().join("configs/rl"))
            .unwrap();

        let encoder = builder.into_inner().unwrap();
        encoder.finish().unwrap();
        bytes
    }

    #[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 download_rl_bundle_errors_when_server_unreachable() {
        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_unpacks_entries_into_the_destination() {
        let bytes = sample_bundle_archive();

        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"
        );
    }

    #[test]
    fn unpack_rl_bundle_archive_creates_missing_destination_directory() {
        let bytes = sample_bundle_archive();

        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());
    }
}