Skip to main content

rusty_stack/platform/
registry.rs

1//! Component registry and installed component detection.
2//!
3//! Provides:
4//! - **Component registry** — 19 known components with display names and installer mappings
5//!   (VAL-PLAT-009, VAL-PLAT-010)
6//! - **Installed component detection** — path-based, Python module, and git-based strategies
7//!   (VAL-PLAT-011, VAL-PLAT-012, VAL-PLAT-013)
8//! - **Version querying** — per-component-type version retrieval
9//!   (VAL-PLAT-014)
10//! - **Batch Python operations** — single-subprocess module detection and version queries
11//!   (VAL-PLAT-015, VAL-PLAT-016)
12
13use std::collections::HashMap;
14use std::path::{Path, PathBuf};
15use std::process::Command;
16
17use serde::{Deserialize, Serialize};
18
19use super::environment::resolve_user_home;
20
21// ===========================================================================
22// Public Types
23// ===========================================================================
24
25/// How a component's installation status is detected.
26#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
27pub enum DetectionMethod {
28    /// Check for files/paths on the filesystem.
29    PathBased,
30    /// Try importing a Python module via `importlib`.
31    PythonModule,
32    /// Check for a `.git` directory under an expected clone path.
33    GitBased,
34    /// Check for a command on `PATH`.
35    CommandBased,
36}
37
38/// Metadata for a known component in the registry.
39#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
40pub struct ComponentInfo {
41    /// Canonical component identifier (e.g., `"pytorch"`, `"flash-attn"`).
42    pub id: String,
43    /// Human-readable display name (e.g., `"PyTorch"`, `"Flash Attention"`).
44    pub display_name: String,
45    /// How to detect whether this component is installed.
46    pub detection_method: DetectionMethod,
47    /// Installer script name (relative to scripts directory).
48    pub installer_script: String,
49    /// Python import name, if this is a Python module.
50    pub python_import: Option<String>,
51    /// Expected clone directory name (relative to user home), for git-based components.
52    pub clone_dir: Option<String>,
53}
54
55/// Result of querying a component's version.
56#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
57pub struct VersionInfo {
58    /// The component identifier.
59    pub component_id: String,
60    /// Version string, or `"not installed"` / `"unknown"`.
61    pub version: String,
62}
63
64// ===========================================================================
65// Component Registry (VAL-PLAT-009, VAL-PLAT-010)
66// ===========================================================================
67
68/// Returns the canonical list of all 19 known components.
69pub fn known_components() -> &'static [ComponentInfo] {
70    // Built once, reused across calls.
71    static COMPONENTS: std::sync::OnceLock<Vec<ComponentInfo>> = std::sync::OnceLock::new();
72    COMPONENTS.get_or_init(|| {
73        vec![
74            // Foundation
75            ComponentInfo {
76                id: "rocm".into(),
77                display_name: "ROCm".into(),
78                detection_method: DetectionMethod::PathBased,
79                installer_script: "install_rocm.sh".into(),
80                python_import: None,
81                clone_dir: None,
82            },
83            ComponentInfo {
84                id: "rocm-smi".into(),
85                display_name: "ROCm SMI".into(),
86                detection_method: DetectionMethod::CommandBased,
87                installer_script: String::new(), // managed by ROCm
88                python_import: None,
89                clone_dir: None,
90            },
91            // Core
92            ComponentInfo {
93                id: "pytorch".into(),
94                display_name: "PyTorch".into(),
95                detection_method: DetectionMethod::PythonModule,
96                installer_script: "install_pytorch_rocm.sh".into(),
97                python_import: Some("torch".into()),
98                clone_dir: None,
99            },
100            ComponentInfo {
101                id: "triton".into(),
102                display_name: "Triton".into(),
103                detection_method: DetectionMethod::PythonModule,
104                installer_script: "install_triton_multi.sh".into(),
105                python_import: Some("triton".into()),
106                clone_dir: None,
107            },
108            ComponentInfo {
109                id: "onnx".into(),
110                display_name: "ONNX Runtime".into(),
111                detection_method: DetectionMethod::PythonModule,
112                installer_script: "build_onnxruntime_multi.sh".into(),
113                python_import: Some("onnxruntime".into()),
114                clone_dir: None,
115            },
116            ComponentInfo {
117                id: "migraphx".into(),
118                display_name: "MIGraphX".into(),
119                detection_method: DetectionMethod::PythonModule,
120                installer_script: "install_migraphx_multi.sh".into(),
121                python_import: Some("migraphx".into()),
122                clone_dir: None,
123            },
124            ComponentInfo {
125                id: "flash-attn".into(),
126                display_name: "Flash Attention".into(),
127                detection_method: DetectionMethod::PythonModule,
128                installer_script: "install_flash_attention_ck.sh".into(),
129                python_import: Some("flash_attn".into()),
130                clone_dir: None,
131            },
132            ComponentInfo {
133                id: "mpi4py".into(),
134                display_name: "MPI4Py".into(),
135                detection_method: DetectionMethod::PythonModule,
136                installer_script: "install_mpi4py.sh".into(),
137                python_import: Some("mpi4py".into()),
138                clone_dir: None,
139            },
140            ComponentInfo {
141                id: "megatron".into(),
142                display_name: "Megatron-LM".into(),
143                detection_method: DetectionMethod::GitBased,
144                installer_script: "install_megatron.sh".into(),
145                python_import: None,
146                clone_dir: Some("Megatron-LM".into()),
147            },
148            ComponentInfo {
149                id: "aiter".into(),
150                display_name: "AITER".into(),
151                detection_method: DetectionMethod::PythonModule,
152                installer_script: "install_aiter.sh".into(),
153                python_import: Some("aiter".into()),
154                clone_dir: None,
155            },
156            // Extensions
157            ComponentInfo {
158                id: "deepspeed".into(),
159                display_name: "DeepSpeed".into(),
160                detection_method: DetectionMethod::PythonModule,
161                installer_script: "install_deepspeed.sh".into(),
162                python_import: Some("deepspeed".into()),
163                clone_dir: None,
164            },
165            ComponentInfo {
166                id: "vllm".into(),
167                display_name: "vLLM".into(),
168                detection_method: DetectionMethod::PythonModule,
169                installer_script: "install_vllm_multi.sh".into(),
170                python_import: Some("vllm".into()),
171                clone_dir: None,
172            },
173            ComponentInfo {
174                id: "bitsandbytes".into(),
175                display_name: "bitsandbytes".into(),
176                detection_method: DetectionMethod::PythonModule,
177                installer_script: "install_bitsandbytes_multi.sh".into(),
178                python_import: Some("bitsandbytes".into()),
179                clone_dir: None,
180            },
181            ComponentInfo {
182                id: "wandb".into(),
183                display_name: "Weights & Biases".into(),
184                detection_method: DetectionMethod::PythonModule,
185                installer_script: "install_wandb.sh".into(),
186                python_import: Some("wandb".into()),
187                clone_dir: None,
188            },
189            // Extensions
190            ComponentInfo {
191                id: "llama-cpp".into(),
192                display_name: "llama.cpp (HIP)".into(),
193                detection_method: DetectionMethod::CommandBased,
194                installer_script: String::new(), // Native Rust installer
195                python_import: None,
196                clone_dir: None,
197            },
198            // UI/UX
199            ComponentInfo {
200                id: "comfyui".into(),
201                display_name: "ComfyUI".into(),
202                detection_method: DetectionMethod::GitBased,
203                installer_script: "install_comfyui.sh".into(),
204                python_import: None,
205                clone_dir: Some("ComfyUI".into()),
206            },
207            ComponentInfo {
208                id: "vllm-studio".into(),
209                display_name: "vLLM Studio".into(),
210                detection_method: DetectionMethod::GitBased,
211                installer_script: "install_vllm_studio.sh".into(),
212                python_import: None,
213                clone_dir: Some("vllm-studio".into()),
214            },
215            ComponentInfo {
216                id: "textgen".into(),
217                display_name: "text-generation-webui".into(),
218                detection_method: DetectionMethod::GitBased,
219                installer_script: "install_textgen.sh".into(),
220                python_import: None,
221                clone_dir: Some("text-generation-webui".into()),
222            },
223            // Environment
224            ComponentInfo {
225                id: "permanent-env".into(),
226                display_name: "ML Stack Environment".into(),
227                detection_method: DetectionMethod::PathBased,
228                installer_script: String::new(), // configured by installer
229                python_import: None,
230                clone_dir: None,
231            },
232        ]
233    })
234}
235
236/// Check whether a component ID is recognized (VAL-PLAT-009).
237pub fn is_known_component(id: &str) -> bool {
238    known_components().iter().any(|c| c.id == id)
239}
240
241/// Look up a component by ID.
242pub fn get_component(id: &str) -> Option<&'static ComponentInfo> {
243    known_components().iter().find(|c| c.id == id)
244}
245
246/// Return the human-readable display name for a component ID.
247///
248/// Unknown IDs are returned as-is (passthrough).
249/// VAL-PLAT-010: The mapping must be bijective for all 19 known components.
250pub fn display_name(id: &str) -> String {
251    get_component(id)
252        .map(|c| c.display_name.clone())
253        .unwrap_or_else(|| id.to_string())
254}
255
256/// Build a lookup map from display name back to component ID.
257///
258/// Used to verify bijectivity of the display name mapping.
259pub fn display_name_to_id_map() -> HashMap<String, String> {
260    known_components()
261        .iter()
262        .map(|c| (c.display_name.clone(), c.id.clone()))
263        .collect()
264}
265
266// ===========================================================================
267// Installed Component Detection
268// ===========================================================================
269
270/// Check whether a specific component is installed.
271///
272/// Dispatches to the appropriate detection strategy based on the component's
273/// `detection_method`:
274///
275/// - **PathBased**: ROCm checks version file + `rocminfo`; `permanent-env` checks `~/.mlstack_env`.
276/// - **CommandBased**: Checks that the binary exists and executes.
277/// - **PythonModule**: Uses `importlib.util.find_spec` in candidate interpreters.
278/// - **GitBased**: Checks for `.git` under expected clone path.
279pub fn is_component_installed(id: &str) -> bool {
280    is_component_installed_with_home(id, &resolve_user_home())
281}
282
283/// Check whether a component is installed using an explicit home directory.
284///
285/// This is the testable version that doesn't depend on global state.
286pub fn is_component_installed_with_home(id: &str, home: &Path) -> bool {
287    let Some(info) = get_component(id) else {
288        return false;
289    };
290
291    match info.detection_method {
292        DetectionMethod::PathBased => detect_path_based(info, home),
293        DetectionMethod::CommandBased => detect_command_based(info, home),
294        DetectionMethod::PythonModule => detect_python_module_single(info),
295        DetectionMethod::GitBased => detect_git_based(info, home),
296    }
297}
298
299// -----------------------------------------------------------------------
300// Path-based detection (VAL-PLAT-011)
301// -----------------------------------------------------------------------
302
303/// Path-based detection for ROCm and permanent-env.
304fn detect_path_based(info: &ComponentInfo, home: &Path) -> bool {
305    match info.id.as_str() {
306        "rocm" => {
307            // ROCm: version file must exist AND rocminfo must be runnable
308            let version_file_exists = Path::new("/opt/rocm/.info/version").exists();
309            let rocminfo_works = Command::new("rocminfo")
310                .arg("--version")
311                .output()
312                .map(|o| o.status.success())
313                .unwrap_or(false);
314            version_file_exists && rocminfo_works
315        }
316        "permanent-env" => {
317            // Check for marker string in ~/.mlstack_env
318            let env_file = home.join(".mlstack_env");
319            if !env_file.exists() {
320                return false;
321            }
322            std::fs::read_to_string(&env_file)
323                .map(|content| content.contains("MLSTACK_PYTHON_BIN"))
324                .unwrap_or(false)
325        }
326        _ => false,
327    }
328}
329
330// -----------------------------------------------------------------------
331// Command-based detection
332// -----------------------------------------------------------------------
333
334/// Command-based detection for tools like rocm-smi.
335fn detect_command_based(info: &ComponentInfo, home: &Path) -> bool {
336    match info.id.as_str() {
337        "rocm-smi" => Command::new("rocm-smi")
338            .arg("--version")
339            .output()
340            .map(|o| o.status.success())
341            .unwrap_or(false),
342        "llama-cpp" => {
343            // Prefer installer contract: binary under ~/.mlstack/components/llama-cpp/bin
344            // must be functional. Fall back to PATH lookup for manual/system installs.
345            let home_str = home.to_string_lossy();
346            crate::installers::components::llama_cpp::is_llama_cli_functional(&home_str)
347                || Command::new("llama-cli")
348                    .arg("--help")
349                    .output()
350                    .map(|o| o.status.success())
351                    .unwrap_or(false)
352        }
353        _ => false,
354    }
355}
356
357// -----------------------------------------------------------------------
358// Python module detection (VAL-PLAT-012)
359// -----------------------------------------------------------------------
360
361/// Detect a single Python module component by trying importlib in candidate interpreters.
362fn detect_python_module_single(info: &ComponentInfo) -> bool {
363    let Some(ref import_name) = info.python_import else {
364        return false;
365    };
366
367    let interpreters = super::environment::python_interpreters();
368    for python in &interpreters {
369        if try_python_import(python, import_name) {
370            return true;
371        }
372    }
373    false
374}
375
376/// Try importing a Python module using the given interpreter.
377fn try_python_import(python: &Path, module: &str) -> bool {
378    let script = format!(
379        "import importlib.util; spec = importlib.util.find_spec('{module}'); exit(0 if spec else 1)"
380    );
381    Command::new(python)
382        .arg("-c")
383        .arg(&script)
384        .output()
385        .map(|o| o.status.success())
386        .unwrap_or(false)
387}
388
389// -----------------------------------------------------------------------
390// Git-based detection (VAL-PLAT-013)
391// -----------------------------------------------------------------------
392
393/// Git-based detection for ComfyUI, vllm-studio, textgen.
394fn detect_git_based(info: &ComponentInfo, home: &Path) -> bool {
395    let Some(ref clone_dir) = info.clone_dir else {
396        return false;
397    };
398
399    let git_dir = home.join(clone_dir).join(".git");
400    git_dir.exists()
401}
402
403// ===========================================================================
404// Version Querying (VAL-PLAT-014)
405// ===========================================================================
406
407/// Query the version of a single installed component.
408///
409/// Returns:
410/// - Semver string for Python modules
411/// - Git short-hash + subject for git-based components
412/// - `"installed"` for command-only tools
413/// - `"not installed"` when absent
414pub fn get_version(id: &str) -> String {
415    get_version_with_home(id, &resolve_user_home())
416}
417
418/// Query version with explicit home directory (testable).
419pub fn get_version_with_home(id: &str, home: &Path) -> String {
420    let Some(info) = get_component(id) else {
421        return "unknown".to_string();
422    };
423
424    match info.detection_method {
425        DetectionMethod::PathBased => get_version_path_based(info, home),
426        DetectionMethod::CommandBased => get_version_command_based(info, home),
427        DetectionMethod::PythonModule => get_version_python_single(info),
428        DetectionMethod::GitBased => get_version_git(info, home),
429    }
430}
431
432/// Get version for path-based components.
433fn get_version_path_based(info: &ComponentInfo, home: &Path) -> String {
434    match info.id.as_str() {
435        "rocm" => {
436            // Try version file first, then rocminfo fallback
437            if let Ok(content) = std::fs::read_to_string("/opt/rocm/.info/version") {
438                let v = content.trim().to_string();
439                if !v.is_empty() {
440                    return v;
441                }
442            }
443            if let Ok(output) = Command::new("rocminfo").arg("--version").output() {
444                if output.status.success() {
445                    let stdout = String::from_utf8_lossy(&output.stdout);
446                    if let Some(line) = stdout.lines().next() {
447                        return line.trim().to_string();
448                    }
449                }
450            }
451            "unknown".to_string()
452        }
453        "permanent-env" => {
454            if is_component_installed_with_home("permanent-env", home) {
455                "installed".to_string()
456            } else {
457                "not installed".to_string()
458            }
459        }
460        _ => "unknown".to_string(),
461    }
462}
463
464/// Get version for command-based components.
465fn get_version_command_based(info: &ComponentInfo, home: &Path) -> String {
466    match info.id.as_str() {
467        "rocm-smi" => {
468            if let Ok(output) = Command::new("rocm-smi").arg("--version").output() {
469                if output.status.success() {
470                    let stdout = String::from_utf8_lossy(&output.stdout);
471                    if let Some(line) = stdout.lines().next() {
472                        return extract_semver(line.trim());
473                    }
474                }
475            }
476            "not installed".to_string()
477        }
478        "llama-cpp" => {
479            if let Ok(output) = Command::new("llama-cli").arg("--version").output() {
480                if output.status.success() {
481                    let stdout = String::from_utf8_lossy(&output.stdout);
482                    let stderr = String::from_utf8_lossy(&output.stderr);
483                    let combined = format!("{stdout}\n{stderr}");
484                    let parsed = extract_semver(combined.trim());
485                    if parsed != "unknown" {
486                        return parsed;
487                    }
488                    return "installed".to_string();
489                }
490            }
491            if detect_command_based(info, home) {
492                return "installed".to_string();
493            }
494            "not installed".to_string()
495        }
496        _ => "unknown".to_string(),
497    }
498}
499
500/// Get version for a single Python module component.
501fn get_version_python_single(info: &ComponentInfo) -> String {
502    let Some(ref import_name) = info.python_import else {
503        return "unknown".to_string();
504    };
505
506    let interpreters = super::environment::python_interpreters();
507    for python in &interpreters {
508        let script = format!(
509            "try:\n    import {import_name}\n    ver = getattr({import_name}, '__version__', None)\n    if ver is None:\n        import importlib.metadata as im\n        for _n in ['{import_name}', '{import_name}'.replace('-', '_'), 'amd_' + '{import_name}'.replace('-', '_'), 'amd-' + '{import_name}']:\n            try:\n                ver = im.version(_n)\n                break\n            except Exception:\n                pass\n    if ver is not None:\n        print(ver)\nexcept (ImportError, AttributeError):\n    pass",
510            import_name = import_name
511        );
512        if let Ok(output) = Command::new(python).arg("-c").arg(&script).output() {
513            if output.status.success() {
514                let stdout = String::from_utf8_lossy(&output.stdout);
515                // Take the last non-empty line that doesn't look like a Python
516                // library warning (some libs like accelerate write warnings to
517                // stdout instead of stderr).
518                let version = stdout
519                    .lines()
520                    .map(|l| l.trim())
521                    .filter(|l| !l.is_empty() && !l.contains("[WARNING]") && !l.contains("[ERROR]"))
522                    .last()
523                    .unwrap_or("")
524                    .to_string();
525                if !version.is_empty() {
526                    return version;
527                }
528            }
529        }
530    }
531    "not installed".to_string()
532}
533
534/// Get version for git-based components (short hash + subject).
535/// Also stores the short hash for update comparison.
536fn get_version_git(info: &ComponentInfo, home: &Path) -> String {
537    let Some(ref clone_dir) = info.clone_dir else {
538        return "unknown".to_string();
539    };
540
541    let repo_path = home.join(clone_dir);
542
543    if !repo_path.join(".git").exists() {
544        return "not installed".to_string();
545    }
546
547    // Try to get a tag-based version first (git describe)
548    if let Ok(output) = Command::new("git")
549        .args(["-C", repo_path.to_str().unwrap_or(".")])
550        .args(["describe", "--tags", "--abbrev=0"])
551        .output()
552    {
553        if output.status.success() {
554            let stdout = String::from_utf8_lossy(&output.stdout);
555            let tag = stdout.trim().to_string();
556            if !tag.is_empty() {
557                return tag;
558            }
559        }
560    }
561
562    // Fallback to short hash + subject
563    if let Ok(output) = Command::new("git")
564        .args(["-C", repo_path.to_str().unwrap_or(".")])
565        .args(["log", "-1", "--format=%h %s"])
566        .output()
567    {
568        if output.status.success() {
569            let stdout = String::from_utf8_lossy(&output.stdout);
570            let line = stdout.trim().to_string();
571            if !line.is_empty() {
572                return line;
573            }
574        }
575    }
576
577    "installed".to_string()
578}
579
580// ===========================================================================
581// Batch Python Operations (VAL-PLAT-015, VAL-PLAT-016)
582// ===========================================================================
583
584/// Extract a semantic version substring from a version output string.
585///
586/// Handles strings like:
587/// - `"ROCM-SMI version: 4.0.0+unknown"` → `"4.0.0"`
588/// - `"4.0.0+unknown"` → `"4.0.0"`
589/// - `"v2.6.0"` → `"2.6.0"`
590/// - `"2.6.0"` → `"2.6.0"`
591fn extract_semver(raw: &str) -> String {
592    // Strip common prefixes
593    let s = raw.trim().trim_start_matches('v').to_string();
594
595    // Try to find a semver pattern (X.Y.Z or X.Y)
596    let re = regex::Regex::new(r"(\d+\.\d+(?:\.\d+)?)").unwrap();
597    if let Some(caps) = re.captures(&s) {
598        return caps[1].to_string();
599    }
600
601    // No match — return cleaned string
602    s
603}
604
605/// Check if a git-based component has updates available.
606///
607/// Performs `git fetch` then compares local HEAD with remote HEAD.
608/// Returns `true` if remote has new commits.
609pub fn git_component_has_updates(id: &str) -> bool {
610    let Some(info) = get_component(id) else {
611        return false;
612    };
613    let Some(ref clone_dir) = info.clone_dir else {
614        return false;
615    };
616    let home = resolve_user_home();
617    let repo_path = home.join(clone_dir);
618
619    if !repo_path.join(".git").exists() {
620        return false;
621    }
622
623    // Fetch without output
624    let _ = Command::new("git")
625        .args(["-C", repo_path.to_str().unwrap_or(".")])
626        .args(["fetch", "--quiet"])
627        .output();
628
629    // Compare local HEAD with remote HEAD
630    if let Ok(output) = Command::new("git")
631        .args(["-C", repo_path.to_str().unwrap_or(".")])
632        .args(["log", "HEAD..@{u}", "--oneline"])
633        .output()
634    {
635        if output.status.success() {
636            let stdout = String::from_utf8_lossy(&output.stdout);
637            return !stdout.trim().is_empty();
638        }
639    }
640
641    false
642}
643
644/// Python module component IDs and their import names.
645fn python_module_mappings() -> &'static [(String, String)] {
646    static MAPPINGS: std::sync::OnceLock<Vec<(String, String)>> = std::sync::OnceLock::new();
647    MAPPINGS.get_or_init(|| {
648        known_components()
649            .iter()
650            .filter(|c| c.detection_method == DetectionMethod::PythonModule)
651            .filter_map(|c| {
652                c.python_import
653                    .as_ref()
654                    .map(|imp| (c.id.clone(), imp.clone()))
655            })
656            .collect()
657    })
658}
659
660/// Detect all installed Python ML modules using a single subprocess.
661///
662/// VAL-PLAT-016: Must use a single Python process invocation.
663/// Returns a set of component IDs that are installed.
664pub fn detect_python_modules() -> Vec<String> {
665    let interpreters = super::environment::python_interpreters();
666    detect_python_modules_with_interpreters(&interpreters)
667}
668
669/// Detect Python modules using the provided interpreter list (testable).
670pub fn detect_python_modules_with_interpreters(interpreters: &[PathBuf]) -> Vec<String> {
671    let mappings = python_module_mappings();
672    if mappings.is_empty() || interpreters.is_empty() {
673        return Vec::new();
674    }
675
676    // Build a single Python script that checks all modules at once
677    let mut script = String::from("mods = [\n");
678    for (comp_id, import_name) in mappings {
679        script.push_str(&format!("    ({comp_id:?}, {import_name:?}),\n"));
680    }
681    script.push_str(
682        "]\nfor comp_id, import_name in mods:\n    try:\n        __import__(import_name)\n        print(comp_id)\n    except ImportError:\n        pass\n",
683    );
684
685    for python in interpreters {
686        if let Ok(output) = Command::new(python).arg("-c").arg(&script).output() {
687            tracing::debug!(
688                "Python detection: Interpreter {} returned status: {}",
689                python.display(),
690                output.status
691            );
692            if output.status.success() {
693                let stdout = String::from_utf8_lossy(&output.stdout);
694                let detected: Vec<String> = stdout
695                    .lines()
696                    .map(|l| l.trim().to_string())
697                    .filter(|l| !l.is_empty() && is_known_component(&l))
698                    .collect();
699                if !detected.is_empty() {
700                    return detected;
701                }
702            }
703        }
704    }
705
706    Vec::new()
707}
708
709/// Query versions of multiple Python modules in a single subprocess.
710///
711/// VAL-PLAT-015: Must use a single subprocess invocation.
712/// Returns `component_id=version` pairs for all requested components.
713pub fn get_versions_batch(component_ids: &[&str]) -> Vec<VersionInfo> {
714    let interpreters = super::environment::python_interpreters();
715    get_versions_batch_with_interpreters(component_ids, &interpreters)
716}
717
718/// Batch version query with provided interpreters (testable).
719pub fn get_versions_batch_with_interpreters(
720    component_ids: &[&str],
721    interpreters: &[PathBuf],
722) -> Vec<VersionInfo> {
723    let mappings = python_module_mappings();
724
725    if interpreters.is_empty() {
726        return component_ids
727            .iter()
728            .map(|id| VersionInfo {
729                component_id: (*id).to_string(),
730                version: "unknown".to_string(),
731            })
732            .collect();
733    }
734
735    // Build the Python script for batch version query
736    let mut script = String::from("import sys\nmapping = {\n");
737    for (comp_id, import_name) in mappings {
738        script.push_str(&format!(
739            "    {:?}: ({:?}, '__version__'),\n",
740            comp_id, import_name
741        ));
742    }
743    script.push_str(
744        "}\nfor comp_id in sys.argv[1:]:\n    if comp_id not in mapping:\n        print(f'{comp_id}=unknown')\n        continue\n    mod_name, ver_attr = mapping[comp_id]\n    try:\n        mod = __import__(mod_name)\n        ver = getattr(mod, ver_attr, 'unknown')\n        if ver == 'unknown':\n            try:\n                import importlib.metadata as im\n                for _n in [mod_name, mod_name.replace('-', '_'), 'amd_' + mod_name.replace('-', '_'), 'amd-' + mod_name]:\n                    try:\n                        ver = im.version(_n)\n                        break\n                    except Exception:\n                        pass\n            except Exception:\n                pass\n        print(f'{comp_id}={ver}')\n    except ImportError:\n        print(f'{comp_id}=not installed')\n",
745    );
746
747    for python in interpreters {
748        if let Ok(output) = Command::new(python)
749            .arg("-c")
750            .arg(&script)
751            .args(component_ids)
752            .output()
753        {
754            if output.status.success() {
755                let stdout = String::from_utf8_lossy(&output.stdout);
756                return stdout
757                    .lines()
758                    .filter_map(|line| {
759                        let line = line.trim();
760                        if line.is_empty() {
761                            return None;
762                        }
763                        let (id, version) = line.split_once('=')?;
764                        // Skip lines where id contains characters typical of
765                        // Python library warnings that leak to stdout instead
766                        // of stderr (e.g. "[2026-05-18 ...] [WARNING] ...").
767                        if id.contains('[') || id.contains(' ') || id.contains(']') {
768                            return None;
769                        }
770                        Some(VersionInfo {
771                            component_id: id.to_string(),
772                            version: version.to_string(),
773                        })
774                    })
775                    .collect();
776            }
777        }
778    }
779
780    // Fallback: return "unknown" for all
781    component_ids
782        .iter()
783        .map(|id| VersionInfo {
784            component_id: (*id).to_string(),
785            version: "unknown".to_string(),
786        })
787        .collect()
788}
789
790/// Detect all installed components using the full detection pipeline.
791///
792/// Combines path-based, command-based, Python module (batch), and git-based
793/// detection into a single scan.
794pub fn detect_all_installed() -> Vec<String> {
795    let home = resolve_user_home();
796    detect_all_installed_with_home(&home)
797}
798
799/// Detect all installed components with explicit home directory (testable).
800pub fn detect_all_installed_with_home(home: &Path) -> Vec<String> {
801    let mut installed = Vec::new();
802
803    // Path-based: ROCm
804    if is_component_installed_with_home("rocm", home) {
805        installed.push("rocm".to_string());
806    }
807
808    // Command-based components
809    for info in known_components() {
810        if info.detection_method == DetectionMethod::CommandBased
811            && is_component_installed(&info.id)
812        {
813            installed.push(info.id.clone());
814        }
815    }
816
817    // Python modules: batch detection
818    let python_installed = detect_python_modules();
819    installed.extend(python_installed);
820
821    // Git-based
822    for info in known_components() {
823        if info.detection_method == DetectionMethod::GitBased && detect_git_based(info, home) {
824            installed.push(info.id.clone());
825        }
826    }
827
828    // Path-based: permanent-env
829    if is_component_installed_with_home("permanent-env", home) {
830        installed.push("permanent-env".to_string());
831    }
832
833    installed.sort();
834    installed.dedup();
835    installed
836}
837
838// ===========================================================================
839// Tests
840// ===========================================================================
841
842#[cfg(test)]
843mod tests {
844    use super::*;
845
846    // -----------------------------------------------------------------------
847    // VAL-PLAT-009: Component registry contains all 19 known components
848    // -----------------------------------------------------------------------
849
850    #[test]
851    fn test_registry_has_exactly_19_components() {
852        let components = known_components();
853        assert_eq!(
854            components.len(),
855            19,
856            "Registry must contain exactly 19 known components, found {}",
857            components.len()
858        );
859    }
860
861    #[test]
862    fn test_registry_contains_all_expected_ids() {
863        let expected_ids = [
864            "rocm",
865            "pytorch",
866            "triton",
867            "megatron",
868            "deepspeed",
869            "vllm",
870            "aiter",
871            "onnx",
872            "bitsandbytes",
873            "migraphx",
874            "flash-attn",
875            "mpi4py",
876            "wandb",
877            "comfyui",
878            "vllm-studio",
879            "textgen",
880            "rocm-smi",
881            "permanent-env",
882            "llama-cpp",
883        ];
884        let components = known_components();
885        for expected in &expected_ids {
886            assert!(
887                components.iter().any(|c| c.id == *expected),
888                "Missing component: {expected}"
889            );
890        }
891    }
892
893    #[test]
894    fn test_is_known_component_accepts_all_19() {
895        let ids = [
896            "rocm",
897            "pytorch",
898            "triton",
899            "megatron",
900            "deepspeed",
901            "vllm",
902            "aiter",
903            "onnx",
904            "bitsandbytes",
905            "migraphx",
906            "flash-attn",
907            "mpi4py",
908            "wandb",
909            "comfyui",
910            "vllm-studio",
911            "textgen",
912            "rocm-smi",
913            "permanent-env",
914            "llama-cpp",
915        ];
916        for id in &ids {
917            assert!(
918                is_known_component(id),
919                "is_known_component should accept '{id}'"
920            );
921        }
922    }
923
924    #[test]
925    fn test_is_known_component_rejects_unknown() {
926        assert!(!is_known_component("nonexistent"));
927        assert!(!is_known_component("cuda"));
928        assert!(!is_known_component("tensorflow"));
929        assert!(!is_known_component(""));
930        assert!(!is_known_component("PyTorch")); // case-sensitive
931    }
932
933    #[test]
934    fn test_get_component_returns_info_for_known() {
935        let info = get_component("pytorch").unwrap();
936        assert_eq!(info.id, "pytorch");
937        assert_eq!(info.display_name, "PyTorch");
938        assert_eq!(info.detection_method, DetectionMethod::PythonModule);
939        assert_eq!(info.python_import.as_deref(), Some("torch"));
940    }
941
942    #[test]
943    fn test_get_component_returns_none_for_unknown() {
944        assert!(get_component("nonexistent").is_none());
945    }
946
947    #[test]
948    fn test_all_components_have_non_empty_ids() {
949        for c in known_components() {
950            assert!(!c.id.is_empty(), "Component has empty id");
951        }
952    }
953
954    #[test]
955    fn test_all_component_ids_are_unique() {
956        let ids: Vec<&str> = known_components().iter().map(|c| c.id.as_str()).collect();
957        let unique: std::collections::HashSet<&str> = ids.iter().copied().collect();
958        assert_eq!(ids.len(), unique.len(), "Component IDs must be unique");
959    }
960
961    // -----------------------------------------------------------------------
962    // VAL-PLAT-010: Display name mapping is bijective for all 19 components
963    // -----------------------------------------------------------------------
964
965    #[test]
966    fn test_all_display_names_non_empty() {
967        for c in known_components() {
968            assert!(
969                !c.display_name.is_empty(),
970                "Display name for '{}' is empty",
971                c.id
972            );
973        }
974    }
975
976    #[test]
977    fn test_display_names_are_distinct() {
978        let names: Vec<&str> = known_components()
979            .iter()
980            .map(|c| c.display_name.as_str())
981            .collect();
982        let unique: std::collections::HashSet<&str> = names.iter().copied().collect();
983        assert_eq!(
984            names.len(),
985            unique.len(),
986            "Display names must all be distinct (bijective)"
987        );
988    }
989
990    #[test]
991    fn test_display_name_to_id_is_bijective() {
992        let map = display_name_to_id_map();
993        assert_eq!(
994            map.len(),
995            19,
996            "Display name to ID map must have exactly 19 entries"
997        );
998
999        // Verify round-trip: display_name(id) -> display_name -> id == original id
1000        for c in known_components() {
1001            let name = display_name(&c.id);
1002            let back_id = map.get(&name);
1003            assert_eq!(
1004                back_id,
1005                Some(&c.id),
1006                "Round-trip failed for '{}': display_name='{}'",
1007                c.id,
1008                name
1009            );
1010        }
1011    }
1012
1013    #[test]
1014    fn test_display_name_unknown_passthrough() {
1015        assert_eq!(display_name("nonexistent"), "nonexistent");
1016        assert_eq!(display_name("custom-tool"), "custom-tool");
1017    }
1018
1019    #[test]
1020    fn test_display_name_known_values() {
1021        assert_eq!(display_name("rocm"), "ROCm");
1022        assert_eq!(display_name("pytorch"), "PyTorch");
1023        assert_eq!(display_name("flash-attn"), "Flash Attention");
1024        assert_eq!(display_name("wandb"), "Weights & Biases");
1025        assert_eq!(display_name("comfyui"), "ComfyUI");
1026        assert_eq!(display_name("permanent-env"), "ML Stack Environment");
1027    }
1028
1029    // -----------------------------------------------------------------------
1030    // VAL-PLAT-011: Path-based detection (ROCm, rocm-smi, permanent-env)
1031    // -----------------------------------------------------------------------
1032
1033    #[test]
1034    fn test_rocm_detection_requires_both_version_file_and_rocminfo() {
1035        // On a real system with ROCm, this should return true
1036        // On a system without ROCm, this should return false
1037        // The test just verifies it doesn't panic
1038        let result = is_component_installed("rocm");
1039        // Just verify it runs without panic
1040        let _ = result;
1041    }
1042
1043    #[test]
1044    fn test_permanent_env_detection_checks_marker() {
1045        let dir = tempfile::tempdir().unwrap();
1046
1047        // Without marker string
1048        std::fs::write(dir.path().join(".mlstack_env"), "# empty\n").unwrap();
1049        assert!(!is_component_installed_with_home(
1050            "permanent-env",
1051            dir.path()
1052        ));
1053
1054        // With marker string
1055        std::fs::write(
1056            dir.path().join(".mlstack_env"),
1057            "export MLSTACK_PYTHON_BIN=/usr/bin/python3\n",
1058        )
1059        .unwrap();
1060        assert!(is_component_installed_with_home(
1061            "permanent-env",
1062            dir.path()
1063        ));
1064    }
1065
1066    #[test]
1067    fn test_permanent_env_not_installed_when_file_missing() {
1068        let dir = tempfile::tempdir().unwrap();
1069        assert!(!is_component_installed_with_home(
1070            "permanent-env",
1071            dir.path()
1072        ));
1073    }
1074
1075    // -----------------------------------------------------------------------
1076    // VAL-PLAT-012: Python module detection via importlib
1077    // -----------------------------------------------------------------------
1078
1079    #[test]
1080    fn test_python_module_detection_for_pytorch() {
1081        // On a real system with PyTorch installed, this should return true
1082        // On a system without PyTorch, it should return false
1083        // The test just verifies it doesn't panic
1084        let _ = is_component_installed("pytorch");
1085    }
1086
1087    #[test]
1088    fn test_detect_python_modules_with_interpreters_empty() {
1089        let result = detect_python_modules_with_interpreters(&[]);
1090        assert!(result.is_empty());
1091    }
1092
1093    #[test]
1094    fn test_detect_python_modules_uses_single_subprocess() {
1095        // Verify the function runs and returns a Vec<String>
1096        // The "single subprocess" guarantee is structural (one Command invocation)
1097        let result = detect_python_modules();
1098        // All entries should be valid component IDs
1099        for id in &result {
1100            assert!(is_known_component(id), "Unknown component ID: {id}");
1101        }
1102    }
1103
1104    // -----------------------------------------------------------------------
1105    // VAL-PLAT-013: Git-based detection (ComfyUI, vllm-studio, textgen)
1106    // -----------------------------------------------------------------------
1107
1108    #[test]
1109    fn test_git_based_detection_comfyui_with_git_dir() {
1110        let dir = tempfile::tempdir().unwrap();
1111
1112        // Create ComfyUI/.git
1113        std::fs::create_dir_all(dir.path().join("ComfyUI/.git")).unwrap();
1114        assert!(is_component_installed_with_home("comfyui", dir.path()));
1115    }
1116
1117    #[test]
1118    fn test_git_based_detection_comfyui_without_git_dir() {
1119        let dir = tempfile::tempdir().unwrap();
1120
1121        // ComfyUI dir exists but no .git
1122        std::fs::create_dir_all(dir.path().join("ComfyUI")).unwrap();
1123        assert!(!is_component_installed_with_home("comfyui", dir.path()));
1124    }
1125
1126    #[test]
1127    fn test_git_based_detection_vllm_studio() {
1128        let dir = tempfile::tempdir().unwrap();
1129
1130        std::fs::create_dir_all(dir.path().join("vllm-studio/.git")).unwrap();
1131        assert!(is_component_installed_with_home("vllm-studio", dir.path()));
1132    }
1133
1134    #[test]
1135    fn test_git_based_detection_textgen() {
1136        let dir = tempfile::tempdir().unwrap();
1137
1138        std::fs::create_dir_all(dir.path().join("text-generation-webui/.git")).unwrap();
1139        assert!(is_component_installed_with_home("textgen", dir.path()));
1140    }
1141
1142    // -----------------------------------------------------------------------
1143    // VAL-PLAT-014: Version querying returns correct format per component type
1144    // -----------------------------------------------------------------------
1145
1146    #[test]
1147    fn test_version_query_rocm_format() {
1148        let version = get_version("rocm");
1149        // Should be non-empty and not "unknown" on a system with ROCm
1150        assert!(!version.is_empty());
1151    }
1152
1153    #[test]
1154    fn test_version_query_permanent_env() {
1155        let dir = tempfile::tempdir().unwrap();
1156
1157        // Not installed
1158        let version = get_version_with_home("permanent-env", dir.path());
1159        assert_eq!(version, "not installed");
1160
1161        // Installed
1162        std::fs::write(
1163            dir.path().join(".mlstack_env"),
1164            "export MLSTACK_PYTHON_BIN=/usr/bin/python3\n",
1165        )
1166        .unwrap();
1167        let version = get_version_with_home("permanent-env", dir.path());
1168        assert_eq!(version, "installed");
1169    }
1170
1171    #[test]
1172    fn test_version_query_unknown_component() {
1173        assert_eq!(get_version("nonexistent"), "unknown");
1174    }
1175
1176    #[test]
1177    fn test_version_query_git_based_not_installed() {
1178        let dir = tempfile::tempdir().unwrap();
1179        let version = get_version_with_home("comfyui", dir.path());
1180        assert_eq!(version, "not installed");
1181    }
1182
1183    // -----------------------------------------------------------------------
1184    // VAL-PLAT-015: Batch Python version uses single subprocess
1185    // -----------------------------------------------------------------------
1186
1187    #[test]
1188    fn test_get_versions_batch_returns_results() {
1189        let ids = vec!["pytorch", "triton", "nonexistent-mod"];
1190        let results = get_versions_batch(&ids);
1191
1192        assert_eq!(results.len(), 3);
1193
1194        // Check that each result has the right component_id
1195        assert_eq!(results[0].component_id, "pytorch");
1196        assert_eq!(results[1].component_id, "triton");
1197        assert_eq!(results[2].component_id, "nonexistent-mod");
1198
1199        // nonexistent-mod should be "unknown" (not in mapping)
1200        assert_eq!(results[2].version, "unknown");
1201    }
1202
1203    #[test]
1204    fn test_get_versions_batch_with_interpreters_empty() {
1205        let ids = vec!["pytorch"];
1206        let results = get_versions_batch_with_interpreters(&ids, &[]);
1207        assert_eq!(results.len(), 1);
1208        assert_eq!(results[0].version, "unknown");
1209    }
1210
1211    #[test]
1212    fn test_get_versions_batch_empty_ids() {
1213        let results = get_versions_batch(&[]);
1214        assert!(results.is_empty());
1215    }
1216
1217    // -----------------------------------------------------------------------
1218    // VAL-PLAT-016: Batch Python module detection uses single subprocess
1219    // -----------------------------------------------------------------------
1220
1221    #[test]
1222    fn test_detect_all_installed_returns_valid_ids() {
1223        let installed = detect_all_installed();
1224        for id in &installed {
1225            assert!(
1226                is_known_component(id),
1227                "detect_all_installed returned unknown component: {id}"
1228            );
1229        }
1230    }
1231
1232    #[test]
1233    fn test_detect_all_installed_no_duplicates() {
1234        let installed = detect_all_installed();
1235        let unique: std::collections::HashSet<&String> = installed.iter().collect();
1236        assert_eq!(
1237            installed.len(),
1238            unique.len(),
1239            "detect_all_installed should not return duplicates"
1240        );
1241    }
1242
1243    // -----------------------------------------------------------------------
1244    // ComponentInfo serde roundtrip
1245    // -----------------------------------------------------------------------
1246
1247    #[test]
1248    fn test_component_info_serde_roundtrip() {
1249        let info = get_component("pytorch").unwrap();
1250        let json = serde_json::to_string(info).unwrap();
1251        let back: ComponentInfo = serde_json::from_str(&json).unwrap();
1252        assert_eq!(info, &back);
1253    }
1254
1255    #[test]
1256    fn test_detection_method_serde_roundtrip() {
1257        let methods = [
1258            DetectionMethod::PathBased,
1259            DetectionMethod::PythonModule,
1260            DetectionMethod::GitBased,
1261            DetectionMethod::CommandBased,
1262        ];
1263        for method in &methods {
1264            let json = serde_json::to_string(method).unwrap();
1265            let back: DetectionMethod = serde_json::from_str(&json).unwrap();
1266            assert_eq!(*method, back);
1267        }
1268    }
1269
1270    #[test]
1271    fn test_version_info_serde_roundtrip() {
1272        let info = VersionInfo {
1273            component_id: "pytorch".to_string(),
1274            version: "2.4.0".to_string(),
1275        };
1276        let json = serde_json::to_string(&info).unwrap();
1277        let back: VersionInfo = serde_json::from_str(&json).unwrap();
1278        assert_eq!(info, back);
1279    }
1280
1281    // -----------------------------------------------------------------------
1282    // Installer script mappings
1283    // -----------------------------------------------------------------------
1284
1285    #[test]
1286    fn test_python_components_have_import_names() {
1287        for c in known_components() {
1288            if c.detection_method == DetectionMethod::PythonModule {
1289                assert!(
1290                    c.python_import.is_some(),
1291                    "Python module component '{}' must have python_import",
1292                    c.id
1293                );
1294                assert!(
1295                    !c.python_import.as_ref().unwrap().is_empty(),
1296                    "python_import for '{}' must be non-empty",
1297                    c.id
1298                );
1299            }
1300        }
1301    }
1302
1303    #[test]
1304    fn test_git_components_have_clone_dirs() {
1305        for c in known_components() {
1306            if c.detection_method == DetectionMethod::GitBased {
1307                assert!(
1308                    c.clone_dir.is_some(),
1309                    "Git-based component '{}' must have clone_dir",
1310                    c.id
1311                );
1312                assert!(
1313                    !c.clone_dir.as_ref().unwrap().is_empty(),
1314                    "clone_dir for '{}' must be non-empty",
1315                    c.id
1316                );
1317            }
1318        }
1319    }
1320
1321    #[test]
1322    fn test_installer_scripts_for_installable_components() {
1323        // rocm-smi, permanent-env, and native Rust components don't have installer scripts
1324        let no_script = ["rocm-smi", "permanent-env", "llama-cpp"];
1325        for c in known_components() {
1326            if !no_script.contains(&c.id.as_str()) {
1327                assert!(
1328                    !c.installer_script.is_empty(),
1329                    "Component '{}' must have an installer script",
1330                    c.id
1331                );
1332            }
1333        }
1334    }
1335
1336    // -----------------------------------------------------------------------
1337    // Real system integration (manual verification helpers)
1338    // -----------------------------------------------------------------------
1339
1340    #[test]
1341    fn test_real_system_detect_all() {
1342        let installed = detect_all_installed();
1343        // On the real system, we expect at least ROCm to be installed
1344        // This test is informational — it just verifies the pipeline doesn't panic
1345        println!("Installed components: {:?}", installed);
1346        // Test passes if we get here without panicking
1347    }
1348
1349    #[test]
1350    fn test_real_system_version_queries() {
1351        // Query versions for all components — informational
1352        for c in known_components() {
1353            let version = get_version(&c.id);
1354            println!("{}: {}", display_name(&c.id), version);
1355        }
1356        // Test passes if we get here without panicking
1357    }
1358}