mecha10-cli 0.7.0

//! AI model availability checks before spawning nodes
//!
//! Used in standalone mode (where node-runner is not available) to make
//! sure any AI models a node's config declares are downloaded before the
//! node is spawned.

use crate::services::ModelService;
use anyhow::Result;
use mecha10_core::context::split_base_and_custom;
use serde::Deserialize;
use std::path::PathBuf;

/// Minimal config structure to extract model requirements from node configs
///
/// Node-specific fields (`model`, `provider`) may live on the shared `NodeConfigBase` or nested
/// under this node config's `custom` object, depending on the field, per
/// `docs/adr/0001-node-config-base-schema.md`. Rather than hand-rolling an opinion about which of
/// those two places each field currently lives (which silently broke auto-download before, PR
/// #1126, and would break again every time a field moves between the two), this struct is
/// deserialized from the *merged* output of `mecha10_core::context::split_base_and_custom`, the
/// one real implementation of that base+`custom` split, also used by
/// `Context::load_environment_aware_config` (and so transitively by `mecha10-node-runner` and
/// the launcher), so this parser is agnostic to where any of these fields currently live.
///
/// `llm-command`'s "which LLM model"/"which LLM backend" concepts are also `model.name`/
/// `model.provider` here, not separate top-level fields - see
/// `mecha10_core::node_config_base::NodeConfigBase`'s docs for why.
#[derive(Deserialize, Default)]
struct NodeModelConfig {
    model: Option<ModelInfo>,
}

#[derive(Deserialize)]
struct ModelInfo {
    name: Option<String>,
    /// LLM provider (`"local"` means Ollama - see `ensure_ollama_models_available`).
    #[serde(default)]
    provider: Option<String>,
    /// `"registry"` means the model is resolved via `mecha10 model set`, not `model pull` -
    /// see `ModelReference::is_catalog()` in `handlers/model.rs`.
    #[serde(default)]
    source: Option<String>,
}

/// Parse a node's `config.json` contents, resolving the active `dev`/`production` environment
/// section and then running it through the real [`split_base_and_custom`] split before reading
/// `model`/`model.provider` off the merged result.
fn parse_node_config(content: &str) -> Option<NodeModelConfig> {
    let profile = std::env::var("MECHA10_ENVIRONMENT").unwrap_or_else(|_| "dev".to_string());

    let value: serde_json::Value = serde_json::from_str(content).ok()?;
    let obj = value.as_object()?;

    let env_key = match profile.as_str() {
        "production" | "prod" => "production",
        _ => "dev",
    };
    let section = obj.get(env_key)?;

    let merged = split_base_and_custom(section).ok()?;
    serde_json::from_value(merged).ok()
}

/// Ensure all required AI models are available for nodes
///
/// Reads node configs and downloads missing models before spawning nodes.
/// This is called in standalone mode where node-runner is not available.
pub(super) async fn ensure_node_models_available(node_names: &[String]) -> Result<()> {
    let model_service = match ModelService::with_models_dir("models".into()) {
        Ok(service) => service,
        Err(e) => {
            eprintln!("⚠️  Warning: Could not initialize model service: {}", e);
            return Ok(());
        }
    };

    for node_name in node_names {
        // A multi-instance node's identifier may carry an encoded instance id (LAB-4131, see
        // `NodeRef::wire_identifier`) - config files are still looked up by the bare node name.
        let (node_name, _instance_id) = node_name.split_once("::").unwrap_or((node_name, ""));
        let short_name = node_name.strip_prefix("@mecha10/").unwrap_or(node_name);

        let config_paths = crate::paths::config::node_config_candidates(short_name);

        for config_path in &config_paths {
            let path = PathBuf::from(config_path);
            if !path.exists() {
                continue;
            }

            let content = match std::fs::read_to_string(&path) {
                Ok(c) => c,
                Err(_) => continue,
            };

            let config = match parse_node_config(&content) {
                Some(c) => c,
                None => continue,
            };

            if let Some(model_info) = config.model {
                if model_info.source.as_deref() == Some("registry") {
                    break;
                }
                if model_info.provider.as_deref() == Some("local") {
                    // Ollama models - handled by `ensure_ollama_models_available` instead of
                    // the ONNX catalog pull below.
                    break;
                }
                if let Some(model_name) = model_info.name {
                    let model_path = PathBuf::from(format!("models/catalog/{}/model.onnx", model_name));
                    if !model_path.exists() {
                        println!("📦 Downloading model '{}' for node '{}'...", model_name, node_name);
                        match model_service.pull(&model_name, None).await {
                            Ok(_) => println!("✅ Model '{}' ready", model_name),
                            Err(e) => {
                                eprintln!("⚠️  Warning: Failed to download model '{}': {}", model_name, e);
                                eprintln!("   You can manually download with: mecha10 model pull {}", model_name);
                            }
                        }
                    }
                }
            }
            break;
        }
    }

    Ok(())
}

