use crate::commands::ModelCommand;
use crate::context::CliContext;
use crate::paths;
use crate::services::credentials::CredentialsService;
use crate::services::ModelService;
use anyhow::{Context, Result};
use chrono::{DateTime, Utc};
use console::style;
use dialoguer::{theme::ColorfulTheme, Confirm, Select};
use indicatif::{ProgressBar, ProgressStyle};
use mecha10_core::context::split_base_and_custom;
use mecha10_models::services::model_service::ModelOrigin;
use serde::Deserialize;
use std::collections::HashMap;
use std::path::PathBuf;
use uuid::Uuid;
#[derive(Debug, Deserialize)]
struct ModelResponse {
id: Uuid,
name: String,
latest_version: i32,
}
#[derive(Debug, Deserialize)]
struct ModelVersionResponse {
version: i32,
created_at: DateTime<Utc>,
hyperparameters: HashMap<String, serde_json::Value>,
}
fn model_by_name_url(base: &str, name: &str) -> String {
format!(
"{}/api/models/by-name/{}",
base.trim_end_matches('/'),
urlencoding_path(name)
)
}
fn model_version_url(base: &str, model_id: Uuid, version: i32) -> String {
format!(
"{}/api/models/{model_id}/versions/number/{version}",
base.trim_end_matches('/')
)
}
fn model_versions_url(base: &str, model_id: Uuid) -> String {
format!("{}/api/models/{model_id}/versions", base.trim_end_matches('/'))
}
#[derive(Debug, Deserialize)]
struct ModelVersionList {
versions: Vec<ModelVersionResponse>,
}
fn slug_component(part: &str) -> String {
let mut slug = String::with_capacity(part.len());
let mut last_was_sep = true;
for c in part.trim().chars() {
if c.is_ascii_alphanumeric() {
slug.push(c.to_ascii_lowercase());
last_was_sep = false;
} else if !last_was_sep {
slug.push('_');
last_was_sep = true;
}
}
slug.trim_end_matches('_').to_string()
}
fn version_picker_label(v: &ModelVersionResponse) -> String {
let mut params: Vec<(&String, &serde_json::Value)> = v.hyperparameters.iter().collect();
params.sort_by(|a, b| a.0.cmp(b.0));
let params_str = params
.iter()
.map(|(k, val)| format!("{k}={val}"))
.collect::<Vec<_>>()
.join(", ");
let params_display = if params_str.is_empty() {
"no hyperparameters recorded".to_string()
} else if params_str.len() > 80 {
format!("{}…", ¶ms_str[..80])
} else {
params_str
};
format!(
"v{} - {} - {}",
v.version,
v.created_at.format("%Y-%m-%d %H:%M UTC"),
params_display
)
}
fn select_version(versions: &[ModelVersionResponse]) -> Result<i32> {
let items: Vec<String> = versions.iter().map(version_picker_label).collect();
let selection = Select::with_theme(&ColorfulTheme::default())
.with_prompt("Select a model version")
.items(&items)
.default(0)
.interact()
.map_err(|e| anyhow::anyhow!("failed to get version selection: {e}"))?;
Ok(versions[selection].version)
}
fn urlencoding_path(segment: &str) -> String {
segment.replace(' ', "%20").replace('/', "%2F")
}
fn merge_registry_selection(
mut config: serde_json::Value,
model_name: &str,
version: i32,
model_path: &str,
) -> Result<serde_json::Value> {
let obj = config
.as_object_mut()
.context("existing node config is not a JSON object")?;
let env_sections: Vec<&str> = ["dev", "production"]
.into_iter()
.filter(|key| obj.contains_key(*key))
.collect();
if env_sections.is_empty() {
merge_model_block(obj, model_name, version, model_path)?;
} else {
for env_key in env_sections {
let section_obj = obj
.get_mut(env_key)
.and_then(|v| v.as_object_mut())
.with_context(|| format!("existing node config's \"{env_key}\" section is not a JSON object"))?;
merge_model_block(section_obj, model_name, version, model_path)?;
}
}
Ok(config)
}
fn merge_model_block(
obj: &mut serde_json::Map<String, serde_json::Value>,
model_name: &str,
version: i32,
model_path: &str,
) -> Result<()> {
let model_obj = obj
.entry("model".to_string())
.or_insert_with(|| serde_json::json!({}))
.as_object_mut()
.context("existing node config's \"model\" field is not a JSON object")?;
model_obj.insert("name".to_string(), serde_json::json!(model_name));
model_obj.insert("version".to_string(), serde_json::json!(version));
model_obj.insert("source".to_string(), serde_json::json!("registry"));
model_obj.insert("path".to_string(), serde_json::json!(model_path));
Ok(())
}
fn local_artifact_path(model_name: &str, version: i32) -> PathBuf {
PathBuf::from(paths::project::MODELS_DIR)
.join("registry")
.join(model_name)
.join(version.to_string())
.join("model.onnx")
}
async fn control_plane_base_url(ctx: &mut CliContext) -> Result<String> {
let project_config = ctx.load_project_config().await?;
Ok(project_config.environments.control_plane_url())
}
pub async fn handle_model(ctx: &mut CliContext, cmd: &ModelCommand) -> Result<()> {
if let ModelCommand::Set { node, model, version } = cmd {
return handle_set(ctx, node, model, *version).await;
}
println!();
println!("{}", style("🤖 AI Model Management").bold());
println!("{}", style("━".repeat(50)).dim());
println!();
let models_dir = if let Ok(project) = ctx.project() {
project.root().join(paths::project::MODELS_DIR)
} else {
PathBuf::from(paths::project::MODELS_DIR)
};
let service = ModelService::with_models_dir(models_dir).context("Failed to initialize model service")?;
match cmd {
ModelCommand::List { verbose, recommended } => {
handle_list(&service, *verbose, *recommended).await?;
}
ModelCommand::Pull {
name,
all,
full_catalog,
repo,
file,
custom_name,
} => {
handle_pull(
ctx,
&service,
name.as_deref(),
*all,
*full_catalog,
repo,
file,
custom_name,
)
.await?;
}
ModelCommand::Installed { verbose } => {
handle_installed(&service, *verbose).await?;
}
ModelCommand::Remove { name, yes } => {
handle_remove(&service, name, *yes).await?;
}
ModelCommand::Info { name, version } => {
handle_info(&service, name, *version).await?;
}
ModelCommand::Set { .. } => unreachable!("handled above"),
}
Ok(())
}
async fn handle_list(service: &ModelService, verbose: bool, _recommended: bool) -> Result<()> {
let catalog = service.list_catalog().await?;
if catalog.is_empty() {
println!("No models found in catalog.");
return Ok(());
}
println!("{} Models available in catalog:", style("📦").bold());
println!();
for model in &catalog {
if verbose {
println!(" {} {}", style("•").cyan(), style(&model.name).bold());
println!(" Description: {}", model.description);
println!(" Task: {}", model.task);
println!(" Repository: {}", model.repo);
println!(" File: {}", model.filename);
if let Some(preset) = &model.preprocessing_preset {
println!(" Preset: {}", preset);
}
println!();
} else {
println!(
" {} {} - {}",
style("•").cyan(),
style(&model.name).bold(),
model.description
);
}
}
if !verbose {
println!();
println!("{}", style("Use --verbose for detailed information").dim());
}
println!();
println!("To download a model: {}", style("mecha10 model pull <name>").cyan());
println!();
Ok(())
}
#[allow(clippy::too_many_arguments)]
async fn handle_pull(
ctx: &mut CliContext,
service: &ModelService,
name: Option<&str>,
all: bool,
full_catalog: bool,
repo: &Option<String>,
file: &Option<String>,
custom_name: &Option<String>,
) -> Result<()> {
let pb = ProgressBar::new_spinner();
pb.set_style(
ProgressStyle::default_spinner()
.template("{spinner:.green} {msg}")
.unwrap(),
);
pb.enable_steady_tick(std::time::Duration::from_millis(100));
if full_catalog {
println!("Downloading the entire model catalog...");
println!();
let catalog = service.list_catalog().await?;
let mut paths = Vec::new();
for entry in catalog {
pb.set_message(format!("Downloading {}", entry.name));
match service.pull(&entry.name, Some(&pb)).await {
Ok(path) => paths.push(path),
Err(e) => {
eprintln!("⚠️ Failed to download {}: {}", entry.name, e);
}
}
}
pb.finish_and_clear();
println!("{} Downloaded {} models:", style("✓").green().bold(), paths.len());
for path in &paths {
println!(" • {}", path.display());
}
println!();
} else if all {
let model_names = project_referenced_catalog_models(ctx)
.await
.context("failed to scan this project's node configs for referenced models")?;
if model_names.is_empty() {
pb.finish_and_clear();
println!("No catalog models are referenced by this project's node configs.");
println!();
println!(
" {} to download a specific model",
style("mecha10 model pull <name>").cyan()
);
println!(
" {} to download the entire curated catalog",
style("mecha10 model pull --full-catalog").cyan()
);
println!();
return Ok(());
}
println!(
"Downloading {} model(s) referenced by this project's node configs...",
model_names.len()
);
println!();
let mut paths = Vec::new();
for model_name in &model_names {
pb.set_message(format!("Downloading {}", model_name));
match service.pull(model_name, Some(&pb)).await {
Ok(path) => paths.push(path),
Err(e) => {
eprintln!("⚠️ Failed to download {}: {}", model_name, e);
}
}
}
pb.finish_and_clear();
println!("{} Downloaded {} models:", style("✓").green().bold(), paths.len());
for path in &paths {
println!(" • {}", path.display());
}
println!();
} else if let (Some(repo), Some(file), Some(custom_name)) = (repo, file, custom_name) {
pb.set_message(format!("Downloading {} from {}", custom_name, repo));
let path = service.pull_from_repo(repo, file, custom_name, Some(&pb)).await?;
pb.finish_and_clear();
println!(
"{} Downloaded custom model '{}'",
style("✓").green().bold(),
custom_name
);
println!(" Location: {}", path.display());
println!();
} else if let Some(name) = name {
pb.set_message(format!("Downloading {}", name));
let path = service.pull(name, Some(&pb)).await?;
pb.finish_and_clear();
println!("{} Downloaded '{}'", style("✓").green().bold(), name);
println!(" Location: {}", path.display());
println!();
} else {
pb.finish_and_clear();
return Err(anyhow::anyhow!(
"Please specify a model name, --all, --full-catalog, or --repo with --file and --custom-name"
));
}
Ok(())
}
async fn project_referenced_catalog_models(ctx: &mut CliContext) -> Result<Vec<String>> {
let node_names = ctx
.project()
.context(
"`mecha10 model pull --all` must be run inside a mecha10 project (no mecha10.json found) - \
use `mecha10 model pull --full-catalog` to download the entire catalog instead",
)?
.list_enabled_nodes()
.await?;
let project_root = ctx.project()?.root().to_path_buf();
let mut names = std::collections::BTreeSet::new();
for node_name in &node_names {
for candidate in paths::config::node_config_candidates(node_name) {
let path = project_root.join(&candidate);
if !path.exists() {
continue;
}
let content = tokio::fs::read_to_string(&path)
.await
.with_context(|| format!("failed to read node config: {}", path.display()))?;
if let Some(model_name) = extract_catalog_model_reference(&content, node_name, &path)? {
names.insert(model_name);
}
break;
}
}
Ok(names.into_iter().collect())
}
#[derive(Debug, Deserialize, Default)]
struct NodeModelSection {
model: Option<ModelReference>,
}
#[derive(Debug, Deserialize)]
struct ModelReference {
name: Option<String>,
#[serde(default)]
source: Option<String>,
}
impl ModelReference {
fn is_catalog(&self) -> bool {
!matches!(self.source.as_deref(), Some("registry"))
}
}
fn extract_catalog_model_reference(
content: &str,
node_name: &str,
config_path: &std::path::Path,
) -> Result<Option<String>> {
let profile = std::env::var("MECHA10_ENV").unwrap_or_else(|_| "dev".to_string());
let Ok(value) = serde_json::from_str::<serde_json::Value>(content) else {
return Ok(None);
};
let Some(obj) = value.as_object() else {
return Ok(None);
};
let section_value = if obj.contains_key("dev") || obj.contains_key("production") {
let env_key = match profile.as_str() {
"production" | "prod" => "production",
_ => "dev",
};
let Some(section) = obj.get(env_key) else {
anyhow::bail!(
"Node '{}' config is missing its '{}' environment section (MECHA10_ENV='{}'): {}",
node_name,
env_key,
profile,
config_path.display()
);
};
section.clone()
} else {
value
};
let Ok(merged) = split_base_and_custom(§ion_value) else {
return Ok(None);
};
let Ok(section) = serde_json::from_value::<NodeModelSection>(merged) else {
return Ok(None);
};
let Some(model) = section.model else {
return Ok(None);
};
if !model.is_catalog() {
return Ok(None);
}
Ok(model.name)
}
async fn handle_installed(service: &ModelService, verbose: bool) -> Result<()> {
let installed = service.list_installed().await?;
if installed.is_empty() {
println!("No models installed.");
println!();
println!("To download models: {}", style("mecha10 model pull <name>").cyan());
println!();
return Ok(());
}
let count_str = if installed.len() == 1 {
"1 model".to_string()
} else {
format!("{} models", installed.len())
};
println!("{} {} installed:", style("📦").bold(), count_str);
println!();
for model in &installed {
let source_label = match (model.source, model.version) {
(ModelOrigin::Registry, Some(version)) => format!("registry v{version}"),
(ModelOrigin::Registry, None) => "registry".to_string(),
(ModelOrigin::Catalog, _) => "catalog".to_string(),
};
if verbose {
println!(" {} {}", style("•").cyan(), style(&model.name).bold());
println!(" Path: {}", model.path.display());
println!(" Size: {} bytes", model.size);
println!(" Source: {}", source_label);
if let Some(entry) = &model.catalog_entry {
println!(" From catalog: {}", entry.description);
println!(" Task: {}", entry.task);
} else if model.source == ModelOrigin::Catalog {
println!(" {} Custom model", style("ℹ").dim());
}
println!();
} else {
let size_mb = model.size as f64 / 1_048_576.0;
let extra_info = if let Some(entry) = &model.catalog_entry {
format!(" ({})", entry.task)
} else if model.source == ModelOrigin::Catalog {
" (custom)".to_string()
} else {
format!(" ({source_label})")
};
println!(
" {} {} - {:.2} MB{}",
style("•").cyan(),
style(&model.name).bold(),
size_mb,
style(extra_info).dim()
);
}
}
if !verbose {
println!();
println!("{}", style("Use --verbose for detailed information").dim());
}
println!();
Ok(())
}
async fn handle_remove(service: &ModelService, name: &str, skip_confirm: bool) -> Result<()> {
let installed = service.list_installed().await?;
let model = installed
.iter()
.find(|m| m.name == name)
.ok_or_else(|| anyhow::anyhow!("Model '{}' is not installed", name))?;
if !skip_confirm {
let confirm = Confirm::new()
.with_prompt(format!("Remove model '{}'? ({} bytes)", name, model.size))
.default(false)
.interact()?;
if !confirm {
println!("Cancelled.");
return Ok(());
}
}
service.remove(name).await?;
println!("{} Removed model '{}'", style("✓").green().bold(), name);
println!();
Ok(())
}
async fn handle_info(service: &ModelService, name: &str, version: Option<i32>) -> Result<()> {
let info = service.info(name, version).await?;
println!("{} Model: {}", style("ℹ").blue().bold(), style(name).bold());
println!();
if let Some(entry) = &info.catalog_entry {
println!("{}", style("Catalog Information:").underlined());
println!(" Description: {}", entry.description);
println!(" Task: {}", entry.task);
println!(" Repository: {}", entry.repo);
println!(" File: {}", entry.filename);
if let Some(preset) = &entry.preprocessing_preset {
println!(" Preset: {}", preset);
}
println!();
} else {
println!("{}", style("Not found in catalog (custom model)").dim());
println!();
}
if let Some(installed) = &info.installed_info {
let source_label = match installed.source {
ModelOrigin::Registry => "registry",
ModelOrigin::Catalog => "catalog",
};
let version_label = match installed.version {
Some(v) => v.to_string(),
None => "latest".to_string(),
};
println!("{}", style("Installation:").underlined());
println!(" Status: {}", style("Installed").green());
println!(" Source: {}", source_label);
println!(" Version: {}", version_label);
println!(" Path: {}", installed.path.display());
println!(
" Size: {} bytes ({:.2} MB)",
installed.size,
installed.size as f64 / 1_048_576.0
);
println!();
} else {
println!("{}", style("Installation:").underlined());
println!(" Status: {}", style("Not installed").yellow());
println!();
println!(" To install: {}", style(format!("mecha10 model pull {}", name)).cyan());
println!(
" {} A registry model can be fetched with: {}",
style("ℹ").dim(),
style(format!("mecha10 model set <node> --model {}", name)).cyan()
);
println!();
}
Ok(())
}
async fn handle_set(ctx: &mut CliContext, node: &str, model_name: &str, version: Option<i32>) -> Result<()> {
println!();
println!("{}", style("Resolving model from registry...").bold());
let credentials_service = CredentialsService::new();
let credentials = if ctx.is_ci() {
credentials_service
.load()?
.ok_or_else(|| anyhow::anyhow!("Not logged in. Run `mecha10 auth login` first."))?
} else {
mecha10_auth::ensure_authenticated_interactive(&credentials_service).await?
};
let base_url = control_plane_base_url(ctx).await?;
let client = reqwest::Client::new();
let project_config = ctx.load_project_config().await?;
let prefix = slug_component(&project_config.name);
let full_name = if prefix.is_empty() || model_name.starts_with(&format!("{prefix}_")) {
model_name.to_string()
} else {
format!("{prefix}_{model_name}")
};
let model_resp = client
.get(model_by_name_url(&base_url, &full_name))
.bearer_auth(&credentials.api_key)
.send()
.await
.context("failed to reach the model registry")?
.error_for_status()
.context(format!("model '{full_name}' was not found in the registry"))?
.json::<ModelResponse>()
.await
.context("model registry returned an unexpected response")?;
let resolved_version = match version {
Some(v) => v,
None => {
let versions_resp = client
.get(model_versions_url(&base_url, model_resp.id))
.bearer_auth(&credentials.api_key)
.send()
.await
.context("failed to reach the model registry")?
.error_for_status()
.context("failed to list versions for this model")?
.json::<ModelVersionList>()
.await
.context("model registry returned an unexpected response")?;
match versions_resp.versions.as_slice() {
[] => model_resp.latest_version,
[only] => only.version,
many => select_version(many)?,
}
}
};
let version_resp = client
.get(model_version_url(&base_url, model_resp.id, resolved_version))
.bearer_auth(&credentials.api_key)
.send()
.await
.context("failed to reach the model registry")?
.error_for_status()
.context(format!(
"version {resolved_version} of '{}' was not found",
model_resp.name
))?
.json::<ModelVersionResponse>()
.await
.context("model registry returned an unexpected response")?;
println!(
" {} {} v{}",
style("Found:").dim(),
style(&model_resp.name).bold(),
version_resp.version
);
println!("{}", style("Downloading artifact...").bold());
let project_root = ctx.project()?.root().to_path_buf();
let models_dir = project_root.join(paths::project::MODELS_DIR);
let model_service = ModelService::with_models_dir(models_dir).context("Failed to initialize model service")?;
let absolute_path = model_service
.pull_registry_version(
&base_url,
&credentials.api_key,
&model_resp.name,
version_resp.version,
None,
)
.await
.context("failed to download the model artifact")?;
let relative_path = local_artifact_path(&model_resp.name, version_resp.version);
let artifact_size = tokio::fs::metadata(&absolute_path)
.await
.context("failed to read downloaded artifact metadata")?
.len();
println!(
" {} {} ({} bytes)",
style("Saved:").dim(),
absolute_path.display(),
artifact_size
);
println!("{}", style("Updating node config...").bold());
let framework_config_path = project_root.join(paths::config::framework_node(node));
let project_config_path = project_root.join(paths::config::project_node(node));
let config_path = if framework_config_path.exists() {
framework_config_path
} else {
project_config_path
};
let existing = if config_path.exists() {
let content = tokio::fs::read_to_string(&config_path)
.await
.context("failed to read existing node config")?;
serde_json::from_str(&content).context("existing node config is not valid JSON")?
} else {
serde_json::json!({})
};
let relative_path_str = relative_path.to_string_lossy().replace('\\', "/");
let updated = merge_registry_selection(existing, &model_resp.name, version_resp.version, &relative_path_str)?;
if let Some(parent) = config_path.parent() {
tokio::fs::create_dir_all(parent)
.await
.context("failed to create node config directory")?;
}
tokio::fs::write(&config_path, serde_json::to_string_pretty(&updated)?)
.await
.context("failed to write node config")?;
println!(" {} {}", style("Updated:").dim(), config_path.display());
println!();
println!(
"{} {} is now set to {} v{} (registry)",
style("✔").green().bold(),
node,
model_resp.name,
version_resp.version
);
println!();
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use crate::context::CliContextBuilder;
use serial_test::serial;
#[test]
fn slug_component_lowercases_and_collapses_non_alphanumerics() {
assert_eq!(slug_component("Find Target!"), "find_target");
}
#[test]
fn slug_component_trims_trailing_underscores() {
assert_eq!(slug_component(" test-rover "), "test_rover");
}
#[test]
fn slug_component_is_empty_for_an_empty_input() {
assert_eq!(slug_component(""), "");
}
fn version(v: i32, params: &[(&str, serde_json::Value)]) -> ModelVersionResponse {
ModelVersionResponse {
version: v,
created_at: DateTime::parse_from_rfc3339("2026-09-30T14:22:00Z")
.unwrap()
.with_timezone(&Utc),
hyperparameters: params.iter().map(|(k, v)| (k.to_string(), v.clone())).collect(),
}
}
#[test]
fn version_picker_label_includes_version_datetime_and_sorted_hyperparameters() {
let v = version(
3,
&[("gamma", serde_json::json!(0.95)), ("lr", serde_json::json!(0.0003))],
);
assert_eq!(
version_picker_label(&v),
"v3 - 2026-09-30 14:22 UTC - gamma=0.95, lr=0.0003"
);
}
#[test]
fn version_picker_label_notes_when_no_hyperparameters_were_recorded() {
let v = version(1, &[]);
assert_eq!(
version_picker_label(&v),
"v1 - 2026-09-30 14:22 UTC - no hyperparameters recorded"
);
}
#[test]
fn model_by_name_url_appends_the_expected_path() {
assert_eq!(
model_by_name_url("https://control-plane.example.com", "rover-detector"),
"https://control-plane.example.com/api/models/by-name/rover-detector"
);
}
#[test]
fn model_by_name_url_strips_a_trailing_slash_from_the_base() {
assert_eq!(
model_by_name_url("https://control-plane.example.com/", "rover-detector"),
"https://control-plane.example.com/api/models/by-name/rover-detector"
);
}
#[test]
fn model_version_url_is_derived_consistently() {
let id = Uuid::nil();
assert_eq!(
model_version_url("https://cp.example.com", id, 3),
format!("https://cp.example.com/api/models/{id}/versions/number/3")
);
}
#[test]
fn local_artifact_path_is_namespaced_by_source_and_version() {
assert_eq!(
local_artifact_path("rover-detector", 4),
PathBuf::from("models/registry/rover-detector/4/model.onnx")
);
}
#[test]
fn local_artifact_path_never_collides_across_versions() {
assert_ne!(
local_artifact_path("rover-detector", 1),
local_artifact_path("rover-detector", 2)
);
}
#[test]
fn merge_registry_selection_preserves_unrelated_fields() {
let existing = serde_json::json!({
"confidence_threshold": 0.6,
"model": {
"name": "yolov8n",
"source": "catalog",
},
});
let merged = merge_registry_selection(
existing,
"rover-detector",
4,
"models/registry/rover-detector/4/model.onnx",
)
.unwrap();
assert_eq!(merged["confidence_threshold"], serde_json::json!(0.6));
assert_eq!(merged["model"]["name"], serde_json::json!("rover-detector"));
assert_eq!(merged["model"]["version"], serde_json::json!(4));
assert_eq!(merged["model"]["source"], serde_json::json!("registry"));
assert_eq!(
merged["model"]["path"],
serde_json::json!("models/registry/rover-detector/4/model.onnx")
);
}
#[test]
fn merge_registry_selection_works_on_an_empty_config() {
let merged = merge_registry_selection(serde_json::json!({}), "m", 1, "models/registry/m/1/model.onnx").unwrap();
assert_eq!(merged["model"]["source"], serde_json::json!("registry"));
}
#[test]
fn merge_registry_selection_writes_into_dev_and_production_sections() {
let existing = serde_json::json!({
"dev": {
"model": { "name": "find-target" },
"target_class": "cat",
},
"production": {
"model": { "name": "find-target" },
"target_class": "cat",
},
});
let merged = merge_registry_selection(
existing,
"test_find_target_ppo",
2,
"models/registry/test_find_target_ppo/2/model.onnx",
)
.unwrap();
for env_key in ["dev", "production"] {
assert_eq!(
merged[env_key]["model"]["name"],
serde_json::json!("test_find_target_ppo")
);
assert_eq!(merged[env_key]["model"]["version"], serde_json::json!(2));
assert_eq!(merged[env_key]["model"]["source"], serde_json::json!("registry"));
assert_eq!(
merged[env_key]["model"]["path"],
serde_json::json!("models/registry/test_find_target_ppo/2/model.onnx")
);
assert_eq!(merged[env_key]["target_class"], serde_json::json!("cat"));
}
assert!(
merged.get("model").is_none(),
"must not also write a flat top-level model key"
);
}
#[test]
fn merge_registry_selection_only_touches_env_sections_that_exist() {
let existing = serde_json::json!({
"dev": { "model": { "name": "find-target" } },
});
let merged = merge_registry_selection(existing, "m", 1, "models/registry/m/1/model.onnx").unwrap();
assert_eq!(merged["dev"]["model"]["source"], serde_json::json!("registry"));
assert!(merged.get("production").is_none());
}
#[test]
fn merge_registry_selection_rejects_a_non_object_config() {
let err = merge_registry_selection(serde_json::json!([1, 2, 3]), "m", 1, "models/registry/m/1/model.onnx")
.unwrap_err();
assert!(err.to_string().contains("not a JSON object"));
}
#[test]
#[serial]
fn extract_catalog_model_reference_reads_the_dev_section_by_default() {
std::env::remove_var("MECHA10_ENV");
let content = r#"{"dev": {"custom": {"model": {"name": "yolov8n"}}}, "production": {"custom": {"model": {"name": "yolov8n-prod"}}}}"#;
assert_eq!(
extract_catalog_model_reference(content, "node", &PathBuf::from("config.json")).unwrap(),
Some("yolov8n".to_string())
);
}
#[test]
#[serial]
fn extract_catalog_model_reference_honors_mecha10_env_production() {
std::env::set_var("MECHA10_ENV", "production");
let content = r#"{"dev": {"custom": {"model": {"name": "yolov8n"}}}, "production": {"custom": {"model": {"name": "yolov8n-prod"}}}}"#;
let result = extract_catalog_model_reference(content, "node", &PathBuf::from("config.json"));
std::env::remove_var("MECHA10_ENV");
assert_eq!(result.unwrap(), Some("yolov8n-prod".to_string()));
}
#[test]
#[serial]
fn extract_catalog_model_reference_errors_when_active_environment_section_is_missing() {
std::env::set_var("MECHA10_ENV", "production");
let content = r#"{"dev": {"custom": {"model": {"name": "yolov8n"}}}}"#;
let err =
extract_catalog_model_reference(content, "object-detector", &PathBuf::from("config.json")).unwrap_err();
std::env::remove_var("MECHA10_ENV");
assert!(err.to_string().contains("object-detector"));
assert!(err.to_string().contains("production"));
}
#[test]
#[serial]
fn extract_catalog_model_reference_ignores_a_registry_sourced_model() {
std::env::remove_var("MECHA10_ENV");
let content = r#"{"dev": {"custom": {"model": {"name": "my-detector", "source": "registry"}}}}"#;
assert_eq!(
extract_catalog_model_reference(content, "node", &PathBuf::from("config.json")).unwrap(),
None
);
}
#[test]
#[serial]
fn extract_catalog_model_reference_treats_missing_source_as_catalog() {
std::env::remove_var("MECHA10_ENV");
let content = r#"{"dev": {"custom": {"model": {"name": "mobilenet-v2"}}}}"#;
assert_eq!(
extract_catalog_model_reference(content, "node", &PathBuf::from("config.json")).unwrap(),
Some("mobilenet-v2".to_string())
);
}
#[test]
#[serial]
fn extract_catalog_model_reference_returns_none_for_a_node_with_no_model_field() {
std::env::remove_var("MECHA10_ENV");
let content =
r#"{"dev": {"custom": {"camera_id": "camera0"}}, "production": {"custom": {"camera_id": "camera0"}}}"#;
assert_eq!(
extract_catalog_model_reference(content, "node", &PathBuf::from("config.json")).unwrap(),
None
);
}
#[test]
#[serial]
fn extract_catalog_model_reference_returns_none_for_a_non_object_model_field() {
std::env::remove_var("MECHA10_ENV");
let content = r#"{"dev": {"custom": {"model": "rover"}}, "production": {"custom": {"model": "rover"}}}"#;
assert_eq!(
extract_catalog_model_reference(content, "node", &PathBuf::from("config.json")).unwrap(),
None
);
}
#[test]
#[serial]
fn extract_catalog_model_reference_falls_back_to_flat_non_environment_aware_parsing() {
std::env::remove_var("MECHA10_ENV");
let content = r#"{"custom": {"model": {"name": "yolov8n"}}}"#;
assert_eq!(
extract_catalog_model_reference(content, "node", &PathBuf::from("config.json")).unwrap(),
Some("yolov8n".to_string())
);
}
#[test]
#[serial]
fn extract_catalog_model_reference_delegates_to_the_real_base_and_custom_split() {
std::env::remove_var("MECHA10_ENV");
let dev_section = serde_json::json!({
"custom": { "model": { "name": "yolov8n" } },
"relay_topics": ["/x"],
});
let content = serde_json::json!({ "dev": dev_section.clone(), "production": dev_section.clone() }).to_string();
let merged = split_base_and_custom(&dev_section).expect("split should succeed");
assert_eq!(
extract_catalog_model_reference(&content, "node", &PathBuf::from("config.json")).unwrap(),
merged["model"]["name"].as_str().map(|s| s.to_string())
);
}
#[test]
#[serial]
fn extract_catalog_model_reference_reads_model_nested_under_custom_like_object_detector() {
std::env::remove_var("MECHA10_ENV");
let content = r#"{
"dev": {
"custom": {
"model": { "name": "yolov8n" },
"input_size": 640,
"confidence_threshold": 0.5,
"topics": { "publishes": [], "subscribes": [] }
},
"relay_topics": []
},
"production": {
"custom": {
"model": { "name": "yolov8n" },
"input_size": 640,
"confidence_threshold": 0.5,
"topics": { "publishes": [], "subscribes": [] }
},
"relay_topics": []
}
}"#;
assert_eq!(
extract_catalog_model_reference(content, "object-detector", &PathBuf::from("config.json")).unwrap(),
Some("yolov8n".to_string())
);
}
#[test]
#[serial]
fn extract_catalog_model_reference_reads_model_nested_under_custom_like_image_classifier() {
std::env::remove_var("MECHA10_ENV");
let content = r#"{
"dev": {
"custom": {
"processing_rate_hz": 10,
"top_k": 5,
"model": { "name": "mobilenet-v2" },
"topics": { "publishes": [], "subscribes": [] }
},
"relay_topics": []
},
"production": {
"custom": {
"processing_rate_hz": 30,
"top_k": 3,
"model": { "name": "mobilenet-v2" },
"topics": { "publishes": [], "subscribes": [] }
},
"relay_topics": []
}
}"#;
assert_eq!(
extract_catalog_model_reference(content, "image-classifier", &PathBuf::from("config.json")).unwrap(),
Some("mobilenet-v2".to_string())
);
}
fn scaffold_project(node_configs: &[(&str, &str, serde_json::Value)]) -> (tempfile::TempDir, CliContext) {
let dir = tempfile::tempdir().unwrap();
let nodes: std::collections::HashMap<&str, &str> = node_configs
.iter()
.map(|(identifier, _, _)| (*identifier, "0.1.0"))
.collect();
let mecha10_json = serde_json::json!({
"name": "proj",
"version": "1.0.0",
"robot_id": "robot-1",
"nodes": nodes,
});
let config_path = dir.path().join("mecha10.json");
std::fs::write(&config_path, mecha10_json.to_string()).unwrap();
for (identifier, config_dir_name, body) in node_configs {
let is_framework = identifier.starts_with("@mecha10/");
let node_dir = if is_framework {
dir.path().join("configs/nodes/@mecha10").join(config_dir_name)
} else {
dir.path().join("configs/nodes").join(config_dir_name)
};
std::fs::create_dir_all(&node_dir).unwrap();
std::fs::write(node_dir.join("config.json"), body.to_string()).unwrap();
}
let ctx = CliContextBuilder::new()
.working_dir(dir.path().to_path_buf())
.config_path(Some(config_path))
.build()
.unwrap();
(dir, ctx)
}
#[tokio::test]
#[serial]
async fn project_referenced_catalog_models_scopes_to_a_projects_own_node_configs() {
std::env::remove_var("MECHA10_ENV");
let (_dir, mut ctx) = scaffold_project(&[
(
"@mecha10/object-detector",
"object-detector",
serde_json::json!({
"dev": {"model": {"name": "yolov8n"}},
"production": {"model": {"name": "yolov8n"}},
}),
),
(
"@mecha10/image-classifier",
"image-classifier",
serde_json::json!({
"dev": {"model": {"name": "mobilenet-v2"}},
"production": {"model": {"name": "mobilenet-v2"}},
}),
),
]);
let models = project_referenced_catalog_models(&mut ctx).await.unwrap();
assert_eq!(models, vec!["mobilenet-v2".to_string(), "yolov8n".to_string()]);
}
#[tokio::test]
#[serial]
async fn project_referenced_catalog_models_excludes_registry_sourced_models() {
std::env::remove_var("MECHA10_ENV");
let (_dir, mut ctx) = scaffold_project(&[
(
"@mecha10/object-detector",
"object-detector",
serde_json::json!({
"dev": {"model": {"name": "yolov8n"}},
"production": {"model": {"name": "yolov8n"}},
}),
),
(
"custom-registry-node",
"custom-registry-node",
serde_json::json!({
"dev": {"model": {"name": "my-registry-model", "source": "registry"}},
"production": {"model": {"name": "my-registry-model", "source": "registry"}},
}),
),
]);
let models = project_referenced_catalog_models(&mut ctx).await.unwrap();
assert_eq!(models, vec!["yolov8n".to_string()]);
}
#[tokio::test]
#[serial]
async fn project_referenced_catalog_models_deduplicates_a_model_shared_by_two_nodes() {
std::env::remove_var("MECHA10_ENV");
let (_dir, mut ctx) = scaffold_project(&[
(
"@mecha10/object-detector",
"object-detector",
serde_json::json!({
"dev": {"model": {"name": "yolov8n"}},
"production": {"model": {"name": "yolov8n"}},
}),
),
(
"second-detector",
"second-detector",
serde_json::json!({
"dev": {"model": {"name": "yolov8n"}},
"production": {"model": {"name": "yolov8n"}},
}),
),
]);
let models = project_referenced_catalog_models(&mut ctx).await.unwrap();
assert_eq!(models, vec!["yolov8n".to_string()]);
}
#[tokio::test]
#[serial]
async fn project_referenced_catalog_models_returns_empty_when_no_node_declares_a_model() {
std::env::remove_var("MECHA10_ENV");
let (_dir, mut ctx) = scaffold_project(&[(
"@mecha10/camera",
"camera",
serde_json::json!({
"dev": {"camera_id": "camera0"},
"production": {"camera_id": "camera0"},
}),
)]);
let models = project_referenced_catalog_models(&mut ctx).await.unwrap();
assert!(models.is_empty());
}
#[tokio::test]
async fn project_referenced_catalog_models_errors_outside_a_project() {
let dir = tempfile::tempdir().unwrap();
let mut ctx = CliContextBuilder::new()
.working_dir(dir.path().to_path_buf())
.build()
.unwrap();
let err = project_referenced_catalog_models(&mut ctx).await.unwrap_err();
assert!(
err.to_string().contains("--full-catalog"),
"error should point users at the --full-catalog escape hatch: {err}"
);
}
}