Skip to main content

inferd_daemon/
store.rs

1//! Shared content-addressable model store (ADR 0011).
2//!
3//! Layout:
4//!
5//! ```text
6//! $MODELS_HOME/
7//! ├── blobs/
8//! │   └── sha256/
9//! │       └── <aa>/                          # 2-char fanout
10//! │           └── <full-hash>/
11//! │               └── data
12//! ├── manifests/
13//! │   └── <name>.json
14//! └── locks/
15//!     ├── <name>.lock                        # advisory write lock
16//!     └── quarantine/
17//!         └── <ts>-<reason>.bin              # bad blobs, kept for forensics
18//! ```
19//!
20//! Resolution order for `$MODELS_HOME` (first hit wins):
21//!
22//! 1. `models_home` field in the operator config file.
23//! 2. `MODELS_HOME` env var.
24//! 3. Platform default:
25//!    - Linux/*BSD: `${XDG_DATA_HOME:-$HOME/.local/share}/models/`
26//!    - macOS:      `~/Library/Application Support/models/`
27//!    - Windows:    `%LOCALAPPDATA%\models\`
28//!
29//! Windows MUST NOT default to `%APPDATA%` (Roaming). Roaming
30//! profiles upload `%APPDATA%` to the domain controller / OneDrive,
31//! which would replicate multi-GB blobs to every machine the user
32//! signs into.
33
34use serde::{Deserialize, Serialize};
35use std::fs::File;
36use std::io::{self, BufReader, BufWriter, Write};
37use std::path::{Path, PathBuf};
38
39/// Default `$MODELS_HOME` for the running platform.
40///
41/// Resolution chain documented in the module header.
42pub fn default_models_home() -> PathBuf {
43    if let Some(p) = std::env::var_os("MODELS_HOME") {
44        let pb = PathBuf::from(p);
45        if !pb.as_os_str().is_empty() {
46            return pb;
47        }
48    }
49    platform_default()
50}
51
52#[cfg(target_os = "linux")]
53fn platform_default() -> PathBuf {
54    if let Some(xdg) = std::env::var_os("XDG_DATA_HOME") {
55        let pb = PathBuf::from(xdg);
56        if !pb.as_os_str().is_empty() {
57            return pb.join("models");
58        }
59    }
60    home_dir()
61        .unwrap_or_else(|| PathBuf::from("."))
62        .join(".local")
63        .join("share")
64        .join("models")
65}
66
67#[cfg(target_os = "macos")]
68fn platform_default() -> PathBuf {
69    home_dir()
70        .unwrap_or_else(|| PathBuf::from("."))
71        .join("Library")
72        .join("Application Support")
73        .join("models")
74}
75
76#[cfg(windows)]
77fn platform_default() -> PathBuf {
78    if let Some(p) = std::env::var_os("LOCALAPPDATA") {
79        let pb = PathBuf::from(p);
80        if !pb.as_os_str().is_empty() {
81            return pb.join("models");
82        }
83    }
84    // Sensible fallback if LOCALAPPDATA is somehow unset.
85    home_dir()
86        .unwrap_or_else(|| PathBuf::from("."))
87        .join("AppData")
88        .join("Local")
89        .join("models")
90}
91
92#[cfg(not(any(target_os = "linux", target_os = "macos", windows)))]
93fn platform_default() -> PathBuf {
94    home_dir()
95        .unwrap_or_else(|| PathBuf::from("."))
96        .join(".local")
97        .join("share")
98        .join("models")
99}
100
101fn home_dir() -> Option<PathBuf> {
102    #[cfg(unix)]
103    {
104        std::env::var_os("HOME").map(PathBuf::from)
105    }
106    #[cfg(not(unix))]
107    {
108        std::env::var_os("USERPROFILE").map(PathBuf::from)
109    }
110}
111
112/// Handle to a model store rooted at one `$MODELS_HOME`.
113///
114/// Cheap to construct (just resolves paths). All disk creation is
115/// lazy: directories are made on first write.
116#[derive(Debug, Clone)]
117pub struct ModelStore {
118    root: PathBuf,
119}
120
121impl ModelStore {
122    /// Open a store rooted at `root`. Resolves `~` to the home dir.
123    pub fn open(root: impl Into<PathBuf>) -> Self {
124        let mut root = root.into();
125        if let Some(stripped) = root
126            .to_str()
127            .and_then(|s| s.strip_prefix("~/").or_else(|| s.strip_prefix("~\\")))
128            && let Some(home) = home_dir()
129        {
130            root = home.join(stripped);
131        }
132        Self { root }
133    }
134
135    /// Open the store at the platform default location.
136    pub fn at_platform_default() -> Self {
137        Self::open(default_models_home())
138    }
139
140    /// Root path of this store (`$MODELS_HOME`).
141    pub fn root(&self) -> &Path {
142        &self.root
143    }
144
145    /// Path to the `blobs/sha256/<aa>/<full-hash>/data` blob for
146    /// a SHA-256 hex string. Does NOT check existence.
147    pub fn blob_path(&self, sha256_hex: &str) -> PathBuf {
148        let aa = sha256_hex.get(..2).unwrap_or("00");
149        self.root
150            .join("blobs")
151            .join("sha256")
152            .join(aa)
153            .join(sha256_hex)
154            .join("data")
155    }
156
157    /// Path to the in-progress `data.tmp` for a download. Lives in
158    /// a `.partial-<hash>` sibling so abandoned downloads are
159    /// trivially distinguishable from finalised blobs.
160    pub fn partial_path(&self, sha256_hex: &str) -> PathBuf {
161        let aa = sha256_hex.get(..2).unwrap_or("00");
162        self.root
163            .join("blobs")
164            .join("sha256")
165            .join(aa)
166            .join(format!(".partial-{sha256_hex}"))
167            .join("data.tmp")
168    }
169
170    /// Path to a manifest by name.
171    pub fn manifest_path(&self, name: &str) -> PathBuf {
172        self.root.join("manifests").join(format!("{name}.json"))
173    }
174
175    /// Path to the advisory lock file for `name`. Producers hold an
176    /// exclusive lock on this file across blob-write + manifest-write
177    /// to keep two daemons from racing on the same name.
178    pub fn lock_path(&self, name: &str) -> PathBuf {
179        self.root.join("locks").join(format!("{name}.lock"))
180    }
181
182    /// Directory where bad blobs are quarantined. Per ADR 0011 we
183    /// keep these for forensic inspection rather than deleting them.
184    pub fn quarantine_dir(&self) -> PathBuf {
185        self.root.join("locks").join("quarantine")
186    }
187
188    /// Read a manifest by name. Returns `Ok(None)` if absent so
189    /// callers can branch on present/missing without parsing IO
190    /// kinds.
191    pub fn read_manifest(&self, name: &str) -> io::Result<Option<Manifest>> {
192        let path = self.manifest_path(name);
193        match File::open(&path) {
194            Ok(file) => {
195                let manifest: Manifest = serde_json::from_reader(BufReader::new(file))
196                    .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
197                Ok(Some(manifest))
198            }
199            Err(e) if e.kind() == io::ErrorKind::NotFound => Ok(None),
200            Err(e) => Err(e),
201        }
202    }
203
204    /// Write a manifest atomically: write to `<path>.tmp`, then
205    /// rename. The blob it references must already exist on disk —
206    /// callers are expected to write the blob first.
207    pub fn write_manifest(&self, manifest: &Manifest) -> io::Result<PathBuf> {
208        let dir = self.root.join("manifests");
209        std::fs::create_dir_all(&dir)?;
210        let final_path = self.manifest_path(&manifest.name);
211        let tmp_path = final_path.with_extension("json.tmp");
212        {
213            let file = File::create(&tmp_path)?;
214            let mut writer = BufWriter::new(file);
215            serde_json::to_writer_pretty(&mut writer, manifest)
216                .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
217            writer.write_all(b"\n")?;
218            writer.flush()?;
219        }
220        // POSIX `rename` is atomic when target exists. Windows
221        // `MoveFileExW(MOVEFILE_REPLACE_EXISTING)` (which Rust's
222        // `std::fs::rename` uses on Windows) is the closest
223        // equivalent and is the conventional choice here.
224        std::fs::rename(&tmp_path, &final_path)?;
225        Ok(final_path)
226    }
227
228    /// Move a bad-bytes file into the quarantine dir under a
229    /// timestamped name. Returns the new path.
230    pub fn quarantine(&self, src: &Path, reason: &str) -> io::Result<PathBuf> {
231        let qdir = self.quarantine_dir();
232        std::fs::create_dir_all(&qdir)?;
233        let ts = chrono::Utc::now().format("%Y%m%dT%H%M%SZ").to_string();
234        let safe_reason: String = reason
235            .chars()
236            .map(|c| if c.is_ascii_alphanumeric() { c } else { '-' })
237            .collect();
238        let dest = qdir.join(format!("{ts}-{safe_reason}.bin"));
239        std::fs::rename(src, &dest)?;
240        Ok(dest)
241    }
242
243    /// Ensure the on-disk skeleton (`blobs/`, `manifests/`,
244    /// `locks/`) exists. Idempotent.
245    pub fn ensure_layout(&self) -> io::Result<()> {
246        std::fs::create_dir_all(self.root.join("blobs").join("sha256"))?;
247        std::fs::create_dir_all(self.root.join("manifests"))?;
248        std::fs::create_dir_all(self.root.join("locks"))?;
249        Ok(())
250    }
251}
252
253/// Manifest schema v1. Wire-compatible with the cross-tool
254/// Shared Local Model Store proposal — other tools that adopt
255/// the convention can read inferd's manifests and vice versa.
256#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
257pub struct Manifest {
258    /// Always 1 for v1.
259    pub schema_version: u32,
260    /// Stable name (e.g. `"gemma-4-e4b"`). Maps 1:1 to the
261    /// manifest filename.
262    pub name: String,
263    /// Format token: `"gguf"` for llama.cpp-family weights.
264    pub format: String,
265    /// `sha256:<64-hex>` reference into `blobs/`.
266    pub blob: String,
267    /// Size of the blob in bytes. Diagnostic, not authoritative.
268    pub size_bytes: u64,
269    /// SPDX-style license id when known. Diagnostic only.
270    #[serde(default)]
271    pub license: Option<String>,
272    /// Where the blob came from. Diagnostic + future migration aid.
273    pub source: ManifestSource,
274    /// What wrote this manifest, in `<tool>/<version>` form.
275    pub produced_by: String,
276    /// RFC 3339 UTC timestamp.
277    pub produced_at: String,
278}
279
280/// Diagnostic provenance metadata for a manifest. Not consulted at
281/// runtime — the daemon trusts the SHA in `Manifest::blob`, not the
282/// upstream registry.
283#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
284pub struct ManifestSource {
285    /// Hostname the blob was downloaded from.
286    pub registry: String,
287    /// Repo path on that host (e.g. `unsloth/gemma-4-E4B-it-GGUF`).
288    pub repo: String,
289    /// Branch / tag / commit identifier.
290    pub revision: String,
291    /// Filename in the upstream repo, for migration breadcrumbs.
292    pub filename: String,
293}
294
295/// Parse a `sha256:<hex>` string into the bare hex.
296pub fn parse_blob_ref(s: &str) -> Option<&str> {
297    s.strip_prefix("sha256:")
298}
299
300/// Format a hex SHA-256 as a `sha256:<hex>` blob ref.
301pub fn format_blob_ref(sha256_hex: &str) -> String {
302    format!("sha256:{sha256_hex}")
303}
304
305#[cfg(test)]
306#[allow(unsafe_code)] // Edition 2024 made env::set_var/remove_var unsafe.
307mod tests {
308    use super::*;
309    use tempfile::tempdir;
310
311    #[test]
312    fn blob_path_uses_two_char_fanout() {
313        let store = ModelStore::open("/x");
314        let p = store.blob_path("abcdef0123456789abcdef0123456789abcdef0123456789abcdef0123456789");
315        let s = p.to_string_lossy();
316        assert!(s.contains("blobs"));
317        assert!(s.contains("sha256"));
318        assert!(s.ends_with("data") || s.ends_with("data\\") || s.ends_with("data/"));
319        // Fanout dir is "ab" (first two chars).
320        let parts: Vec<_> = p.components().collect();
321        assert!(
322            parts
323                .iter()
324                .any(|c| c.as_os_str() == std::ffi::OsStr::new("ab"))
325        );
326    }
327
328    #[test]
329    fn partial_path_lives_in_dot_partial_sibling() {
330        let store = ModelStore::open("/x");
331        let p =
332            store.partial_path("abcdef0123456789abcdef0123456789abcdef0123456789abcdef0123456789");
333        let s = p.to_string_lossy();
334        assert!(s.contains(".partial-abcdef"));
335        assert!(s.ends_with("data.tmp"));
336    }
337
338    #[test]
339    fn manifest_path_uses_name_dot_json() {
340        let store = ModelStore::open("/x");
341        let p = store.manifest_path("gemma-4-e4b");
342        assert!(
343            p.ends_with("manifests/gemma-4-e4b.json") || p.ends_with("manifests\\gemma-4-e4b.json")
344        );
345    }
346
347    #[test]
348    fn ensure_layout_creates_dirs() {
349        let dir = tempdir().unwrap();
350        let store = ModelStore::open(dir.path());
351        store.ensure_layout().unwrap();
352        assert!(dir.path().join("blobs").join("sha256").is_dir());
353        assert!(dir.path().join("manifests").is_dir());
354        assert!(dir.path().join("locks").is_dir());
355    }
356
357    #[test]
358    fn write_then_read_manifest_round_trip() {
359        let dir = tempdir().unwrap();
360        let store = ModelStore::open(dir.path());
361        let m = Manifest {
362            schema_version: 1,
363            name: "gemma-4-e4b".into(),
364            format: "gguf".into(),
365            blob: "sha256:30d1e7949597a3446726064e80b876fd1b5cba4aa6eec53d27afa420e731fb36".into(),
366            size_bytes: 5_126_304_928,
367            license: Some("apache-2.0".into()),
368            source: ManifestSource {
369                registry: "huggingface.co".into(),
370                repo: "unsloth/gemma-4-E4B-it-GGUF".into(),
371                revision: "main".into(),
372                filename: "gemma-4-E4B-it-UD-Q4_K_XL.gguf".into(),
373            },
374            produced_by: "inferd/0.1.0-alpha.0".into(),
375            produced_at: "2026-05-18T17:06:10Z".into(),
376        };
377        store.write_manifest(&m).unwrap();
378        let got = store.read_manifest("gemma-4-e4b").unwrap().unwrap();
379        assert_eq!(got, m);
380    }
381
382    #[test]
383    fn read_missing_manifest_returns_none() {
384        let dir = tempdir().unwrap();
385        let store = ModelStore::open(dir.path());
386        assert!(store.read_manifest("nope").unwrap().is_none());
387    }
388
389    #[test]
390    fn quarantine_moves_file_under_quarantine_dir() {
391        let dir = tempdir().unwrap();
392        let store = ModelStore::open(dir.path());
393        store.ensure_layout().unwrap();
394        let bad = dir.path().join("bad.bin");
395        std::fs::write(&bad, b"bytes").unwrap();
396        let qpath = store.quarantine(&bad, "sha-mismatch").unwrap();
397        assert!(!bad.exists());
398        assert!(qpath.exists());
399        assert!(
400            qpath
401                .to_string_lossy()
402                .contains(&format!("locks{}quarantine", std::path::MAIN_SEPARATOR))
403        );
404    }
405
406    #[test]
407    fn parse_blob_ref_strips_prefix() {
408        assert_eq!(parse_blob_ref("sha256:abc"), Some("abc"));
409        assert_eq!(parse_blob_ref("nope"), None);
410    }
411
412    #[test]
413    fn default_models_home_honours_models_home_env() {
414        // SAFETY: edition 2024 made env::set_var/remove_var unsafe
415        // because environment mutation isn't thread-safe. Cargo runs
416        // unit tests in a multi-threaded harness by default. This
417        // test is fine if it's the only test in the process touching
418        // MODELS_HOME — which it is — and we restore the saved value
419        // before returning so we don't leak state to siblings.
420        let saved = std::env::var_os("MODELS_HOME");
421        unsafe {
422            std::env::set_var("MODELS_HOME", "/tmp/inferd-test-models-home");
423        }
424        let p = default_models_home();
425        assert_eq!(p, PathBuf::from("/tmp/inferd-test-models-home"));
426        unsafe {
427            if let Some(v) = saved {
428                std::env::set_var("MODELS_HOME", v);
429            } else {
430                std::env::remove_var("MODELS_HOME");
431            }
432        }
433    }
434}