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 =
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 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 std::fs::rename(&tmp_path, &final_path)?;
230 Ok(final_path)
231 }
232
233 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 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#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
262pub struct Manifest {
263 pub schema_version: u32,
265 pub name: String,
268 pub format: String,
270 pub blob: String,
272 pub size_bytes: u64,
274 #[serde(default)]
276 pub license: Option<String>,
277 pub source: ManifestSource,
279 pub produced_by: String,
281 pub produced_at: String,
283}
284
285#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
289pub struct ManifestSource {
290 pub registry: String,
292 pub repo: String,
294 pub revision: String,
296 pub filename: String,
298}
299
300pub fn parse_blob_ref(s: &str) -> Option<&str> {
302 s.strip_prefix("sha256:")
303}
304
305pub fn format_blob_ref(sha256_hex: &str) -> String {
307 format!("sha256:{sha256_hex}")
308}
309
310#[cfg(test)]
311#[allow(unsafe_code)] mod 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 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 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}