use anyhow::{Context, Result};
use flate2::read::GzDecoder;
use flate2::write::GzEncoder;
use flate2::Compression;
use std::path::Path;
use tar::{Archive, Builder};
const SIMULATION_ASSET_DIRS: [&str; 2] = ["models", "environments"];
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))
}
pub fn rl_bundle_tarball_name(project_name: &str, version: &str) -> String {
format!("{project_name}-v{version}-rl-bundle.tar.gz")
}
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
)
}
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(())
}
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())
}
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() {
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() {
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());
}
}