modelshelf 0.1.0

A shared local LLM model registry: discover, deduplicate, download, and update models across desktop apps.
Documentation
//! End-to-end scan tests against fake-but-structurally-real ecosystem trees
//! built inside a temp directory. No real user data is touched: every root
//! is injected through `ScanEnv`.

use std::collections::HashMap;
use std::path::Path;

use modelshelf::gguf::fixtures::{gguf_bytes, write_gguf};
use modelshelf::scan::fixtures;
use modelshelf::{Ecosystem, ModelQuery, OpenOptions, ScanEnv, ScanOptions, Shelf};

struct Fixture {
    _tmp: tempfile::TempDir,
    shelf: Shelf,
    options: ScanOptions,
    ollama_root: std::path::PathBuf,
    lmstudio_root: std::path::PathBuf,
    hf_root: std::path::PathBuf,
}

/// A shelf plus empty fake ecosystem roots, fully isolated in a tempdir.
fn fixture() -> Fixture {
    let tmp = tempfile::tempdir().unwrap();
    let base = tmp.path();
    let shelf_home = base.join("shelf");
    let fake_home = base.join("home"); // empty: no accidental real scans
    let ollama_root = base.join("ollama");
    let lmstudio_root = base.join("lmstudio");
    let hf_root = base.join("hf");
    for d in [&fake_home, &ollama_root, &lmstudio_root, &hf_root] {
        std::fs::create_dir_all(d).unwrap();
    }

    let shelf = Shelf::open(OpenOptions::at(&shelf_home).app_id("com.example.test")).unwrap();
    let options = ScanOptions {
        env: ScanEnv {
            home_dir: Some(fake_home),
            roots: HashMap::from([
                (Ecosystem::Ollama, ollama_root.clone()),
                (Ecosystem::Lmstudio, lmstudio_root.clone()),
                (Ecosystem::HfCache, hf_root.clone()),
            ]),
        },
        ..Default::default()
    };
    Fixture {
        _tmp: tmp,
        shelf,
        options,
        ollama_root,
        lmstudio_root,
        hf_root,
    }
}

fn lmstudio_path(root: &Path, repo: &str, file: &str) -> std::path::PathBuf {
    root.join(repo).join(file)
}

#[test]
fn scan_unifies_models_across_ecosystems_and_detects_duplicates() {
    let fx = fixture();

    // The same model content in three ecosystems...
    let shared = gguf_bytes("Shared Llama", "llama", 15, 8192, b"SHARED-CONTENT");
    let shared_hex =
        fixtures::ollama::add_model(&fx.ollama_root, "shared-llama", "latest", &shared);
    write_gguf(
        &lmstudio_path(
            &fx.lmstudio_root,
            "bartowski/Shared-Llama-GGUF",
            "Shared-Llama-Q4_K_M.gguf",
        ),
        "Shared Llama",
        "llama",
        15,
        8192,
        b"SHARED-CONTENT",
    );
    fixtures::hf_cache::add_model(
        &fx.hf_root,
        "bartowski/Shared-Llama-GGUF",
        "deadbeefrev",
        "Shared-Llama-Q4_K_M.gguf",
        &shared,
    );

    // ...and one model that exists only in LM Studio.
    write_gguf(
        &lmstudio_path(&fx.lmstudio_root, "Qwen/Solo-GGUF", "Solo-Q8_0.gguf"),
        "Solo Qwen",
        "qwen2",
        7,
        32768,
        b"SOLO-CONTENT",
    );

    let report = fx.shelf.scan(&fx.options).unwrap();
    assert_eq!(report.files_found, 4, "3 shared copies + 1 solo");
    assert_eq!(report.new_models, 2, "same content collapses to one model");
    assert!(report.skipped.is_empty(), "skipped: {:?}", report.skipped);

    let models = fx.shelf.list().unwrap();
    assert_eq!(models.len(), 2);

    let shared_model = models
        .iter()
        .find(|m| m.id.sha256_hex() == Some(shared_hex.as_str()))
        .expect("shared model registered under its content hash");
    assert_eq!(shared_model.locations.len(), 3);
    let ecosystems: Vec<Ecosystem> = shared_model.locations.iter().map(|l| l.ecosystem).collect();
    assert!(ecosystems.contains(&Ecosystem::Ollama));
    assert!(ecosystems.contains(&Ecosystem::Lmstudio));
    assert!(ecosystems.contains(&Ecosystem::HfCache));
    // Source mapping derived from the HF cache or LM Studio layout.
    assert_eq!(
        shared_model.source.as_ref().map(|s| s.repo.as_str()),
        Some("bartowski/Shared-Llama-GGUF")
    );

    // Duplicate detection: three physically distinct copies of `shared`.
    assert_eq!(report.duplicate_groups.len(), 1);
    let dup = &report.duplicate_groups[0];
    assert_eq!(dup.id, shared_model.id);
    assert_eq!(dup.distinct_copies, 3);
    assert_eq!(
        dup.reclaimable_bytes,
        shared_model.size_bytes * 2,
        "two of three copies are reclaimable"
    );

    // GGUF metadata was extracted for the plain-file model.
    let solo = models.iter().find(|m| m.id != shared_model.id).unwrap();
    let gguf = solo.gguf.as_ref().expect("gguf metadata parsed");
    assert_eq!(gguf.general_name.as_deref(), Some("Solo Qwen"));
    assert_eq!(gguf.quantization.as_deref(), Some("Q8_0"));
}

