kernel/discovery/
habitat.rs

1//! [`ModelHabitat`]: where models live on this machine. It computes the store
2//! roots (from the environment, the home directory, and user settings) and
3//! assembles the [`StoreScanner`]s that sweep them. This is the single place the
4//! four scanners are wired together and pointed at their default locations.
5//!
6//! The Apple-Foundation (builtin) scanner is platform-gated and not yet ported;
7//! `scanners` simply omits it.
8
9use std::collections::{HashMap, HashSet};
10use std::path::{Path, PathBuf};
11
12use crate::discovery::hf_scanner::HFCacheScanner;
13use crate::discovery::lm_studio_scanner::LMStudioScanner;
14use crate::discovery::loose_file_scanner::LooseFileScanner;
15use crate::discovery::ollama_scanner::OllamaStoreScanner;
16use crate::discovery::scanner::StoreScanner;
17use crate::fs::expand_tilde;
18use crate::records::SourceKind;
19
20/// The discovery-relevant model settings: user-added folders to watch and extra
21/// Hugging Face cache roots. (The full settings domain is not yet ported.)
22#[derive(Debug, Clone, Default, PartialEq, Eq)]
23pub struct ModelsSettings {
24    /// Extra directories the user asked to scan for loose files.
25    pub watched_folders: Vec<String>,
26    /// Extra Hugging Face cache roots the user configured.
27    pub hf_cache_roots: Vec<String>,
28}
29
30/// The machine's model locations, resolved from `home` and `environment`.
31#[derive(Debug, Clone)]
32pub struct ModelHabitat {
33    home: PathBuf,
34    environment: HashMap<String, String>,
35}
36
37impl ModelHabitat {
38    /// A habitat rooted at `home` with the given `environment`.
39    pub fn new(home: impl Into<PathBuf>, environment: HashMap<String, String>) -> Self {
40        Self {
41            home: home.into(),
42            environment,
43        }
44    }
45
46    /// A habitat detected from the process: `$HOME` and the current environment.
47    pub fn detect() -> Self {
48        let home = std::env::var_os("HOME")
49            .map(PathBuf::from)
50            .unwrap_or_default();
51        let environment = std::env::vars().collect();
52        Self { home, environment }
53    }
54
55    /// Every `(kind, root)` this habitat would scan, in order.
56    pub fn roots(&self, settings: &ModelsSettings) -> Vec<(SourceKind, PathBuf)> {
57        let mut roots = vec![(SourceKind::ollama(), self.ollama_root())];
58        for url in self.hf_default_roots(&settings.hf_cache_roots) {
59            roots.push((SourceKind::huggingface_cache(), url));
60        }
61        for url in lm_studio_roots(&self.home) {
62            roots.push((SourceKind::lm_studio(), url));
63        }
64        for url in loose_directories(&self.home) {
65            roots.push((SourceKind::file(), url));
66        }
67        for path in &settings.watched_folders {
68            roots.push((SourceKind::file(), PathBuf::from(path)));
69        }
70        roots
71    }
72
73    /// The scanners to run. If `kinds` is given, only scanners producing at least
74    /// one of those kinds are included; `None` includes all.
75    pub fn scanners(
76        &self,
77        kinds: Option<&[SourceKind]>,
78        settings: &ModelsSettings,
79    ) -> Vec<Box<dyn StoreScanner>> {
80        let wanted = |produced: &[SourceKind]| match kinds {
81            None => true,
82            Some(kinds) => produced.iter().any(|kind| kinds.contains(kind)),
83        };
84        let mut scanners: Vec<Box<dyn StoreScanner>> = Vec::new();
85        if wanted(&[SourceKind::ollama()]) {
86            scanners.push(Box::new(OllamaStoreScanner::new(self.ollama_root())));
87        }
88        if wanted(&[SourceKind::huggingface_cache()]) {
89            // The user roots go ONLY through `user_roots` (scanned as required);
90            // the default roots must exclude them (empty user list) so the scanner
91            // doesn't sweep a user root twice. `roots()` above intentionally does
92            // include them, since it enumerates the full set.
93            scanners.push(Box::new(HFCacheScanner::with_user_roots(
94                self.hf_default_roots(&[]),
95                self.hf_user_roots(&settings.hf_cache_roots),
96            )));
97        }
98        if wanted(&[SourceKind::lm_studio()]) {
99            scanners.push(Box::new(LMStudioScanner::new(lm_studio_roots(&self.home))));
100        }
101        if wanted(&[SourceKind::file(), SourceKind::folder()]) {
102            let watched = settings.watched_folders.iter().map(PathBuf::from).collect();
103            scanners.push(Box::new(LooseFileScanner::with_user_directories(
104                loose_directories(&self.home),
105                watched,
106            )));
107        }
108        scanners
109    }
110
111    fn ollama_root(&self) -> PathBuf {
112        match self.environment.get("OLLAMA_MODELS") {
113            Some(custom) if !custom.is_empty() => expand_tilde(custom, &self.home),
114            _ => self.home.join(".ollama/models"),
115        }
116    }
117
118    /// The standard Hugging Face roots (env + the default cache) plus the user's
119    /// configured roots, de-duplicated in order.
120    fn hf_default_roots(&self, user: &[String]) -> Vec<PathBuf> {
121        let mut candidates = Vec::new();
122        if let Some(cache) = self.environment.get("HF_HUB_CACHE")
123            && !cache.is_empty()
124        {
125            candidates.push(expand_tilde(cache, &self.home));
126        }
127        if let Some(hf_home) = self.environment.get("HF_HOME")
128            && !hf_home.is_empty()
129        {
130            candidates.push(expand_tilde(hf_home, &self.home).join("hub"));
131        }
132        candidates.push(self.home.join(".cache/huggingface/hub"));
133        candidates.extend(self.hf_user_roots(user));
134        dedup(candidates)
135    }
136
137    /// For each user path, the hub subdirectories that exist (`hub`,
138    /// `huggingface/hub`, or the path itself), falling back to the bare path.
139    fn hf_user_roots(&self, paths: &[String]) -> Vec<PathBuf> {
140        let mut roots = Vec::new();
141        for path in paths {
142            let base = expand_tilde(path, &self.home);
143            let candidates = [base.join("hub"), base.join("huggingface/hub"), base.clone()];
144            let existing: Vec<PathBuf> = candidates
145                .into_iter()
146                .filter(|url| is_hub_directory(url))
147                .collect();
148            if existing.is_empty() {
149                roots.push(base);
150            } else {
151                roots.extend(existing);
152            }
153        }
154        dedup(roots)
155    }
156}
157
158fn lm_studio_roots(home: &Path) -> Vec<PathBuf> {
159    vec![
160        home.join(".lmstudio/models"),
161        home.join(".cache/lm-studio/models"),
162    ]
163}
164
165fn loose_directories(home: &Path) -> Vec<PathBuf> {
166    vec![home.join("Downloads"), home.join("Models")]
167}
168
169fn is_hub_directory(url: &Path) -> bool {
170    url.is_dir()
171}
172
173/// De-duplicate paths, preserving first-seen order.
174fn dedup(paths: Vec<PathBuf>) -> Vec<PathBuf> {
175    let mut seen = HashSet::new();
176    paths
177        .into_iter()
178        .filter(|path| seen.insert(path.clone()))
179        .collect()
180}
kernel/discovery/habitat.rs

kernel/discovery/
habitat.rs
