Skip to main content

mecha10_cli/services/
rl_bundle.rs

1//! RL project bundle download (LAB-2769)
2//!
3//! Downloads and unpacks a robot-scoped RL project bundle - a tar.gz of a project's
4//! `configs/rl/` and `configs/simulation/{models,environments}` - before training starts.
5
6use anyhow::{Context, Result};
7use flate2::read::GzDecoder;
8use std::path::Path;
9use tar::Archive;
10
11/// URL for LAB-2767's robot-scoped RL project bundle upload endpoint.
12///
13/// Assumed shape: `PUT {control_plane_url}/api/robots/{robot_id}/rl-bundle`.
14fn build_rl_bundle_url(control_plane_url: &str, robot_id: &str) -> String {
15    format!(
16        "{}/api/robots/{}/rl-bundle",
17        control_plane_url.trim_end_matches('/'),
18        robot_id
19    )
20}
21
22/// Download the RL project bundle previously uploaded for `robot_id`, returning the raw tar.gz
23/// bytes for [`unpack_rl_bundle_archive`] to unpack.
24pub async fn download_rl_bundle(
25    control_plane_url: &str,
26    robot_id: &str,
27    api_key: &str,
28    user_id: &str,
29) -> Result<Vec<u8>> {
30    let url = build_rl_bundle_url(control_plane_url, robot_id);
31
32    let client = reqwest::Client::new();
33    let response = client
34        .get(&url)
35        .header("Authorization", format!("Bearer {}", api_key))
36        .header("X-User-ID", user_id)
37        .send()
38        .await
39        .context("Failed to reach mecha10-api")?;
40
41    let status = response.status();
42    if !status.is_success() {
43        let body = response.text().await.unwrap_or_default();
44        anyhow::bail!("RL bundle download failed with status {}: {}", status, body);
45    }
46
47    let bytes = response
48        .bytes()
49        .await
50        .context("Failed to read RL bundle download response body")?;
51
52    Ok(bytes.to_vec())
53}
54
55/// Unpack an RL project bundle archive (as returned by [`download_rl_bundle`]) into `dest_root`.
56/// Overwrites files already present at the destination; does not delete anything at the
57/// destination that's absent from the archive.
58pub fn unpack_rl_bundle_archive(archive_bytes: &[u8], dest_root: &Path) -> Result<()> {
59    std::fs::create_dir_all(dest_root)
60        .with_context(|| format!("Failed to create destination directory {}", dest_root.display()))?;
61
62    let decoder = GzDecoder::new(archive_bytes);
63    let mut archive = Archive::new(decoder);
64    archive
65        .unpack(dest_root)
66        .context("Failed to unpack RL project bundle")?;
67
68    Ok(())
69}
70
71#[cfg(test)]
72mod tests {
73    use super::*;
74    use flate2::write::GzEncoder;
75    use flate2::Compression;
76    use tempfile::tempdir;
77
78    /// Builds a minimal tar.gz archive containing a single `configs/rl/jobs/train.yml` entry.
79    fn sample_bundle_archive() -> Vec<u8> {
80        let mut bytes = Vec::new();
81        let encoder = GzEncoder::new(&mut bytes, Compression::default());
82        let mut builder = tar::Builder::new(encoder);
83
84        let src = tempdir().unwrap();
85        std::fs::create_dir_all(src.path().join("configs/rl/jobs")).unwrap();
86        std::fs::write(src.path().join("configs/rl/jobs/train.yml"), "flow_type: train").unwrap();
87        builder
88            .append_dir_all("configs/rl", src.path().join("configs/rl"))
89            .unwrap();
90
91        let encoder = builder.into_inner().unwrap();
92        encoder.finish().unwrap();
93        bytes
94    }
95
96    #[test]
97    fn build_rl_bundle_url_scopes_by_robot_id_and_strips_trailing_slash() {
98        assert_eq!(
99            build_rl_bundle_url("https://api.example.com/", "robot-1"),
100            "https://api.example.com/api/robots/robot-1/rl-bundle"
101        );
102    }
103
104    #[tokio::test]
105    async fn download_rl_bundle_errors_when_server_unreachable() {
106        let result = download_rl_bundle("http://127.0.0.1:9", "robot-1", "test-api-key", "user-1").await;
107        assert!(result.is_err());
108    }
109
110    #[test]
111    fn unpack_rl_bundle_archive_unpacks_entries_into_the_destination() {
112        let bytes = sample_bundle_archive();
113
114        let dest = tempdir().unwrap();
115        unpack_rl_bundle_archive(&bytes, dest.path()).unwrap();
116
117        assert_eq!(
118            std::fs::read_to_string(dest.path().join("configs/rl/jobs/train.yml")).unwrap(),
119            "flow_type: train"
120        );
121    }
122
123    #[test]
124    fn unpack_rl_bundle_archive_creates_missing_destination_directory() {
125        let bytes = sample_bundle_archive();
126
127        let dest_parent = tempdir().unwrap();
128        let dest = dest_parent.path().join("does-not-exist-yet");
129        assert!(!dest.exists());
130
131        unpack_rl_bundle_archive(&bytes, &dest).unwrap();
132
133        assert!(dest.join("configs/rl/jobs/train.yml").exists());
134    }
135
136    #[test]
137    fn unpack_rl_bundle_archive_errors_on_invalid_archive_bytes() {
138        let dest = tempdir().unwrap();
139        let result = unpack_rl_bundle_archive(b"not a gzip archive", dest.path());
140        assert!(result.is_err());
141    }
142}