/// Ollama's default host-mapped endpoint - see `${OLLAMA_PORT:-11434}` in the `ollama` Compose
/// service (`mecha10-templates`' `docker-compose.yml.template`/`docker-compose.remote.yml.template`).
/// This preflight always runs on the host (the CLI process itself, not inside any node's
/// container), so `localhost` is correct here even though the node process reaching Ollama from
/// inside a container needs `OLLAMA_ENDPOINT` pointed at the `ollama` service name instead - see
/// `llm_command::OpenAIReasoningNode::new`.
fn ollama_base_url() -> String {
    std::env::var("OLLAMA_ENDPOINT")
        .map(|endpoint| endpoint.trim_end_matches("/api/chat").to_string())
        .unwrap_or_else(|_| "http://localhost:11434".to_string())
}

/// Ensure every node declaring `"provider": "local"` (Ollama) has its `model.name` pulled before
/// nodes are spawned - the same "check disk, download if missing" preflight
/// [`ensure_node_models_available`] already does for ONNX vision models, adapted to Ollama's own
/// model store. That store is bind-mounted at `<project_root>/models/catalog` - the same
/// directory the ONNX catalog above already uses - by the `ollama` Compose service, so nothing
/// here needs to know about that path directly; Ollama manages its own layout under it once a
/// model is pulled.
///
/// A no-op (with a warning, not an error) if Ollama isn't reachable at all - this mirrors
/// `ensure_node_models_available`'s "warn and continue" behavior for the same reason: a
/// preflight check should never be the thing that blocks `mecha10 dev` from starting.
pub(super) async fn ensure_ollama_models_available(node_names: &[String]) -> Result<()> {
    let client = reqwest::Client::new();
    let base_url = ollama_base_url();

    let mut wanted_models = Vec::new();
    for node_name in node_names {
        // See the matching comment in `ensure_node_models_available` above (LAB-4131).
        let (node_name, _instance_id) = node_name.split_once("::").unwrap_or((node_name, ""));
        let short_name = node_name.strip_prefix("@mecha10/").unwrap_or(node_name);

        for config_path in crate::paths::config::node_config_candidates(short_name) {
            let path = PathBuf::from(&config_path);
            if !path.exists() {
                continue;
            }

            let Ok(content) = std::fs::read_to_string(&path) else {
                continue;
            };

            let Some(config) = parse_node_config(&content) else {
                continue;
            };

            if let Some(model_info) = config.model {
                if model_info.provider.as_deref() == Some("local") {
                    if let Some(model) = model_info.name {
                        wanted_models.push((node_name.to_string(), model));
                    }
                }
            }
            break;
        }
    }

    if wanted_models.is_empty() {
        return Ok(());
    }

    let tags_url = format!("{base_url}/api/tags");
    let installed: Vec<String> = match client.get(&tags_url).send().await {
        Ok(response) => response
            .json::<serde_json::Value>()
            .await
            .ok()
            .map(|body| mecha10_core::ollama::parse_installed_model_names(&body))
            .unwrap_or_default(),
        Err(e) => {
            eprintln!("⚠️  Warning: Could not reach Ollama at {base_url}: {e}");
            eprintln!("   Local LLM nodes will fail until Ollama is running and the model is pulled.");
            return Ok(());
        }
    };

    for (node_name, model) in wanted_models {
        let already_installed = installed
            .iter()
            .any(|name| name == &model || name.starts_with(&format!("{model}:")));
        if already_installed {
            continue;
        }

        println!("📦 Pulling Ollama model '{}' for node '{}'...", model, node_name);
        let pull_result = client
            .post(format!("{base_url}/api/pull"))
            .json(&serde_json::json!({ "model": model, "stream": false }))
            .send()
            .await;

        match pull_result {
            Ok(response) if response.status().is_success() => println!("✅ Ollama model '{}' ready", model),
            Ok(response) => {
                let status = response.status();
                let body = response.text().await.unwrap_or_default();
                eprintln!("⚠️  Warning: Ollama rejected pulling '{}': {} {}", model, status, body);
                eprintln!(
                    "   You can manually pull with: docker compose exec ollama ollama pull {}",
                    model
                );
            }
            Err(e) => {
                eprintln!("⚠️  Warning: Failed to pull Ollama model '{}': {}", model, e);
                eprintln!(
                    "   You can manually pull with: docker compose exec ollama ollama pull {}",
                    model
                );
            }
        }
    }

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use serial_test::serial;

    /// Proves `parse_node_config` delegates to the real [`split_base_and_custom`] rather than
    /// hardcoding its own opinion about which of `model`/`model.provider` live on the shared
    /// base vs. under `custom` - it reads whatever the real split function actually produces
    /// for a given section, so a future field move between the two can't silently break this
    /// parser the way it broke the hand-rolled version twice before.
    #[test]
    #[serial]
    fn parse_node_config_delegates_to_the_real_base_and_custom_split() {
        std::env::remove_var("MECHA10_ENVIRONMENT");

        let dev_section = serde_json::json!({
            "custom": {
                "model": { "name": "hf.co/bartowski/Qwen2.5-1.5B-Instruct-GGUF:Q4_K_M", "provider": "local" },
            },
            "relay_topics": ["/x"],
        });
        let content = serde_json::json!({ "dev": dev_section.clone(), "production": dev_section.clone() }).to_string();

        let config = parse_node_config(&content).expect("config should parse");
        let merged = split_base_and_custom(&dev_section).expect("split should succeed");

        assert_eq!(
            config.model.as_ref().and_then(|m| m.name.clone()),
            merged["model"]["name"].as_str().map(|s| s.to_string())
        );
        assert_eq!(
            config.model.and_then(|m| m.provider),
            merged["model"]["provider"].as_str().map(|s| s.to_string())
        );
    }

    /// `parse_node_config` must correctly read `model` for a config shaped like
    /// `object-detector`'s real `config.json` (`packages/nodes/object-detector/configs/
    /// config.json`) today - `model` nested under `custom` alongside other node-specific
    /// fields. This doesn't hardcode an opinion that nested-under-`custom` is the "right"
    /// shape (see `parse_node_config_delegates_to_the_real_base_and_custom_split` below for a
    /// shape-agnostic test) - it just proves this one real, currently-shipping shape still
    /// works through the delegation to [`split_base_and_custom`].
    #[test]
    #[serial]
    fn parse_node_config_reads_model_nested_under_custom_like_object_detector() {
        std::env::remove_var("MECHA10_ENVIRONMENT");

        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": []
            }
        }"#;

        let config = parse_node_config(content).expect("config should parse");
        assert_eq!(config.model.and_then(|m| m.name), Some("yolov8n".to_string()));
    }

    /// Same regression as above, for `image-classifier`'s real shape (`packages/nodes/
    /// image-classifier/configs/config.json`), covering the `production` section too.
    #[test]
    #[serial]
    fn parse_node_config_reads_model_nested_under_custom_like_image_classifier_in_production() {
        std::env::set_var("MECHA10_ENVIRONMENT", "production");

        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": []
            }
        }"#;

        let config = parse_node_config(content).expect("config should parse");
        std::env::remove_var("MECHA10_ENVIRONMENT");

        assert_eq!(config.model.and_then(|m| m.name), Some("mobilenet-v2".to_string()));
    }

    /// Same as above, for `llm-command`'s Ollama-backed fields (`model.provider`/`model.name`)
    /// instead of the vision-model-only shape.
    #[test]
    #[serial]
    fn parse_node_config_reads_ollama_provider_and_model_name_nested_under_custom() {
        std::env::remove_var("MECHA10_ENVIRONMENT");

        let content = r#"{
            "dev": {
                "custom": {
                    "model": { "name": "hf.co/bartowski/Qwen2.5-1.5B-Instruct-GGUF:Q4_K_M", "provider": "local" }
                }
            },
            "production": {
                "custom": {
                    "model": { "name": "hf.co/bartowski/Qwen2.5-1.5B-Instruct-GGUF:Q4_K_M", "provider": "local" }
                }
            }
        }"#;

        let config = parse_node_config(content).expect("config should parse");
        let model = config.model.expect("model should be present");
        assert_eq!(model.provider.as_deref(), Some("local"));
        assert_eq!(
            model.name,
            Some("hf.co/bartowski/Qwen2.5-1.5B-Instruct-GGUF:Q4_K_M".to_string())
        );
    }

    /// A node with no `model` field at all (e.g. `speaker`) must still parse cleanly with no
    /// model declared - `custom` defaulting via `#[serde(default)]` must not turn a missing
    /// `custom` block, or a `custom` block with no `model` key, into a parse failure.
    #[test]
    #[serial]
    fn parse_node_config_returns_no_model_for_a_node_with_no_model_field() {
        std::env::remove_var("MECHA10_ENVIRONMENT");

        let content = r#"{
            "dev": { "custom": { "interval_secs": 5 } },
            "production": { "custom": { "interval_secs": 1 } }
        }"#;

        let config = parse_node_config(content).expect("config should parse");
        assert!(config.model.is_none());
    }
}