Skip to main content

rusty_stack/
config.rs

1use anyhow::{Context, Result};
2use serde::{Deserialize, Serialize};
3use serde_json::{json, Value};
4use std::fs;
5use std::path::PathBuf;
6
7/// Python backend for managed environments.
8#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
9#[serde(rename_all = "lowercase")]
10pub enum PythonBackend {
11    /// uv — fast, deterministic, preferred (default if on PATH)
12    Uv,
13    /// Python — venv module fallback
14    Python,
15}
16
17impl Default for PythonBackend {
18    fn default() -> Self {
19        // Default to Uv if available, else Python
20        if crate::platform::environment::command_on_path("uv") {
21            PythonBackend::Uv
22        } else {
23            PythonBackend::Python
24        }
25    }
26}
27
28/// Resolve [`PythonBackend`] from the `MLSTACK_PYTHON_BACKEND` env override.
29///
30/// Centralizes the `"uv"` / `"python"` / default-fallback match so the two
31/// historical call sites (`apply_from_value`'s env override and
32/// `default_with_paths`'s construction-time read) cannot silently diverge if
33/// the accepted spellings ever change. Unknown / empty values fall back to the
34/// environment-aware default (Uv if installed, else Python).
35fn python_backend_from_env() -> PythonBackend {
36    match std::env::var("MLSTACK_PYTHON_BACKEND")
37        .ok()
38        .map(|v| v.trim().to_ascii_lowercase())
39        .as_deref()
40    {
41        Some("uv") => PythonBackend::Uv,
42        Some("python") => PythonBackend::Python,
43        _ => PythonBackend::default(),
44    }
45}
46
47/// Default vLLM source-build version lag — skip the newest release (supply-chain
48/// gate). 1 = second-newest (mild breach/bug buffer). Overridden by user config.
49fn default_vllm_version_lag() -> u32 {
50    1
51}
52
53/// Default ONNX Runtime install method: PyPI prebuilt `onnxruntime-migraphx`.
54fn default_onnx_install_method() -> String {
55    "migraphx".to_string()
56}
57
58#[derive(Debug, Clone, Serialize, Deserialize)]
59pub struct InstallerConfig {
60    pub scripts_dir: String,
61    pub log_dir: String,
62    pub install_path: String,
63    pub batch_mode: bool,
64    pub install_method: String,
65    pub auto_confirm: bool,
66    pub star_repos: bool,
67    pub force_reinstall: bool,
68    pub theme: String,
69    pub performance_profile: String,
70    /// vLLM source-build version gate (supply-chain lag): skip this many of the
71    /// newest PyPI releases (0 = latest, 1 = skip newest, …). Default 1.
72    #[serde(default = "default_vllm_version_lag")]
73    pub vllm_version_lag: u32,
74    /// vLLM source-build age gate: only adopt releases at least this many days
75    /// old (0 = off). Combined with `vllm_version_lag` as "newest satisfying both".
76    #[serde(default)]
77    pub vllm_version_min_age_days: u32,
78    /// ONNX Runtime install method (mirrors the vLLM version-lag config-row
79    /// pattern). "migraphx" = PyPI prebuilt onnxruntime-migraphx (default);
80    /// "prebuilt" = legacy onnxruntime-rocm; "source" = build from source.
81    /// The TUI sets `MLSTACK_ONNX_INSTALL_METHOD` from this before install.
82    #[serde(default = "default_onnx_install_method")]
83    pub onnx_install_method: String,
84    /// ONNX Runtime version override. `None` = pinned default (1.25.0). The TUI
85    /// sets `MLSTACK_ONNX_VERSION` from this (when `Some`) before install.
86    #[serde(default)]
87    pub onnx_version: Option<String>,
88    #[serde(default)]
89    pub python_backend: PythonBackend,
90    #[serde(skip)]
91    pub config_path: PathBuf,
92}
93
94impl InstallerConfig {
95    pub fn load_or_default(scripts_dir: &str) -> Result<Self> {
96        let config_path = config_file_path()?;
97        let log_dir = config_path
98            .parent()
99            .and_then(|p| p.parent())
100            .map(|p| p.join("logs"))
101            .unwrap_or_else(|| PathBuf::from("/tmp/mlstack/logs"));
102
103        let mut config = InstallerConfig::default_with_paths(
104            scripts_dir,
105            log_dir.to_string_lossy().to_string(),
106            config_path.clone(),
107        );
108
109        if config_path.exists() {
110            let raw = fs::read_to_string(&config_path).context("Failed to read config file")?;
111            let value: Value = serde_json::from_str(&raw).unwrap_or_else(|_| json!({}));
112            config.apply_from_value(&value);
113        } else {
114            // Only try to save if we can, but don't crash if we can't
115            let _ = config.save(None);
116        }
117
118        Ok(config)
119    }
120
121    pub fn save(&self, existing: Option<Value>) -> Result<()> {
122        if let Some(parent) = self.config_path.parent() {
123            fs::create_dir_all(parent).context("Failed to create config directory")?;
124        }
125        let mut value = existing.unwrap_or_else(|| json!({}));
126
127        value["scripts_dir"] = json!(self.scripts_dir);
128        value["log_dir"] = json!(self.log_dir);
129        value["batch_mode"] = json!(self.batch_mode);
130        value["install_method"] = json!(self.install_method);
131
132        let user_prefs = value
133            .get("user_preferences")
134            .cloned()
135            .unwrap_or_else(|| json!({}));
136        let mut user_prefs = user_prefs;
137        user_prefs["installation_path"] = json!(self.install_path);
138        user_prefs["install_method"] = json!(self.install_method);
139        user_prefs["auto_confirm"] = json!(self.auto_confirm);
140        user_prefs["star_repos"] = json!(self.star_repos);
141        user_prefs["force_reinstall"] = json!(self.force_reinstall);
142        user_prefs["theme"] = json!(self.theme);
143        user_prefs["performance_profile"] = json!(self.performance_profile);
144        user_prefs["python_backend"] = json!(self.python_backend);
145        user_prefs["vllm_version_lag"] = json!(self.vllm_version_lag);
146        user_prefs["vllm_version_min_age_days"] = json!(self.vllm_version_min_age_days);
147        user_prefs["onnx_install_method"] = json!(self.onnx_install_method);
148        user_prefs["onnx_version"] = json!(self.onnx_version);
149        value["user_preferences"] = user_prefs;
150
151        fs::write(&self.config_path, serde_json::to_string_pretty(&value)?)
152            .context("Failed to save config")?;
153
154        Ok(())
155    }
156
157    fn apply_from_value(&mut self, value: &Value) {
158        if let Some(dir) = value.get("scripts_dir").and_then(|v| v.as_str()) {
159            self.scripts_dir = dir.to_string();
160        }
161        if let Some(dir) = value.get("log_dir").and_then(|v| v.as_str()) {
162            self.log_dir = dir.to_string();
163        }
164        if let Some(batch) = value.get("batch_mode").and_then(|v| v.as_bool()) {
165            self.batch_mode = batch;
166        }
167        if let Some(method) = value.get("install_method").and_then(|v| v.as_str()) {
168            self.install_method = normalize_install_method(method);
169        }
170
171        if let Some(prefs) = value.get("user_preferences") {
172            if let Some(path) = prefs.get("installation_path").and_then(|v| v.as_str()) {
173                self.install_path = path.to_string();
174            }
175            if let Some(method) = prefs.get("install_method").and_then(|v| v.as_str()) {
176                self.install_method = normalize_install_method(method);
177            }
178            if let Some(auto) = prefs.get("auto_confirm").and_then(|v| v.as_bool()) {
179                self.auto_confirm = auto;
180            }
181            if let Some(star) = prefs.get("star_repos").and_then(|v| v.as_bool()) {
182                self.star_repos = star;
183            }
184            if let Some(force) = prefs.get("force_reinstall").and_then(|v| v.as_bool()) {
185                self.force_reinstall = force;
186            }
187            if let Some(theme) = prefs.get("theme").and_then(|v| v.as_str()) {
188                self.theme = theme.to_string();
189            }
190            if let Some(profile) = prefs.get("performance_profile").and_then(|v| v.as_str()) {
191                self.performance_profile = profile.to_string();
192            }
193            if let Some(backend) = prefs.get("python_backend").and_then(|v| v.as_str()) {
194                self.python_backend = match backend.to_lowercase().as_str() {
195                    "uv" => PythonBackend::Uv,
196                    "python" => PythonBackend::Python,
197                    _ => PythonBackend::default(),
198                };
199            }
200            if let Some(lag) = prefs.get("vllm_version_lag").and_then(|v| v.as_u64()) {
201                self.vllm_version_lag = lag as u32;
202            }
203            if let Some(age) = prefs
204                .get("vllm_version_min_age_days")
205                .and_then(|v| v.as_u64())
206            {
207                self.vllm_version_min_age_days = age as u32;
208            }
209            if let Some(method) = prefs.get("onnx_install_method").and_then(|v| v.as_str()) {
210                self.onnx_install_method = method.to_string();
211            }
212            if let Some(v) = prefs.get("onnx_version").and_then(|v| v.as_str()) {
213                let v = v.trim();
214                self.onnx_version = if v.is_empty() {
215                    None
216                } else {
217                    Some(v.to_string())
218                };
219            }
220        }
221
222        // Env override: MLSTACK_PYTHON_BACKEND (highest priority)
223        if std::env::var("MLSTACK_PYTHON_BACKEND").is_ok() {
224            self.python_backend = python_backend_from_env();
225        }
226    }
227
228    pub fn default_with_paths(scripts_dir: &str, log_dir: String, config_path: PathBuf) -> Self {
229        // Apply env override for python_backend at construction time
230        let python_backend = if std::env::var("MLSTACK_PYTHON_BACKEND").is_ok() {
231            python_backend_from_env()
232        } else {
233            PythonBackend::default()
234        };
235
236        Self {
237            scripts_dir: scripts_dir.to_string(),
238            log_dir,
239            install_path: "/opt/rocm".into(),
240            batch_mode: std::env::var("MLSTACK_BATCH_MODE").unwrap_or_default() == "1",
241            install_method: normalize_install_method(
242                &std::env::var("MLSTACK_INSTALL_METHOD").unwrap_or_else(|_| "auto".into()),
243            ),
244            auto_confirm: false,
245            star_repos: true,
246            force_reinstall: false,
247            theme: "dark".into(),
248            performance_profile: "balanced".into(),
249            vllm_version_lag: 1,
250            vllm_version_min_age_days: 0,
251            onnx_install_method: std::env::var("MLSTACK_ONNX_INSTALL_METHOD")
252                .unwrap_or_else(|_| "migraphx".into()),
253            onnx_version: std::env::var("MLSTACK_ONNX_VERSION")
254                .ok()
255                .map(|s| s.trim().to_string())
256                .filter(|s| !s.is_empty()),
257            python_backend,
258            config_path,
259        }
260    }
261}
262
263fn normalize_install_method(input: &str) -> String {
264    match input.trim().to_ascii_lowercase().as_str() {
265        "global" => "global".into(),
266        "venv" => "venv".into(),
267        _ => "auto".into(),
268    }
269}
270
271pub fn config_file_path() -> Result<PathBuf> {
272    let home = std::env::var("HOME").unwrap_or_else(|_| ".".into());
273    let config_dir = PathBuf::from(home).join(".mlstack").join("config");
274    fs::create_dir_all(&config_dir).context("Failed to create config directory")?;
275    Ok(config_dir.join("config.json"))
276}
277
278#[cfg(test)]
279mod tests {
280    use super::*;
281
282    #[test]
283    fn onnx_config_defaults_and_apply() {
284        let _env = crate::test_support::lock_env();
285        std::env::remove_var("MLSTACK_ONNX_INSTALL_METHOD");
286        std::env::remove_var("MLSTACK_ONNX_VERSION");
287
288        let mut cfg = InstallerConfig::default_with_paths(
289            "scripts",
290            "logs".into(),
291            PathBuf::from("/tmp/rusty-config-test.json"),
292        );
293        // Defaults: PyPI migraphx wheel, no version override (pinned 1.25.0).
294        assert_eq!(cfg.onnx_install_method, "migraphx");
295        assert_eq!(cfg.onnx_version, None);
296
297        // apply_from_value reads the user_preferences fields the TUI saves.
298        let raw = serde_json::json!({
299            "user_preferences": {
300                "onnx_install_method": "source",
301                "onnx_version": "1.27.1",
302            }
303        });
304        cfg.apply_from_value(&raw);
305        assert_eq!(cfg.onnx_install_method, "source");
306        assert_eq!(cfg.onnx_version.as_deref(), Some("1.27.1"));
307
308        // Empty version string → None (treated as "use default").
309        let raw = serde_json::json!({ "user_preferences": { "onnx_version": "  " } });
310        cfg.apply_from_value(&raw);
311        assert_eq!(cfg.onnx_version, None);
312    }
313}