1use serde::{Deserialize, Serialize};
35use std::fs::File;
36use std::io::{self, BufReader, BufWriter, Write};
37use std::path::{Path, PathBuf};
38
39pub 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 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#[derive(Debug, Clone)]
117pub struct ModelStore {
118 root: PathBuf,
119}
120
121impl ModelStore {
122 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 pub fn at_platform_default() -> Self {
137 Self::open(default_models_home())
138 }
139
140 pub fn root(&self) -> &Path {
142 &self.root
143 }
144
145 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 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 pub fn manifest_path(&self, name: &str) -> PathBuf {
172 self.root.join("manifests").join(format!("{name}.json"))
173 }
174
175 pub fn lock_path(&self, name: &str) -> PathBuf {
179 self.root.join("locks").join(format!("{name}.lock"))
180 }
181
182 pub fn quarantine_dir(&self) -> PathBuf {
185 self.root.join("locks").join("quarantine")
186 }
187
188 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 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 std::fs::rename(&tmp_path, &final_path)?;
225 Ok(final_path)
226 }
227
228 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 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#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
257pub struct Manifest {
258 pub schema_version: u32,
260 pub name: String,
263 pub format: String,
265 pub blob: String,
267 pub size_bytes: u64,
269 #[serde(default)]
271 pub license: Option<String>,
272 pub source: ManifestSource,
274 pub produced_by: String,
276 pub produced_at: String,
278}
279
280#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
284pub struct ManifestSource {
285 pub registry: String,
287 pub repo: String,
289 pub revision: String,
291 pub filename: String,
293}
294
295pub fn parse_blob_ref(s: &str) -> Option<&str> {
297 s.strip_prefix("sha256:")
298}
299
300pub fn format_blob_ref(sha256_hex: &str) -> String {
302 format!("sha256:{sha256_hex}")
303}
304
305#[cfg(test)]
306#[allow(unsafe_code)] mod 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 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 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}