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 =
196                    serde_json::from_reader(BufReader::new(file)).map_err(|e| {
197                        io::Error::new(
198                            io::ErrorKind::InvalidData,
199                            format!("malformed manifest at {}: {e}", path.display()),
200                        )
201                    })?;
202                Ok(Some(manifest))
203            }
204            Err(e) if e.kind() == io::ErrorKind::NotFound => Ok(None),
205            Err(e) => Err(e),
206        }
207    }
208
209    /// Write a manifest atomically: write to `<path>.tmp`, then
210    /// rename. The blob it references must already exist on disk —
211    /// callers are expected to write the blob first.
212    pub fn write_manifest(&self, manifest: &Manifest) -> io::Result<PathBuf> {
213        let dir = self.root.join("manifests");
214        std::fs::create_dir_all(&dir)?;
215        let final_path = self.manifest_path(&manifest.name);
216        let tmp_path = final_path.with_extension("json.tmp");
217        {
218            let file = File::create(&tmp_path)?;
219            let mut writer = BufWriter::new(file);
220            serde_json::to_writer_pretty(&mut writer, manifest)
221                .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
222            writer.write_all(b"\n")?;
223            writer.flush()?;
224        }
225        // POSIX `rename` is atomic when target exists. Windows
226        // `MoveFileExW(MOVEFILE_REPLACE_EXISTING)` (which Rust's
227        // `std::fs::rename` uses on Windows) is the closest
228        // equivalent and is the conventional choice here.
229        std::fs::rename(&tmp_path, &final_path)?;
230        Ok(final_path)
231    }
232
233    /// Move a bad-bytes file into the quarantine dir under a
234    /// timestamped name. Returns the new path.
235    pub fn quarantine(&self, src: &Path, reason: &str) -> io::Result<PathBuf> {
236        let qdir = self.quarantine_dir();
237        std::fs::create_dir_all(&qdir)?;
238        let ts = chrono::Utc::now().format("%Y%m%dT%H%M%SZ").to_string();
239        let safe_reason: String = reason
240            .chars()
241            .map(|c| if c.is_ascii_alphanumeric() { c } else { '-' })
242            .collect();
243        let dest = qdir.join(format!("{ts}-{safe_reason}.bin"));
244        std::fs::rename(src, &dest)?;
245        Ok(dest)
246    }
247
248    /// Ensure the on-disk skeleton (`blobs/`, `manifests/`,
249    /// `locks/`) exists. Idempotent.
250    pub fn ensure_layout(&self) -> io::Result<()> {
251        std::fs::create_dir_all(self.root.join("blobs").join("sha256"))?;
252        std::fs::create_dir_all(self.root.join("manifests"))?;
253        std::fs::create_dir_all(self.root.join("locks"))?;
254        Ok(())
255    }
256}
257
258/// Manifest schema v1. Wire-compatible with the cross-tool
259/// Shared Local Model Store proposal — other tools that adopt
260/// the convention can read inferd's manifests and vice versa.
261#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
262pub struct Manifest {
263    /// Always 1 for v1.
264    pub schema_version: u32,
265    /// Stable name (e.g. `"gemma-4-e4b"`). Maps 1:1 to the
266    /// manifest filename.
267    pub name: String,
268    /// Format token: `"gguf"` for llama.cpp-family weights.
269    pub format: String,
270    /// `sha256:<64-hex>` reference into `blobs/`.
271    pub blob: String,
272    /// Size of the blob in bytes. Diagnostic, not authoritative.
273    pub size_bytes: u64,
274    /// SPDX-style license id when known. Diagnostic only.
275    #[serde(default)]
276    pub license: Option<String>,
277    /// Where the blob came from. Diagnostic + future migration aid.
278    pub source: ManifestSource,
279    /// What wrote this manifest, in `<tool>/<version>` form.
280    pub produced_by: String,
281    /// RFC 3339 UTC timestamp.
282    pub produced_at: String,
283}
284
285/// Diagnostic provenance metadata for a manifest. Not consulted at
286/// runtime — the daemon trusts the SHA in `Manifest::blob`, not the
287/// upstream registry.
288#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
289pub struct ManifestSource {
290    /// Hostname the blob was downloaded from.
291    pub registry: String,
292    /// Repo path on that host (e.g. `unsloth/gemma-4-E4B-it-GGUF`).
293    pub repo: String,
294    /// Branch / tag / commit identifier.
295    pub revision: String,
296    /// Filename in the upstream repo, for migration breadcrumbs.
297    pub filename: String,
298}
299
300/// Parse a `sha256:<hex>` string into the bare hex.
301pub fn parse_blob_ref(s: &str) -> Option<&str> {
302    s.strip_prefix("sha256:")
303}
304
305/// Format a hex SHA-256 as a `sha256:<hex>` blob ref.
306pub fn format_blob_ref(sha256_hex: &str) -> String {
307    format!("sha256:{sha256_hex}")
308}
309
310#[cfg(test)]
311#[allow(unsafe_code)] // Edition 2024 made env::set_var/remove_var unsafe.
312mod tests {
313    use super::*;
314    use tempfile::tempdir;
315
316    #[test]
317    fn blob_path_uses_two_char_fanout() {
318        let store = ModelStore::open("/x");
319        let p = store.blob_path("abcdef0123456789abcdef0123456789abcdef0123456789abcdef0123456789");
320        let s = p.to_string_lossy();
321        assert!(s.contains("blobs"));
322        assert!(s.contains("sha256"));
323        assert!(s.ends_with("data") || s.ends_with("data\\") || s.ends_with("data/"));
324        // Fanout dir is "ab" (first two chars).
325        let parts: Vec<_> = p.components().collect();
326        assert!(
327            parts
328                .iter()
329                .any(|c| c.as_os_str() == std::ffi::OsStr::new("ab"))
330        );
331    }
332
333    #[test]
334    fn partial_path_lives_in_dot_partial_sibling() {
335        let store = ModelStore::open("/x");
336        let p =
337            store.partial_path("abcdef0123456789abcdef0123456789abcdef0123456789abcdef0123456789");
338        let s = p.to_string_lossy();
339        assert!(s.contains(".partial-abcdef"));
340        assert!(s.ends_with("data.tmp"));
341    }
342
343    #[test]
344    fn manifest_path_uses_name_dot_json() {
345        let store = ModelStore::open("/x");
346        let p = store.manifest_path("gemma-4-e4b");
347        assert!(
348            p.ends_with("manifests/gemma-4-e4b.json") || p.ends_with("manifests\\gemma-4-e4b.json")
349        );
350    }
351
352    #[test]
353    fn ensure_layout_creates_dirs() {
354        let dir = tempdir().unwrap();
355        let store = ModelStore::open(dir.path());
356        store.ensure_layout().unwrap();
357        assert!(dir.path().join("blobs").join("sha256").is_dir());
358        assert!(dir.path().join("manifests").is_dir());
359        assert!(dir.path().join("locks").is_dir());
360    }
361
362    #[test]
363    fn write_then_read_manifest_round_trip() {
364        let dir = tempdir().unwrap();
365        let store = ModelStore::open(dir.path());
366        let m = Manifest {
367            schema_version: 1,
368            name: "gemma-4-e4b".into(),
369            format: "gguf".into(),
370            blob: "sha256:30d1e7949597a3446726064e80b876fd1b5cba4aa6eec53d27afa420e731fb36".into(),
371            size_bytes: 5_126_304_928,
372            license: Some("apache-2.0".into()),
373            source: ManifestSource {
374                registry: "huggingface.co".into(),
375                repo: "unsloth/gemma-4-E4B-it-GGUF".into(),
376                revision: "main".into(),
377                filename: "gemma-4-E4B-it-UD-Q4_K_XL.gguf".into(),
378            },
379            produced_by: "inferd/0.1.0-alpha.0".into(),
380            produced_at: "2026-05-18T17:06:10Z".into(),
381        };
382        store.write_manifest(&m).unwrap();
383        let got = store.read_manifest("gemma-4-e4b").unwrap().unwrap();
384        assert_eq!(got, m);
385    }
386
387    #[test]
388    fn read_missing_manifest_returns_none() {
389        let dir = tempdir().unwrap();
390        let store = ModelStore::open(dir.path());
391        assert!(store.read_manifest("nope").unwrap().is_none());
392    }
393
394    #[test]
395    fn quarantine_moves_file_under_quarantine_dir() {
396        let dir = tempdir().unwrap();
397        let store = ModelStore::open(dir.path());
398        store.ensure_layout().unwrap();
399        let bad = dir.path().join("bad.bin");
400        std::fs::write(&bad, b"bytes").unwrap();
401        let qpath = store.quarantine(&bad, "sha-mismatch").unwrap();
402        assert!(!bad.exists());
403        assert!(qpath.exists());
404        assert!(
405            qpath
406                .to_string_lossy()
407                .contains(&format!("locks{}quarantine", std::path::MAIN_SEPARATOR))
408        );
409    }
410
411    #[test]
412    fn parse_blob_ref_strips_prefix() {
413        assert_eq!(parse_blob_ref("sha256:abc"), Some("abc"));
414        assert_eq!(parse_blob_ref("nope"), None);
415    }
416
417    #[test]
418    fn default_models_home_honours_models_home_env() {
419        // SAFETY: edition 2024 made env::set_var/remove_var unsafe
420        // because environment mutation isn't thread-safe. Cargo runs
421        // unit tests in a multi-threaded harness by default. This
422        // test is fine if it's the only test in the process touching
423        // MODELS_HOME — which it is — and we restore the saved value
424        // before returning so we don't leak state to siblings.
425        let saved = std::env::var_os("MODELS_HOME");
426        unsafe {
427            std::env::set_var("MODELS_HOME", "/tmp/inferd-test-models-home");
428        }
429        let p = default_models_home();
430        assert_eq!(p, PathBuf::from("/tmp/inferd-test-models-home"));
431        unsafe {
432            if let Some(v) = saved {
433                std::env::set_var("MODELS_HOME", v);
434            } else {
435                std::env::remove_var("MODELS_HOME");
436            }
437        }
438    }
439}