reddb_file/
ai_model_cache.rs

1//! Local AI model cache file contracts.
2//!
3//! The server owns registry policy, HTTP handlers, fixture acquisition, and
4//! process locks. This module owns the persisted cache layout and manifest
5//! JSON shape.
6
7use std::fs;
8use std::io;
9use std::path::{Path, PathBuf};
10
11use serde_json::Value as JsonValue;
12
13pub const AI_MODEL_CACHE_DIR_NAME: &str = "ai_models_cache";
14pub const AI_MODEL_CACHE_STAGING_DIR_NAME: &str = ".staging";
15pub const AI_MODEL_CACHE_PURGE_DIR_NAME: &str = ".purge";
16pub const AI_MODEL_CACHE_MANIFEST_FILE: &str = "manifest.json";
17
18#[derive(Debug, Clone, PartialEq, Eq)]
19pub struct AiModelCacheManifestFile {
20    pub path: String,
21    pub sha256_hex: String,
22    pub size_bytes: u64,
23}
24
25#[derive(Debug, Clone, PartialEq, Eq)]
26pub struct AiModelCacheManifest {
27    pub name: String,
28    pub source: String,
29    pub revision: String,
30    pub task: String,
31    pub engine: String,
32    pub dimensions: u32,
33    pub installed_at_unix_ms: u64,
34    pub total_size_bytes: u64,
35    pub files: Vec<AiModelCacheManifestFile>,
36}
37
38pub fn ai_model_cache_root(base: &Path) -> PathBuf {
39    base.join(AI_MODEL_CACHE_DIR_NAME)
40}
41
42pub fn ai_model_cache_staging_root(cache_root: &Path) -> PathBuf {
43    cache_root.join(AI_MODEL_CACHE_STAGING_DIR_NAME)
44}
45
46pub fn ai_model_cache_purge_root(cache_root: &Path) -> PathBuf {
47    cache_root.join(AI_MODEL_CACHE_PURGE_DIR_NAME)
48}
49
50pub fn ai_model_cache_staging_dir(cache_root: &Path, name: &str, unique: &str) -> PathBuf {
51    ai_model_cache_staging_root(cache_root).join(format!("{name}-{unique}"))
52}
53
54pub fn ai_model_cache_purge_dir(cache_root: &Path, name: &str, unique: &str) -> PathBuf {
55    ai_model_cache_purge_root(cache_root).join(format!("{name}-{unique}"))
56}
57
58pub fn ai_model_cache_manifest_path(model_dir: &Path) -> PathBuf {
59    model_dir.join(AI_MODEL_CACHE_MANIFEST_FILE)
60}
61
62pub fn ai_model_cache_manifest_temp_path(dir: &Path) -> PathBuf {
63    dir.join(format!("{AI_MODEL_CACHE_MANIFEST_FILE}.tmp"))
64}
65
66pub fn copy_ai_model_cache_artifact(source: &Path, destination: &Path) -> io::Result<u64> {
67    if let Some(parent) = destination.parent() {
68        fs::create_dir_all(parent)?;
69    }
70    fs::copy(source, destination)
71}
72
73pub fn write_ai_model_cache_manifest(dir: &Path, bytes: &[u8]) -> io::Result<()> {
74    let manifest_tmp = ai_model_cache_manifest_temp_path(dir);
75    fs::write(&manifest_tmp, bytes)?;
76    fs::rename(&manifest_tmp, ai_model_cache_manifest_path(dir))
77}
78
79pub fn promote_ai_model_cache_staging(
80    cache_root: &Path,
81    model_name: &str,
82    unique: &str,
83    staging_dir: &Path,
84    model_dir: &Path,
85) -> io::Result<()> {
86    let purge_root = ai_model_cache_purge_root(cache_root);
87    fs::create_dir_all(&purge_root)?;
88    let purge_dir = ai_model_cache_purge_dir(cache_root, model_name, unique);
89    if model_dir.exists() {
90        fs::rename(model_dir, &purge_dir)?;
91    }
92    if let Err(err) = fs::rename(staging_dir, model_dir) {
93        if purge_dir.exists() {
94            let _ = fs::rename(&purge_dir, model_dir);
95        }
96        let _ = fs::remove_dir_all(staging_dir);
97        return Err(err);
98    }
99    if purge_dir.exists() {
100        let _ = fs::remove_dir_all(&purge_dir);
101    }
102    Ok(())
103}
104
105pub fn drop_ai_model_cache_dir(
106    cache_root: &Path,
107    model_dir: &Path,
108    unique: &str,
109) -> io::Result<bool> {
110    if !model_dir.exists() {
111        return Ok(false);
112    }
113    let purge_root = ai_model_cache_purge_root(cache_root);
114    fs::create_dir_all(&purge_root)?;
115    let model_name = model_dir
116        .file_name()
117        .and_then(|s| s.to_str())
118        .unwrap_or("model");
119    let purge_dir = ai_model_cache_purge_dir(cache_root, model_name, unique);
120    fs::rename(model_dir, &purge_dir)?;
121    let _ = fs::remove_dir_all(&purge_dir);
122    Ok(true)
123}
124
125pub fn encode_ai_model_cache_manifest_json(manifest: &AiModelCacheManifest) -> io::Result<Vec<u8>> {
126    serde_json::to_vec(&manifest_to_json(manifest)).map_err(|err| {
127        io::Error::new(
128            io::ErrorKind::InvalidData,
129            format!("encode AI model cache manifest: {err}"),
130        )
131    })
132}
133
134pub fn decode_ai_model_cache_manifest_json(bytes: &[u8]) -> io::Result<AiModelCacheManifest> {
135    let value: JsonValue = serde_json::from_slice(bytes).map_err(|err| {
136        io::Error::new(
137            io::ErrorKind::InvalidData,
138            format!("AI model cache manifest is not valid JSON: {err}"),
139        )
140    })?;
141    manifest_from_json(&value)
142}
143
144fn manifest_to_json(manifest: &AiModelCacheManifest) -> JsonValue {
145    let mut object = serde_json::Map::new();
146    object.insert("name".to_string(), JsonValue::String(manifest.name.clone()));
147    object.insert(
148        "source".to_string(),
149        JsonValue::String(manifest.source.clone()),
150    );
151    object.insert(
152        "revision".to_string(),
153        JsonValue::String(manifest.revision.clone()),
154    );
155    object.insert("task".to_string(), JsonValue::String(manifest.task.clone()));
156    object.insert(
157        "engine".to_string(),
158        JsonValue::String(manifest.engine.clone()),
159    );
160    object.insert(
161        "dimensions".to_string(),
162        JsonValue::Number(manifest.dimensions.into()),
163    );
164    object.insert(
165        "installed_at_unix_ms".to_string(),
166        JsonValue::Number(manifest.installed_at_unix_ms.into()),
167    );
168    object.insert(
169        "total_size_bytes".to_string(),
170        JsonValue::Number(manifest.total_size_bytes.into()),
171    );
172    let files = manifest
173        .files
174        .iter()
175        .map(|file| {
176            let mut object = serde_json::Map::new();
177            object.insert("path".to_string(), JsonValue::String(file.path.clone()));
178            object.insert(
179                "sha256".to_string(),
180                JsonValue::String(file.sha256_hex.clone()),
181            );
182            object.insert(
183                "size_bytes".to_string(),
184                JsonValue::Number(file.size_bytes.into()),
185            );
186            JsonValue::Object(object)
187        })
188        .collect();
189    object.insert("files".to_string(), JsonValue::Array(files));
190    JsonValue::Object(object)
191}
192
193fn manifest_from_json(value: &JsonValue) -> io::Result<AiModelCacheManifest> {
194    let object = value
195        .as_object()
196        .ok_or_else(|| invalid("manifest is not an object"))?;
197    let name = required_str(object, "name")?;
198    let source = required_str(object, "source")?;
199    let revision = required_str(object, "revision")?;
200    let task = required_str(object, "task")?;
201    let engine = required_str(object, "engine")?;
202    let dimensions = required_u64(object, "dimensions")? as u32;
203    let installed_at_unix_ms = required_u64(object, "installed_at_unix_ms")?;
204    let total_size_bytes = required_u64(object, "total_size_bytes")?;
205    let files_raw = object
206        .get("files")
207        .and_then(JsonValue::as_array)
208        .ok_or_else(|| invalid("manifest field 'files' must be an array"))?;
209    let mut files = Vec::with_capacity(files_raw.len());
210    for (idx, raw) in files_raw.iter().enumerate() {
211        let entry = raw
212            .as_object()
213            .ok_or_else(|| invalid(format!("manifest files[{idx}] is not an object")))?;
214        files.push(AiModelCacheManifestFile {
215            path: required_str_at(entry, "path", idx)?,
216            sha256_hex: required_str_at(entry, "sha256", idx)?,
217            size_bytes: required_u64_at(entry, "size_bytes", idx)?,
218        });
219    }
220    Ok(AiModelCacheManifest {
221        name,
222        source,
223        revision,
224        task,
225        engine,
226        dimensions,
227        installed_at_unix_ms,
228        total_size_bytes,
229        files,
230    })
231}
232
233fn required_str(object: &serde_json::Map<String, JsonValue>, key: &str) -> io::Result<String> {
234    object
235        .get(key)
236        .and_then(JsonValue::as_str)
237        .map(str::to_string)
238        .ok_or_else(|| invalid(format!("manifest field '{key}' missing or not a string")))
239}
240
241fn required_u64(object: &serde_json::Map<String, JsonValue>, key: &str) -> io::Result<u64> {
242    object
243        .get(key)
244        .and_then(JsonValue::as_u64)
245        .ok_or_else(|| invalid(format!("manifest field '{key}' missing or not a number")))
246}
247
248fn required_str_at(
249    object: &serde_json::Map<String, JsonValue>,
250    key: &str,
251    index: usize,
252) -> io::Result<String> {
253    object
254        .get(key)
255        .and_then(JsonValue::as_str)
256        .map(str::to_string)
257        .ok_or_else(|| invalid(format!("manifest files[{index}].{key} missing")))
258}
259
260fn required_u64_at(
261    object: &serde_json::Map<String, JsonValue>,
262    key: &str,
263    index: usize,
264) -> io::Result<u64> {
265    object
266        .get(key)
267        .and_then(JsonValue::as_u64)
268        .ok_or_else(|| invalid(format!("manifest files[{index}].{key} missing")))
269}
270
271fn invalid(message: impl Into<String>) -> io::Error {
272    io::Error::new(io::ErrorKind::InvalidData, message.into())
273}
274
275#[cfg(test)]
276mod tests {
277    use super::*;
278
279    #[test]
280    fn ai_model_cache_manifest_round_trips() {
281        let manifest = AiModelCacheManifest {
282            name: "mini".to_string(),
283            source: "fixture".to_string(),
284            revision: "abc".to_string(),
285            task: "embedding".to_string(),
286            engine: "candle".to_string(),
287            dimensions: 384,
288            installed_at_unix_ms: 42,
289            total_size_bytes: 11,
290            files: vec![AiModelCacheManifestFile {
291                path: "model.bin".to_string(),
292                sha256_hex: "00ff".to_string(),
293                size_bytes: 11,
294            }],
295        };
296
297        let bytes = encode_ai_model_cache_manifest_json(&manifest).expect("encode");
298        let decoded = decode_ai_model_cache_manifest_json(&bytes).expect("decode");
299        assert_eq!(decoded, manifest);
300        assert!(String::from_utf8(bytes)
301            .unwrap()
302            .contains("\"sha256\":\"00ff\""));
303    }
304
305    #[test]
306    fn ai_model_cache_paths_are_canonical() {
307        let root = Path::new("/tmp/reddb");
308        assert_eq!(
309            ai_model_cache_root(root),
310            Path::new("/tmp/reddb").join("ai_models_cache")
311        );
312        assert_eq!(
313            ai_model_cache_staging_dir(root, "m", "u"),
314            Path::new("/tmp/reddb").join(".staging").join("m-u")
315        );
316        assert_eq!(
317            ai_model_cache_purge_dir(root, "m", "u"),
318            Path::new("/tmp/reddb").join(".purge").join("m-u")
319        );
320        assert_eq!(
321            ai_model_cache_manifest_temp_path(root),
322            Path::new("/tmp/reddb").join("manifest.json.tmp")
323        );
324    }
325}
reddb_file/ai_model_cache.rs

reddb_file/
ai_model_cache.rs