#[test]
fn rescan_is_idempotent_and_prunes_deleted_files() {
    let fx = fixture();
    let solo_path = lmstudio_path(&fx.lmstudio_root, "org/Repo-GGUF", "model-Q4_0.gguf");
    write_gguf(&solo_path, "Prunable", "llama", 2, 2048, b"PRUNE-ME");
    write_gguf(
        &lmstudio_path(&fx.lmstudio_root, "org/Keeper-GGUF", "keeper-Q4_0.gguf"),
        "Keeper",
        "llama",
        2,
        2048,
        b"KEEP-ME",
    );

    let first = fx.shelf.scan(&fx.options).unwrap();
    assert_eq!(first.new_models, 2);

    let second = fx.shelf.scan(&fx.options).unwrap();
    assert_eq!(second.new_models, 0, "rescan must not re-register models");
    assert_eq!(fx.shelf.list().unwrap().len(), 2);

    // Delete one file; its model (single location) disappears on rescan.
    std::fs::remove_file(&solo_path).unwrap();
    let third = fx.shelf.scan(&fx.options).unwrap();
    assert_eq!(third.locations_pruned, 1);
    let models = fx.shelf.list().unwrap();
    assert_eq!(models.len(), 1);
    assert_eq!(models[0].display_name, "Keeper");
}

#[test]
fn registry_survives_reopen_and_supports_find_and_refs() {
    let fx = fixture();
    write_gguf(
        &lmstudio_path(
            &fx.lmstudio_root,
            "meta/Llama-3.1-GGUF",
            "Llama-3.1-Q4_K_M.gguf",
        ),
        "Meta Llama 3.1 8B Instruct",
        "llama",
        15,
        131072,
        b"LLAMA31",
    );
    fx.shelf.scan(&fx.options).unwrap();

    // Reopen from disk: state persists.
    let reopened =
        Shelf::open(OpenOptions::at(fx.shelf.paths().home()).app_id("com.example.other")).unwrap();
    let hits = reopened
        .find(&ModelQuery::name_contains("llama-3.1").quant("Q4_K_M"))
        .unwrap();
    assert_eq!(hits.len(), 1);

    // Refs round-trip.
    let id = hits[0].id.clone();
    reopened.add_ref(&id, "summarizer").unwrap();
    let m = &reopened
        .find(&ModelQuery::name_contains("llama-3.1"))
        .unwrap()[0];
    assert_eq!(m.refs.len(), 1);
    assert_eq!(m.refs[0].app_id, "com.example.other");
    reopened.remove_ref(&id).unwrap();
    let m = &reopened
        .find(&ModelQuery::name_contains("llama-3.1"))
        .unwrap()[0];
    assert!(m.refs.is_empty());
}

#[test]
fn custom_roots_are_scanned() {
    let fx = fixture();
    let custom = fx.shelf.paths().home().parent().unwrap().join("my-models");
    write_gguf(
        &custom.join("nested/dir/custom-model.gguf"),
        "Custom Model",
        "phi3",
        18,
        4096,
        b"CUSTOM",
    );
    let mut options = fx.options.clone();
    options.extra_roots.push(custom);
    let report = fx.shelf.scan(&options).unwrap();
    assert_eq!(report.new_models, 1);
    let models = fx.shelf.list().unwrap();
    assert_eq!(models[0].locations[0].ecosystem, Ecosystem::Custom);
}