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

kernel/discovery/
habitat.rs
