Skip to main content

dspk_hub/
cache.rs

1//! Local cache of downloaded artifacts (`~/.cache/dspk/hub/` or `DSPK_HUB_DIR`):
2//! `models/<family>/<slug>/<file>` plus an `index.json` recording each installed file's hash and
3//! size, so status checks never re-read the files ([`HubCache::verify`] does).
4
5use anyhow::{Context, Result};
6use serde::{Deserialize, Serialize};
7use sha2::{Digest, Sha256};
8use std::collections::BTreeMap;
9use std::io::Read;
10use std::path::{Path, PathBuf};
11
12use crate::Artifact;
13
14/// Environment variable overriding the cache root.
15pub const HUB_DIR_ENV: &str = "DSPK_HUB_DIR";
16
17/// Read buffer of the hashing loop.
18const HASH_CHUNK_BYTES: usize = 64 * 1024;
19
20/// Local state of an artifact.
21#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
22pub enum ArtifactStatus {
23    /// Downloaded and recorded with the expected hash and size.
24    Installed,
25    /// Not downloaded.
26    Available,
27    /// A file is present but its record or size does not match the expected artifact.
28    Corrupted,
29}
30
31impl ArtifactStatus {
32    /// Badge for CLI tables.
33    pub fn badge(self) -> &'static str {
34        match self {
35            Self::Installed => "[INSTALLED]",
36            Self::Available => "[AVAILABLE]",
37            Self::Corrupted => "[CORRUPTED]",
38        }
39    }
40}
41
42impl std::fmt::Display for ArtifactStatus {
43    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
44        f.write_str(self.badge())
45    }
46}
47
48/// `index.json` entry of an installed artifact.
49#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
50pub struct CachedArtifactRecord {
51    pub id: String,
52    pub family: String,
53    pub sha256: String,
54    pub size_bytes: u64,
55    pub relative_path: String,
56    pub source_url: String,
57}
58
59/// `index.json` at the cache root.
60#[derive(Debug, Clone, Default, Serialize, Deserialize)]
61pub struct HubCacheIndex {
62    pub entries: BTreeMap<String, CachedArtifactRecord>,
63}
64
65/// The local cache directory. See the module docs.
66#[derive(Debug, Clone)]
67pub struct HubCache {
68    root_dir: PathBuf,
69}
70
71impl HubCache {
72    /// `$DSPK_HUB_DIR`, else `<user cache>/dspk/hub`.
73    pub fn from_env_or_default() -> Self {
74        if let Ok(dir) = std::env::var(HUB_DIR_ENV)
75            && !dir.trim().is_empty()
76        {
77            return Self::new(PathBuf::from(dir.trim()));
78        }
79        let base = dirs::cache_dir()
80            .or_else(|| dirs::home_dir().map(|h| h.join(".cache")))
81            .unwrap_or_else(|| PathBuf::from(".cache"));
82        Self::new(base.join("dspk").join("hub"))
83    }
84
85    pub fn new(root_dir: impl Into<PathBuf>) -> Self {
86        Self { root_dir: root_dir.into() }
87    }
88
89    pub fn root_dir(&self) -> &Path {
90        &self.root_dir
91    }
92
93    pub fn index_path(&self) -> PathBuf {
94        self.root_dir.join("index.json")
95    }
96
97    /// `<root>/models/<family>/<slug>/`.
98    pub fn artifact_dir(&self, artifact: &Artifact) -> PathBuf {
99        self.root_dir.join("models").join(&artifact.family).join(artifact.slug())
100    }
101
102    /// Local path of the artifact's file.
103    pub fn path(&self, artifact: &Artifact) -> PathBuf {
104        self.artifact_dir(artifact).join(&artifact.file_name)
105    }
106
107    /// `index.json`, or an empty index when there is none.
108    pub fn load_index(&self) -> Result<HubCacheIndex> {
109        let path = self.index_path();
110        if !path.exists() {
111            return Ok(HubCacheIndex::default());
112        }
113        let content = std::fs::read_to_string(&path).with_context(|| format!("failed to read hub cache index {}", path.display()))?;
114        serde_json::from_str(&content).with_context(|| format!("failed to parse hub cache index {}", path.display()))
115    }
116
117    pub fn save_index(&self, index: &HubCacheIndex) -> Result<()> {
118        std::fs::create_dir_all(&self.root_dir).with_context(|| format!("failed to create hub cache directory {}", self.root_dir.display()))?;
119        let path = self.index_path();
120        std::fs::write(&path, serde_json::to_string_pretty(index)?).with_context(|| format!("failed to write hub cache index {}", path.display()))
121    }
122
123    /// Records a verified download in `index.json`.
124    pub fn record_installed(&self, artifact: &Artifact, sha256: &str, size_bytes: u64) -> Result<()> {
125        let path = self.path(artifact);
126        let relative_path = path.strip_prefix(&self.root_dir).unwrap_or(&path).to_string_lossy().into_owned();
127        let mut index = self.load_index().unwrap_or_default();
128        index.entries.insert(
129            artifact.id.clone(),
130            CachedArtifactRecord {
131                id: artifact.id.clone(),
132                family: artifact.family.clone(),
133                sha256: sha256.to_ascii_lowercase(),
134                size_bytes,
135                relative_path,
136                source_url: artifact.url.clone(),
137            },
138        );
139        self.save_index(&index)
140    }
141
142    /// Status from `index.json` and the file's size (no hashing; see [`Self::verify`]).
143    pub fn status(&self, artifact: &Artifact) -> ArtifactStatus {
144        let path = self.path(artifact);
145        let Ok(meta) = path.metadata() else { return ArtifactStatus::Available };
146        let recorded = self.load_index().ok().and_then(|i| i.entries.get(&artifact.id).cloned());
147        match recorded {
148            Some(r) if r.sha256.eq_ignore_ascii_case(&artifact.sha256) && r.size_bytes == artifact.size_bytes && meta.len() == artifact.size_bytes => {
149                ArtifactStatus::Installed
150            }
151            _ => ArtifactStatus::Corrupted,
152        }
153    }
154
155    /// Re-hashes the cached file: `Ok(status)` with `Installed` only when hash and size match.
156    pub fn verify(&self, artifact: &Artifact) -> Result<(ArtifactStatus, Option<String>)> {
157        let path = self.path(artifact);
158        if !path.exists() {
159            return Ok((ArtifactStatus::Available, None));
160        }
161        let (digest, size) = compute_file_sha256(&path)?;
162        let ok = digest.eq_ignore_ascii_case(&artifact.sha256) && size == artifact.size_bytes;
163        Ok((if ok { ArtifactStatus::Installed } else { ArtifactStatus::Corrupted }, Some(digest)))
164    }
165
166    /// Removes an artifact's directory and index entry; `true` when anything was removed.
167    pub fn remove(&self, artifact: &Artifact) -> Result<bool> {
168        let mut removed = false;
169        let dir = self.artifact_dir(artifact);
170        if dir.exists() {
171            std::fs::remove_dir_all(&dir).with_context(|| format!("failed to remove cached artifact directory {}", dir.display()))?;
172            removed = true;
173        }
174        if let Ok(mut index) = self.load_index()
175            && index.entries.remove(&artifact.id).is_some()
176        {
177            self.save_index(&index)?;
178            removed = true;
179        }
180        Ok(removed)
181    }
182
183    /// Removes every cached artifact and the index; returns the bytes freed.
184    pub fn clean_all(&self) -> Result<u64> {
185        let models_dir = self.root_dir.join("models");
186        let freed = dir_size_bytes(&models_dir);
187        if models_dir.exists() {
188            std::fs::remove_dir_all(&models_dir).with_context(|| format!("failed to clean hub models directory {}", models_dir.display()))?;
189        }
190        let _ = std::fs::remove_file(self.index_path());
191        Ok(freed)
192    }
193}
194
195/// Lowercase hex SHA-256 and byte length of a file.
196pub fn compute_file_sha256(path: &Path) -> Result<(String, u64)> {
197    let mut file = std::fs::File::open(path).with_context(|| format!("failed to open file for SHA-256 check: {}", path.display()))?;
198    let mut hasher = Sha256::new();
199    let mut buf = vec![0u8; HASH_CHUNK_BYTES];
200    let mut total = 0u64;
201    loop {
202        let n = file.read(&mut buf)?;
203        if n == 0 {
204            break;
205        }
206        hasher.update(&buf[..n]);
207        total += n as u64;
208    }
209    Ok((hex::encode(hasher.finalize()), total))
210}
211
212fn dir_size_bytes(path: &Path) -> u64 {
213    let Ok(entries) = std::fs::read_dir(path) else { return 0 };
214    entries
215        .flatten()
216        .map(|e| {
217            let p = e.path();
218            if p.is_dir() { dir_size_bytes(&p) } else { e.metadata().map_or(0, |m| m.len()) }
219        })
220        .sum()
221}