Skip to main content

studio_worker/
catalog.rs

1//! Local model catalog — the offline equivalent of the studio's
2//! `studioModels` registry.
3//!
4//! The studio is normally the single source of truth for a model's
5//! [`ModelSource`] (which files to download + the CLI defaults). When generating
6//! locally there is no studio, so the worker keeps a small JSON catalog the
7//! operator can edit and extend exactly the way they would add a model in the
8//! studio. It ships seeded with Z-Image-Turbo (the studio's default image
9//! model) so a fresh install can generate out of the box.
10
11use std::path::{Path, PathBuf};
12
13use serde::{Deserialize, Serialize};
14
15/// Tracing target for catalog persistence.  Stable so operators can
16/// filter with `RUST_LOG=studio_worker::catalog=debug`.
17const TRACE_TARGET: &str = "studio_worker::catalog";
18
19use crate::types::{
20    ModelCliDefaults, ModelEngine, ModelFile, ModelFileRole, ModelSource, TaskKind,
21};
22
23/// One catalog entry: a model id plus everything needed to run it. Mirrors the
24/// columns of the studio's `studioModels` row.
25#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
26#[serde(rename_all = "camelCase")]
27pub struct CatalogModel {
28    /// The model id the operator references (e.g. `z-image-turbo-q4_k_m.gguf`).
29    pub id: String,
30    /// Human-readable name shown in the UI.
31    pub display_name: String,
32    /// Task kind this model serves.
33    pub kind: TaskKind,
34    /// Device memory the model needs in GiB; admission checks it before a
35    /// load or transient job (see `docs/runtime/model-lifecycle.md`).
36    #[serde(default)]
37    pub vram_gb_estimate: f32,
38    /// Optional human description.
39    #[serde(default, skip_serializing_if = "Option::is_none")]
40    pub description: Option<String>,
41    /// Download spec + engine + CLI defaults (same shape the studio sends).
42    pub source: ModelSource,
43    /// Whether the model is selectable.
44    #[serde(default = "default_true")]
45    pub enabled: bool,
46    /// Where this entry came from: `"local"` (operator-added / seeded)
47    /// or `"studio"` (mirrored from a studio job offer).  A studio
48    /// re-offer refreshes studio-origin entries; a local-origin entry
49    /// of the same id is never clobbered by the sync.
50    #[serde(default = "default_origin")]
51    pub origin: String,
52    /// Models sharing a group are loaded one at a time: loading one
53    /// unloads the other first.
54    #[serde(default, skip_serializing_if = "Option::is_none")]
55    pub exclusive_group: Option<String>,
56}
57
58fn default_true() -> bool {
59    true
60}
61fn default_origin() -> String {
62    "local".into()
63}
64
65/// A collection of locally-available models.
66#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
67#[serde(rename_all = "camelCase")]
68pub struct Catalog {
69    #[serde(default)]
70    pub models: Vec<CatalogModel>,
71    /// Seed ids the operator deleted; startup seeding never re-adds them.
72    #[serde(default, skip_serializing_if = "Vec::is_empty")]
73    pub dismissed_seeds: Vec<String>,
74}
75
76impl Catalog {
77    /// The built-in catalog: every model the worker ships seeded with.
78    pub fn seed() -> Self {
79        Catalog {
80            models: vec![
81                zimage_turbo(),
82                qwen35_08b(),
83                nemotron_stream(),
84                parakeet_eou(),
85            ],
86            dismissed_seeds: Vec::new(),
87        }
88    }
89
90    /// Add every seed this catalogue lacks (an install that predates it),
91    /// never replacing the operator's entry and never re-adding a seed they
92    /// deleted.  `true` when anything was added.
93    pub fn ensure_seeds(&mut self) -> bool {
94        let missing: Vec<CatalogModel> = Self::seed()
95            .models
96            .into_iter()
97            .filter(|m| self.get(&m.id).is_none() && !self.dismissed_seeds.contains(&m.id))
98            .collect();
99        for model in &missing {
100            tracing::info!(target: TRACE_TARGET, op = "seed", model = %model.id, "added a seed model the catalogue lacked");
101        }
102        let added = !missing.is_empty();
103        self.models.extend(missing);
104        added
105    }
106
107    /// Parse a catalog from a JSON string.
108    pub fn from_json(json: &str) -> serde_json::Result<Self> {
109        serde_json::from_str(json)
110    }
111
112    /// Serialise to pretty JSON.
113    pub fn to_json(&self) -> serde_json::Result<String> {
114        serde_json::to_string_pretty(self)
115    }
116
117    /// Load the catalog from `path`.
118    ///
119    /// * Missing file → seeded with the built-in defaults and written.
120    /// * Corrupt JSON → the file is **quarantined** (renamed to
121    ///   `models.json.corrupt-<unix-ts>`) and a fresh seed written in
122    ///   its place.  The old behaviour — erroring so the caller fell
123    ///   back to an in-memory seed while keeping the save path — meant
124    ///   the next persist silently overwrote the operator's hand-edited
125    ///   catalog; quarantining preserves their bytes for recovery.
126    /// * Any other IO error propagates (nothing is renamed or written).
127    pub fn load_or_seed(path: &Path) -> std::io::Result<Self> {
128        match std::fs::read_to_string(path) {
129            Ok(contents) => match Self::from_json(&contents) {
130                Ok(mut catalog) => {
131                    if catalog.ensure_seeds() {
132                        catalog.save(path)?;
133                    }
134                    Ok(catalog)
135                }
136                Err(parse_err) => {
137                    let quarantine = quarantine_path(path);
138                    std::fs::rename(path, &quarantine)?;
139                    tracing::warn!(
140                        target: TRACE_TARGET,
141                        op = "load",
142                        path = %path.display(),
143                        quarantine = %quarantine.display(),
144                        error = %parse_err,
145                        "catalog is not valid JSON; quarantined the file and reseeded"
146                    );
147                    let seeded = Self::seed();
148                    seeded.save(path)?;
149                    Ok(seeded)
150                }
151            },
152            Err(err) if err.kind() == std::io::ErrorKind::NotFound => {
153                let seeded = Self::seed();
154                seeded.save(path)?;
155                Ok(seeded)
156            }
157            Err(err) => Err(err),
158        }
159    }
160
161    /// Load for serving: the catalog plus the path future saves may
162    /// write to.  A quarantine/seed recovery keeps the path (the file
163    /// is now healthy); an unreadable file (permissions, IO) drops it
164    /// so the worker can never overwrite a file it couldn't read.
165    pub fn load_for_serving(path: Option<PathBuf>) -> (Self, Option<PathBuf>) {
166        match path {
167            Some(path) => match Self::load_or_seed(&path) {
168                Ok(catalog) => (catalog, Some(path)),
169                Err(err) => {
170                    tracing::warn!(
171                        target: TRACE_TARGET,
172                        op = "load",
173                        path = %path.display(),
174                        error = %err,
175                        "catalog unreadable; serving the in-memory seed and \
176                         disabling persistence so the file is never clobbered"
177                    );
178                    (Self::seed(), None)
179                }
180            },
181            None => (Self::seed(), None),
182        }
183    }
184
185    /// Write the catalog to `path` (creating parent dirs) — atomically,
186    /// via the same temp-file + rename dance as `config.toml`, so a
187    /// crash mid-write can't truncate the operator's model catalog.
188    pub fn save(&self, path: &Path) -> std::io::Result<()> {
189        if let Some(parent) = path.parent() {
190            std::fs::create_dir_all(parent)?;
191        }
192        let json = self
193            .to_json()
194            .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
195        crate::config::write_atomic(path, json.as_bytes()).map_err(std::io::Error::other)
196    }
197
198    /// Look up a model by id.
199    pub fn get(&self, id: &str) -> Option<&CatalogModel> {
200        self.models.iter().find(|m| m.id == id)
201    }
202
203    /// All catalog entries.
204    pub fn list(&self) -> &[CatalogModel] {
205        &self.models
206    }
207
208    /// Insert a model, replacing any existing entry with the same id.
209    pub fn upsert(&mut self, model: CatalogModel) {
210        self.dismissed_seeds.retain(|d| *d != model.id);
211        if let Some(existing) = self.models.iter_mut().find(|m| m.id == model.id) {
212            *existing = model;
213        } else {
214            self.models.push(model);
215        }
216    }
217
218    /// Mirror a model seen on a studio job offer into the catalog so
219    /// the local API can serve it too.  Returns whether the catalog
220    /// changed (a no-op returns `false`, so the caller can skip the
221    /// disk write).  A **local-origin** entry of the same id is never
222    /// clobbered — the operator's own edits win; an unchanged
223    /// studio-origin entry is left alone so a re-offer every job
224    /// doesn't churn the file.
225    pub fn sync_studio_model(&mut self, incoming: CatalogModel) -> bool {
226        if let Some(existing) = self.models.iter_mut().find(|m| m.id == incoming.id) {
227            if existing.origin == "local" {
228                return false; // never overwrite operator-owned entries
229            }
230            if *existing == incoming {
231                return false; // already up to date
232            }
233            *existing = incoming;
234            return true;
235        }
236        self.models.push(incoming);
237        true
238    }
239
240    /// Remove a model by id. Returns whether it existed.
241    pub fn remove(&mut self, id: &str) -> bool {
242        let before = self.models.len();
243        self.models.retain(|m| m.id != id);
244        let removed = self.models.len() != before;
245        let is_seed = Self::seed().models.iter().any(|m| m.id == id);
246        if removed && is_seed && !self.dismissed_seeds.iter().any(|d| d == id) {
247            self.dismissed_seeds.push(id.to_string());
248        }
249        removed
250    }
251
252    /// The first enabled image model — used when a request names no model.
253    pub fn default_image_model(&self) -> Option<&CatalogModel> {
254        self.default_model_for(TaskKind::Image)
255    }
256
257    /// The first enabled model of `kind` — used when a request for that
258    /// kind names no explicit model.  Generalises
259    /// [`default_image_model`](Self::default_image_model) so the local
260    /// API can serve every modality the worker's engines support.
261    pub fn default_model_for(&self, kind: TaskKind) -> Option<&CatalogModel> {
262        self.models.iter().find(|m| m.enabled && m.kind == kind)
263    }
264}
265
266/// Where a corrupt catalog gets parked: `<name>.corrupt-<unix-ts>`,
267/// beside the original so the operator can recover their edits.
268pub(crate) fn quarantine_path(path: &Path) -> PathBuf {
269    let name = path
270        .file_name()
271        .map(|n| n.to_string_lossy().into_owned())
272        .unwrap_or_else(|| "models.json".to_string());
273    let ts = std::time::SystemTime::now()
274        .duration_since(std::time::UNIX_EPOCH)
275        .map(|d| d.as_secs())
276        .unwrap_or(0);
277    path.with_file_name(format!("{name}.corrupt-{ts}"))
278}
279
280/// A model file on Hugging Face, pinned to a repository revision.
281fn hf_file(repo: &str, revision: &str, path: &str, bytes: u64, sha256: &str) -> ModelFile {
282    ModelFile {
283        role: ModelFileRole::Model,
284        url: format!("https://huggingface.co/{repo}/resolve/{revision}/{path}"),
285        filename: path.rsplit('/').next().unwrap_or(path).to_string(),
286        approx_bytes: Some(bytes),
287        sha256: Some(sha256.into()),
288    }
289}
290
291/// Revision of `altunenes/parakeet-rs` the streaming speech seeds pin.
292const PARAKEET_RS_REVISION: &str = "4d2a8bc71f5c896ec40faa59732e6716295edaf2";
293
294/// Qwen3.5 0.8B instruct (unsloth Q8_0): a small chat model for one-shot
295/// jobs such as titles and summaries.  Reasoning off (the template thinks
296/// only when asked; a small model left to think runs to its token cap);
297/// 32K context for long transcripts.  Apache-2.0.
298fn qwen35_08b() -> CatalogModel {
299    let mut kwargs = serde_json::Map::new();
300    kwargs.insert("enable_thinking".into(), serde_json::Value::Bool(false));
301    CatalogModel {
302        id: "qwen3.5-0.8b".into(),
303        display_name: "Qwen3.5 0.8B instruct (Q8_0)".into(),
304        kind: TaskKind::Llm,
305        vram_gb_estimate: 1.5,
306        description: Some("Small chat model: reasoning off, 32K context".into()),
307        source: ModelSource {
308            engine: ModelEngine::LlamaCpp,
309            files: vec![hf_file(
310                "unsloth/Qwen3.5-0.8B-GGUF",
311                "6ab461498e2023f6e3c1baea90a8f0fe38ab64d0",
312                "Qwen3.5-0.8B-Q8_0.gguf",
313                811_843_840,
314                "0ad885ffd4bb022fc4f0d33a3308fa108ef8613159d3b3a67e23abca056b7a6c",
315            )],
316            cli_defaults: ModelCliDefaults {
317                context_size: Some(32_768),
318                chat_template_kwargs: Some(kwargs),
319                ..Default::default()
320            },
321        },
322        enabled: true,
323        origin: "local".into(),
324        exclusive_group: None,
325    }
326}
327
328/// Nemotron 3.5 streaming ASR (0.6B, multilingual, punctuated): 560 ms
329/// chunks.  Measured ~3.5 GiB on CUDA.  NVIDIA Open Model License.
330fn nemotron_stream() -> CatalogModel {
331    let dir = "nemotron-3.5-asr-streaming-0.6b-onnx";
332    let f = |name: &str, bytes: u64, sha: &str| {
333        hf_file(
334            "altunenes/parakeet-rs",
335            PARAKEET_RS_REVISION,
336            &format!("{dir}/{name}"),
337            bytes,
338            sha,
339        )
340    };
341    CatalogModel {
342        id: "nemotron-3.5-stream".into(),
343        display_name: "Nemotron 3.5 streaming (0.6B, multilingual)".into(),
344        kind: TaskKind::AudioStt,
345        vram_gb_estimate: 3.5,
346        description: Some("Streaming speech-to-text with punctuation".into()),
347        source: ModelSource {
348            engine: ModelEngine::Parakeet,
349            files: vec![
350                f(
351                    "config.json",
352                    2_979,
353                    "b0289e196d11a17e3c661bbadfe455c87de4baffc1a5e652a5779f5d687c5db0",
354                ),
355                f(
356                    "decoder_joint.onnx",
357                    97_590_054,
358                    "634dfadf24cb4f73c2fae170b36611d68db48186426882cbc8f7e02ed9f2bb29",
359                ),
360                f(
361                    "encoder.onnx",
362                    42_164_972,
363                    "d569fbe78b48fbb04e169d324f5d25463838ceed7b5fc3bfe209872441979bd9",
364                ),
365                f(
366                    "encoder.onnx.data",
367                    2_454_405_120,
368                    "7584f85df76bc9ae6fbdfa53aa8d97b07a842525d1c501d536d77fd9e4f57ac7",
369                ),
370                f(
371                    "tokenizer.model",
372                    406_554,
373                    "ce3895e40806f02a26c3a225161b96ef682d6c0054bae32a245dec4258d7d291",
374                ),
375            ],
376            cli_defaults: ModelCliDefaults::default(),
377        },
378        enabled: true,
379        origin: "local".into(),
380        exclusive_group: Some("stt".into()),
381    }
382}
383
384/// Parakeet realtime EOU (120M, English): 160 ms chunks with
385/// end-of-utterance detection.  Measured ~1.1 GiB on CUDA.  CC-BY-4.0.
386fn parakeet_eou() -> CatalogModel {
387    let dir = "realtime_eou_120m-v1-onnx";
388    let f = |name: &str, bytes: u64, sha: &str| {
389        hf_file(
390            "altunenes/parakeet-rs",
391            PARAKEET_RS_REVISION,
392            &format!("{dir}/{name}"),
393            bytes,
394            sha,
395        )
396    };
397    CatalogModel {
398        id: "parakeet-eou-120m".into(),
399        display_name: "Parakeet EOU (120M, English)".into(),
400        kind: TaskKind::AudioStt,
401        vram_gb_estimate: 1.2,
402        description: Some("Light streaming speech-to-text, English".into()),
403        source: ModelSource {
404            engine: ModelEngine::Parakeet,
405            files: vec![
406                f(
407                    "decoder_joint.onnx",
408                    21_347_639,
409                    "9d2553ac043c2fc5f69e970769b0fb8ab9103fbfdeb7d26a1ea9729d4bd2dddd",
410                ),
411                f(
412                    "encoder.onnx",
413                    459_341_289,
414                    "d472887cc38a784a5bfc21c2dbe247639edc3b3f9992388d8ceceaec07256b5b",
415                ),
416                f(
417                    "tokenizer.json",
418                    20_053,
419                    "f6b0ad8690559351fa478116fe0985a203b76f7c040f3a9381f485c99c0325f8",
420                ),
421            ],
422            cli_defaults: ModelCliDefaults::default(),
423        },
424        enabled: true,
425        origin: "local".into(),
426        exclusive_group: Some("stt".into()),
427    }
428}
429
430/// The canonical Z-Image-Turbo entry, mirroring the studio seed
431/// (`migrations/graphics/0017_seed_registry.sql`).
432fn zimage_turbo() -> CatalogModel {
433    CatalogModel {
434        id: "z-image-turbo-q4_k_m.gguf".into(),
435        display_name: "Z-Image Turbo (Q4_K_M)".into(),
436        kind: TaskKind::Image,
437        vram_gb_estimate: 12.0,
438        description: Some(
439            "Distilled 8-step diffusion model packaged for sd.cpp. Diffusion (Q4_K), \
440             Qwen3-4B text encoder, Flux ae.safetensors VAE."
441                .into(),
442        ),
443        source: ModelSource {
444            engine: ModelEngine::SdCpp,
445            files: vec![
446                // sha256 pins sourced from the HF LFS oids
447                // (`/api/models/<repo>/tree/main`) and cross-checked
448                // against freshly downloaded copies — the out-of-the-box
449                // model must never be swappable in transit or at rest.
450                ModelFile {
451                    role: ModelFileRole::DiffusionModel,
452                    url: "https://huggingface.co/leejet/Z-Image-Turbo-GGUF/resolve/main/z_image_turbo-Q4_K.gguf".into(),
453                    filename: "z_image_turbo-Q4_K.gguf".into(),
454                    approx_bytes: Some(3_864_250_304),
455                    sha256: Some(
456                        "14b375ab4f226bc5378f68f37e899ef3c2242b8541e61e2bc1aff40976086fbd".into(),
457                    ),
458                },
459                ModelFile {
460                    role: ModelFileRole::TextEncoder,
461                    url: "https://huggingface.co/unsloth/Qwen3-4B-Instruct-2507-GGUF/resolve/main/Qwen3-4B-Instruct-2507-Q4_K_M.gguf".into(),
462                    filename: "Qwen3-4B-Instruct-2507-Q4_K_M.gguf".into(),
463                    approx_bytes: Some(2_497_281_120),
464                    sha256: Some(
465                        "3605803b982cb64aead44f6c1b2ae36e3acdb41d8e46c8a94c6533bc4c67e597".into(),
466                    ),
467                },
468                ModelFile {
469                    role: ModelFileRole::Vae,
470                    url: "https://huggingface.co/Comfy-Org/Lumina_Image_2.0_Repackaged/resolve/main/split_files/vae/ae.safetensors".into(),
471                    filename: "ae.safetensors".into(),
472                    approx_bytes: Some(335_304_388),
473                    sha256: Some(
474                        "afc8e28272cd15db3919bacdb6918ce9c1ed22e96cb12c4d5ed0fba823529e38".into(),
475                    ),
476                },
477            ],
478            cli_defaults: ModelCliDefaults {
479                cfg_scale: 1.0,
480                steps: 8,
481                width: 1024,
482                height: 1024,
483                sampling_method: Some("euler".into()),
484                flow_shift: None,
485                zero_cond_t: None,
486                offload_to_cpu: None,
487                context_size: None,
488                chat_template_kwargs: None,
489            },
490        },
491        enabled: true,
492        origin: "local".into(),
493        exclusive_group: None,
494    }
495}
496
497#[cfg(test)]
498mod tests {
499    use super::*;
500
501    fn ids(c: &Catalog) -> Vec<&str> {
502        c.models.iter().map(|m| m.id.as_str()).collect()
503    }
504
505    #[test]
506    fn the_seed_carries_image_llm_and_both_streaming_speech_models() {
507        assert_eq!(
508            ids(&Catalog::seed()),
509            [
510                "z-image-turbo-q4_k_m.gguf",
511                "qwen3.5-0.8b",
512                "nemotron-3.5-stream",
513                "parakeet-eou-120m"
514            ]
515        );
516    }
517
518    #[test]
519    fn every_seed_file_is_pinned_and_checksummed() {
520        for model in Catalog::seed().models {
521            for file in &model.source.files {
522                assert!(
523                    file.sha256.as_deref().is_some_and(|h| h.len() == 64),
524                    "{} {}",
525                    model.id,
526                    file.filename
527                );
528                assert!(
529                    file.approx_bytes.is_some_and(|b| b > 0),
530                    "{} {}",
531                    model.id,
532                    file.filename
533                );
534                if model.id != "z-image-turbo-q4_k_m.gguf" {
535                    // Z-Image mirrors the studio's seed; the rest pin a revision.
536                    assert!(
537                        !file.url.contains("/resolve/main/"),
538                        "{} pins a revision",
539                        file.url
540                    );
541                }
542            }
543        }
544    }
545
546    #[test]
547    fn the_streaming_seeds_swap_with_each_other() {
548        let seed = Catalog::seed();
549        for id in ["nemotron-3.5-stream", "parakeet-eou-120m"] {
550            let m = seed.get(id).unwrap();
551            assert_eq!(m.kind, TaskKind::AudioStt);
552            assert_eq!(m.source.engine, ModelEngine::Parakeet);
553            assert_eq!(m.exclusive_group.as_deref(), Some("stt"));
554        }
555    }
556
557    #[test]
558    fn the_llm_seed_answers_without_reasoning_in_a_long_context() {
559        let seed = Catalog::seed();
560        let m = seed.get("qwen3.5-0.8b").unwrap();
561        assert_eq!(m.kind, TaskKind::Llm);
562        assert_eq!(m.source.engine, ModelEngine::LlamaCpp);
563        assert_eq!(m.source.cli_defaults.context_size, Some(32768));
564        assert_eq!(
565            m.source.cli_defaults.chat_template_kwargs.as_ref().unwrap()["enable_thinking"],
566            false
567        );
568    }
569
570    #[test]
571    fn missing_seeds_are_added_to_an_existing_catalogue() {
572        let mut c = Catalog {
573            models: vec![zimage_turbo()],
574            ..Default::default()
575        };
576        assert!(c.ensure_seeds());
577        assert_eq!(ids(&c), ids(&Catalog::seed()));
578        assert!(!c.ensure_seeds(), "idempotent");
579    }
580
581    #[test]
582    fn seeding_never_overwrites_the_operators_entry() {
583        let mut mine = Catalog::seed().get("qwen3.5-0.8b").unwrap().clone();
584        mine.display_name = "my tuned qwen".into();
585        let mut c = Catalog {
586            models: vec![mine],
587            ..Default::default()
588        };
589        c.ensure_seeds();
590        assert_eq!(c.get("qwen3.5-0.8b").unwrap().display_name, "my tuned qwen");
591    }
592
593    #[test]
594    fn a_deleted_seed_stays_deleted() {
595        let mut c = Catalog::seed();
596        assert!(c.remove("parakeet-eou-120m"));
597        assert_eq!(c.dismissed_seeds, ["parakeet-eou-120m"]);
598        assert!(!c.ensure_seeds());
599        assert!(c.get("parakeet-eou-120m").is_none());
600        let reloaded = Catalog::from_json(&c.to_json().unwrap()).unwrap();
601        assert_eq!(reloaded.dismissed_seeds, ["parakeet-eou-120m"]);
602    }
603
604    #[test]
605    fn re_adding_a_dismissed_seed_forgets_the_dismissal() {
606        let mut c = Catalog::seed();
607        let eou = c.get("parakeet-eou-120m").unwrap().clone();
608        c.remove("parakeet-eou-120m");
609        c.upsert(eou);
610        assert!(c.dismissed_seeds.is_empty());
611    }
612
613    #[test]
614    fn removing_a_non_seed_records_nothing() {
615        let mut c = Catalog::seed();
616        c.upsert(studio_model("mine"));
617        c.remove("mine");
618        assert!(c.dismissed_seeds.is_empty());
619    }
620
621    #[test]
622    fn loading_an_older_catalogue_adds_the_new_seeds_and_saves() {
623        let dir = tempfile::tempdir().unwrap();
624        let path = dir.path().join("models.json");
625        let old = Catalog {
626            models: vec![zimage_turbo()],
627            ..Default::default()
628        };
629        std::fs::write(&path, old.to_json().unwrap()).unwrap();
630        let loaded = Catalog::load_or_seed(&path).unwrap();
631        assert_eq!(ids(&loaded), ids(&Catalog::seed()));
632        let on_disk = Catalog::from_json(&std::fs::read_to_string(&path).unwrap()).unwrap();
633        assert_eq!(ids(&on_disk), ids(&Catalog::seed()));
634    }
635
636    #[test]
637    fn seed_contains_zimage_with_three_files() {
638        let catalog = Catalog::seed();
639        let model = catalog
640            .get("z-image-turbo-q4_k_m.gguf")
641            .expect("z-image seeded");
642        assert_eq!(model.kind, TaskKind::Image);
643        assert_eq!(model.source.engine, ModelEngine::SdCpp);
644        assert_eq!(model.source.files.len(), 3);
645        assert_eq!(model.source.cli_defaults.steps, 8);
646        assert!(model.enabled);
647    }
648
649    #[test]
650    fn every_seeded_file_is_https_and_integrity_pinned() {
651        // The out-of-the-box downloads must be tamper-evident: each
652        // file carries a 64-hex sha256 and a true byte count (used by
653        // the disk-space preflight), served over https.
654        for model in Catalog::seed().list() {
655            for file in &model.source.files {
656                assert!(
657                    file.url.starts_with("https://"),
658                    "{} must be https",
659                    file.url
660                );
661                let sha = file
662                    .sha256
663                    .as_deref()
664                    .unwrap_or_else(|| panic!("{} has no sha256 pin", file.filename));
665                assert_eq!(sha.len(), 64, "{} pin must be 64 hex", file.filename);
666                assert!(sha.chars().all(|c| c.is_ascii_hexdigit()));
667                assert!(
668                    file.approx_bytes.unwrap_or(0) > 0,
669                    "{} needs a real approx_bytes for the disk preflight",
670                    file.filename
671                );
672            }
673        }
674    }
675
676    #[test]
677    fn default_image_model_is_zimage() {
678        let catalog = Catalog::seed();
679        assert_eq!(
680            catalog.default_image_model().map(|m| m.id.as_str()),
681            Some("z-image-turbo-q4_k_m.gguf")
682        );
683    }
684
685    #[test]
686    fn json_round_trips() {
687        let catalog = Catalog::seed();
688        let json = catalog.to_json().unwrap();
689        // camelCase wire keys, mirroring the studio.
690        assert!(json.contains("\"displayName\""));
691        assert!(json.contains("\"cliDefaults\""));
692        assert!(json.contains("\"diffusion-model\""));
693        let parsed = Catalog::from_json(&json).unwrap();
694        assert_eq!(parsed, catalog);
695    }
696
697    #[test]
698    fn upsert_adds_then_replaces() {
699        let mut catalog = Catalog::default();
700        let mut model = zimage_turbo();
701        catalog.upsert(model.clone());
702        assert_eq!(catalog.list().len(), 1);
703
704        model.display_name = "Renamed".into();
705        catalog.upsert(model);
706        assert_eq!(catalog.list().len(), 1);
707        assert_eq!(
708            catalog
709                .get("z-image-turbo-q4_k_m.gguf")
710                .unwrap()
711                .display_name,
712            "Renamed"
713        );
714    }
715
716    #[test]
717    fn remove_reports_presence() {
718        let mut catalog = Catalog::seed();
719        assert!(catalog.remove("z-image-turbo-q4_k_m.gguf"));
720        assert!(!catalog.remove("z-image-turbo-q4_k_m.gguf"));
721        assert!(catalog.get("z-image-turbo-q4_k_m.gguf").is_none());
722    }
723
724    #[test]
725    fn load_or_seed_writes_then_reads_back() {
726        let dir = std::env::temp_dir().join(format!("sw-catalog-{}", std::process::id()));
727        let _ = std::fs::remove_dir_all(&dir);
728        let path = dir.join("models.json");
729
730        // Missing -> seeded + persisted.
731        let seeded = Catalog::load_or_seed(&path).unwrap();
732        assert!(path.exists());
733        assert!(seeded.get("z-image-turbo-q4_k_m.gguf").is_some());
734
735        // Existing -> read back unchanged.
736        let reloaded = Catalog::load_or_seed(&path).unwrap();
737        assert_eq!(reloaded, seeded);
738
739        let _ = std::fs::remove_dir_all(&dir);
740    }
741
742    #[test]
743    fn equality_derives_hold_for_catalog_model() {
744        assert_eq!(zimage_turbo(), zimage_turbo());
745    }
746
747    // -----------------------------------------------------------------
748    // Persistence safety: atomic writes + corrupt-file quarantine.
749    // -----------------------------------------------------------------
750
751    #[test]
752    fn save_atomically_replaces_without_temp_litter() {
753        // A second save must fully replace the file and leave no
754        // temp-file siblings from the write-then-rename dance — a crash
755        // mid-write must never truncate the operator's catalog.
756        let dir = tempfile::tempdir().unwrap();
757        let path = dir.path().join("models.json");
758        Catalog::seed().save(&path).unwrap();
759        let mut small = Catalog::default();
760        small.upsert(CatalogModel {
761            description: None,
762            ..zimage_turbo()
763        });
764        small.save(&path).unwrap();
765
766        // Read the bytes back raw: load_or_seed would top up the seeds.
767        let reloaded = Catalog::from_json(&std::fs::read_to_string(&path).unwrap()).unwrap();
768        assert_eq!(reloaded, small);
769
770        let names: Vec<String> = std::fs::read_dir(dir.path())
771            .unwrap()
772            .map(|e| e.unwrap().file_name().to_string_lossy().to_string())
773            .collect();
774        assert_eq!(
775            names,
776            vec!["models.json".to_string()],
777            "atomic save must leave only the target file, found: {names:?}"
778        );
779    }
780
781    #[test]
782    fn corrupt_catalog_is_quarantined_not_overwritten() {
783        // The exact data-loss shape this guards against: a corrupt
784        // models.json used to make the caller fall back to the seed
785        // while keeping the save path — the next persist silently
786        // destroyed the operator's hand-edited catalog.
787        let dir = tempfile::tempdir().unwrap();
788        let path = dir.path().join("models.json");
789        let operator_bytes = b"{ this is my hand-edited catalog, now corrupt";
790        std::fs::write(&path, operator_bytes).unwrap();
791
792        let logs = crate::test_support::capture({
793            let path = path.clone();
794            move || {
795                let recovered = Catalog::load_or_seed(&path).unwrap();
796                assert_eq!(recovered, Catalog::seed(), "reseeded in place");
797            }
798        });
799
800        // The original bytes survive in a quarantine sibling.
801        let quarantined: Vec<_> = std::fs::read_dir(dir.path())
802            .unwrap()
803            .map(|e| e.unwrap().file_name().to_string_lossy().to_string())
804            .filter(|n| n.starts_with("models.json.corrupt-"))
805            .collect();
806        assert_eq!(quarantined.len(), 1, "exactly one quarantine file");
807        assert_eq!(
808            std::fs::read(dir.path().join(&quarantined[0])).unwrap(),
809            operator_bytes,
810            "the operator's bytes must survive verbatim"
811        );
812        // The live path now holds a healthy seed.
813        assert_eq!(Catalog::load_or_seed(&path).unwrap(), Catalog::seed());
814        // And the recovery left a breadcrumb naming both paths.
815        assert!(logs.contains("quarantined"), "got: {logs}");
816        assert!(logs.contains("models.json.corrupt-"), "got: {logs}");
817    }
818
819    #[test]
820    fn load_for_serving_keeps_the_path_after_quarantine_recovery() {
821        let dir = tempfile::tempdir().unwrap();
822        let path = dir.path().join("models.json");
823        std::fs::write(&path, b"not json").unwrap();
824        let (catalog, save_path) = Catalog::load_for_serving(Some(path.clone()));
825        assert_eq!(catalog, Catalog::seed());
826        assert_eq!(
827            save_path,
828            Some(path),
829            "a quarantined-and-reseeded file is healthy; persistence stays on"
830        );
831    }
832
833    #[test]
834    fn load_for_serving_disables_persistence_when_the_file_is_unreadable() {
835        // A directory where the file should be makes the read fail with
836        // a non-NotFound error on every platform.  The worker must
837        // serve the seed but never gain a path it could clobber.
838        let dir = tempfile::tempdir().unwrap();
839        let path = dir.path().join("models.json");
840        std::fs::create_dir(&path).unwrap();
841        let (catalog, save_path) = Catalog::load_for_serving(Some(path));
842        assert_eq!(catalog, Catalog::seed());
843        assert_eq!(
844            save_path, None,
845            "an unreadable catalog must not be writable"
846        );
847    }
848
849    #[test]
850    fn load_for_serving_without_a_path_serves_the_seed() {
851        let (catalog, save_path) = Catalog::load_for_serving(None);
852        assert_eq!(catalog, Catalog::seed());
853        assert_eq!(save_path, None);
854    }
855
856    // -----------------------------------------------------------------
857    // sync_studio_model — mirror studio-offered models into the catalog
858    // without clobbering the operator's own entries.
859    // -----------------------------------------------------------------
860
861    /// Clone a model with a different id (test helper).
862    fn with_id(mut m: CatalogModel, id: &str) -> CatalogModel {
863        m.id = id.to_string();
864        m
865    }
866
867    fn studio_model(id: &str) -> CatalogModel {
868        with_id(
869            CatalogModel {
870                origin: "studio".into(),
871                ..zimage_turbo()
872            },
873            id,
874        )
875    }
876
877    #[test]
878    fn sync_adds_a_new_studio_model_and_reports_change() {
879        let mut cat = Catalog::default();
880        assert!(cat.sync_studio_model(studio_model("m1")));
881        assert_eq!(cat.get("m1").unwrap().origin, "studio");
882        // Re-syncing the identical model is a no-op (no file churn).
883        assert!(!cat.sync_studio_model(studio_model("m1")));
884    }
885
886    #[test]
887    fn sync_refreshes_a_changed_studio_model() {
888        let mut cat = Catalog::default();
889        cat.sync_studio_model(studio_model("m1"));
890        let mut updated = studio_model("m1");
891        updated.display_name = "Renamed by studio".into();
892        assert!(cat.sync_studio_model(updated));
893        assert_eq!(cat.get("m1").unwrap().display_name, "Renamed by studio");
894    }
895
896    #[test]
897    fn sync_never_clobbers_a_local_origin_entry() {
898        let mut cat = Catalog::default();
899        let mut local = with_id(zimage_turbo(), "m1");
900        local.display_name = "my hand-tuned model".into();
901        // origin defaults to "local".
902        cat.upsert(local);
903        // A studio offer for the same id must not overwrite it.
904        assert!(!cat.sync_studio_model(studio_model("m1")));
905        assert_eq!(cat.get("m1").unwrap().display_name, "my hand-tuned model");
906        assert_eq!(cat.get("m1").unwrap().origin, "local");
907    }
908}