agent-manager/src/models.rs

//! Local model runtimes (axe 11): inventory of models available on
//! ollama, llama.cpp (llama-server) and LM Studio, with an on-disk cache
//! following the probe.rs pattern (keyed by runtime identity + timestamp),
//! and a prune helper for unused models.
//!
//! Rules:
//! * a missing runtime is reported, never fatal;
//! * the cache is bounded and short-lived (default TTL 30 s);
//! * parsing functions are pure and unit-tested with fixtures.

use crate::app::App;
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
use std::path::{Path, PathBuf};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use wait_timeout::ChildExt;

/// Default time-to-live of a cached runtime inventory, in seconds.
pub const CACHE_TTL_SECS: u64 = 30;

/// Default number of days before an unused model becomes a prune candidate.
pub const DEFAULT_PRUNE_DAYS: u64 = 30;

/// A model found on a local runtime.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct RuntimeModel {
    /// Runtime name: "ollama", "llama.cpp" or "lm-studio".
    pub runtime: String,
    /// Model identifier (e.g. "llama3.1:8b").
    pub name: String,
    /// Disk size in bytes when known.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub size_bytes: Option<u64>,
    /// Quantization level when known (e.g. "Q4_K_M").
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub quantization: Option<String>,
    /// RFC 3339 timestamp of the last use/modification when known.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub last_used: Option<String>,
    /// Absolute path of the model file (llama.cpp) or empty.
    #[serde(default, skip_serializing_if = "String::is_empty")]
    pub path: String,
}

/// State of one runtime after detection.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RuntimeInfo {
    pub runtime: String,
    /// Whether the runtime was found (binary on PATH / API reachable).
    pub detected: bool,
    /// Human-readable detail (version, endpoint, or why it was not found).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub note: Option<String>,
}

/// Full inventory: runtime states plus the models of the detected runtimes.
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct ModelsInventory {
    pub runtimes: Vec<RuntimeInfo>,
    pub models: Vec<RuntimeModel>,
}

impl ModelsInventory {
    /// The models of one runtime (when detected).
    pub fn models_of(&self, runtime: &str) -> Vec<&RuntimeModel> {
        self.models
            .iter()
            .filter(|m| m.runtime == runtime)
            .collect()
    }

    /// Look up a model by name across every runtime (name matched exactly,
    /// or by short name when it matches a single model).
    pub fn find(&self, name: &str) -> Option<&RuntimeModel> {
        let exact: Vec<&RuntimeModel> = self
            .models
            .iter()
            .filter(|m| m.name == name)
            .collect();
        match exact.len() {
            0 => {
                // Short name (e.g. "llama3"): unique suffix match wins.
                let short: Vec<&RuntimeModel> = self
                    .models
                    .iter()
                    .filter(|m| m.name.split(':').next() == Some(name))
                    .collect();
                (short.len() == 1).then(|| short[0])
            }
            1 => Some(exact[0]),
            _ => {
                // Ambiguous: prefer the "latest" tag, then the first.
                exact
                    .iter()
                    .find(|m| m.name.ends_with(":latest"))
                    .copied()
                    .or(Some(exact[0]))
            }
        }
    }

    /// Is at least one runtime detected?
    pub fn any_runtime(&self) -> bool {
        self.runtimes.iter().any(|r| r.detected)
    }
}

// ---------------------------------------------------------------------------
// Detection
// ---------------------------------------------------------------------------

/// Detect every local model runtime and their models, using the on-disk
/// cache. A steady-state run does no subprocess and no HTTP request.
pub fn detect(app: &App) -> ModelsInventory {
    let cache_path = app.events_dir().join("models-cache.json");
    let mut cache: ModelsCache = std::fs::read_to_string(&cache_path)
        .ok()
        .and_then(|t| serde_json::from_str(&t).ok())
        .unwrap_or_default();
    let now = now_secs();

    let ollama = cached_or_scan(
        &mut cache,
        "ollama",
        &ollama_key(),
        now,
        CACHE_TTL_SECS,
        detect_ollama,
    );
    let llamacpp = cached_or_scan(
        &mut cache,
        "llama.cpp",
        &llamacpp_key(),
        now,
        CACHE_TTL_SECS,
        detect_llamacpp,
    );
    let lmstudio = cached_or_scan(
        &mut cache,
        "lm-studio",
        &"api".to_string(),
        now,
        CACHE_TTL_SECS,
        detect_lmstudio,
    );

    if let Ok(text) = serde_json::to_string_pretty(&cache) {
        let _ = std::fs::write(&cache_path, text);
    }

    let mut runtimes = Vec::new();
    let mut models = Vec::new();
    for (info, ms) in [ollama, llamacpp, lmstudio] {
        runtimes.push(info);
        models.extend(ms);
    }
    ModelsInventory { runtimes, models }
}

/// Run `scan` unless a fresh cache entry exists; stores the result.
fn cached_or_scan(
    cache: &mut ModelsCache,
    runtime: &str,
    key: &str,
    now: u64,
    ttl: u64,
    scan: fn() -> Vec<RuntimeModel>,
) -> (RuntimeInfo, Vec<RuntimeModel>) {
    let fresh = cache
        .entries
        .get(runtime)
        .filter(|e| e.key == key && now.saturating_sub(e.fetched_at) < ttl);
    if let Some(entry) = fresh {
        return (entry.info.clone(), entry.models.clone());
    }
    let models = scan();
    let info = runtime_info(runtime, &models);
    cache.entries.insert(
        runtime.to_string(),
        CacheEntry {
            key: key.to_string(),
            fetched_at: now,
            info: info.clone(),
            models: models.clone(),
        },
    );
    (info, models)
}

/// Describe a runtime from its models: detected when there is at least one
/// model (or the binary was found even with zero models).
fn runtime_info(runtime: &str, models: &[RuntimeModel]) -> RuntimeInfo {
    match runtime {
        "ollama" => RuntimeInfo {
            runtime: runtime.to_string(),
            detected: which::which("ollama").is_ok(),
            note: if models.is_empty() {
                Some("binary not found on PATH, or 'ollama list' returned nothing".to_string())
            } else {
                Some(format!("{} model(s) via 'ollama list'", models.len()))
            },
        },
        "llama.cpp" => RuntimeInfo {
            runtime: runtime.to_string(),
            detected: !models.is_empty(),
            note: if models.is_empty() {
                Some("llama-server/llama-cli not found, or no .gguf under the model directories".to_string())
            } else {
                Some(format!("{} model(s) found on disk", models.len()))
            },
        },
        _ => RuntimeInfo {
            runtime: runtime.to_string(),
            detected: !models.is_empty(),
            note: if models.is_empty() {
                Some("no response from http://127.0.0.1:1234/v1/models".to_string())
            } else {
                Some(format!("{} model(s) via the local API", models.len()))
            },
        },
    }
}

/// Identity key of the ollama binary (path + mtime + size), so a stale
/// cache entry is never reused after an ollama update.
fn ollama_key() -> String {
    which::which("ollama")
        .map(|p| bin_key(&p))
        .unwrap_or_else(|_| "missing".to_string())
}

/// Identity key of the llama.cpp binaries (llama-server or llama-cli).
fn llamacpp_key() -> String {
    let bin = which::which("llama-server")
        .or_else(|_| which::which("llama-cli"))
        .unwrap_or_default();
    if bin.as_os_str().is_empty() {
        "missing".to_string()
    } else {
        bin_key(&bin)
    }
}

/// Run `ollama list --format json` and parse the model list. Any failure
/// (binary missing, timeout, bad output) yields an empty list.
fn detect_ollama() -> Vec<RuntimeModel> {
    let Ok(bin) = which::which("ollama") else {
        return Vec::new();
    };
    let output = match run_capture(&bin, &["list", "--format", "json"]) {
        Some(o) => o,
        None => return Vec::new(),
    };
    parse_ollama_list(&output).unwrap_or_default()
}

/// Parse the `ollama list --format json` output.
fn parse_ollama_list(json: &str) -> Option<Vec<RuntimeModel>> {
    #[derive(Deserialize)]
    struct OllamaModel {
        name: String,
        #[serde(default)]
        size: Option<u64>,
        #[serde(default)]
        details: Option<OllamaDetails>,
        #[serde(default)]
        modified_at: Option<String>,
    }
    #[derive(Deserialize)]
    struct OllamaDetails {
        #[serde(default)]
        quantization_level: Option<String>,
    }
    #[derive(Deserialize)]
    struct List {
        #[serde(default)]
        models: Vec<OllamaModel>,
    }
    let list: List = serde_json::from_str(json).ok()?;
    Some(
        list.models
            .into_iter()
            .map(|m| RuntimeModel {
                runtime: "ollama".to_string(),
                name: m.name,
                size_bytes: m.size,
                quantization: m.details.and_then(|d| d.quantization_level),
                last_used: m.modified_at,
                path: String::new(),
            })
            .collect(),
    )
}

/// Scan the local disk for llama.cpp models (.gguf files) under the known
/// model directories. The quantization is read from the file name when
/// possible; the last use is the file modification time.
fn detect_llamacpp() -> Vec<RuntimeModel> {
    let _ = which::which("llama-server").or_else(|_| which::which("llama-cli"));
    let mut dirs: Vec<PathBuf> = Vec::new();
    if let Some(h) = crate::config::home_dir() {
        dirs.push(h.join("models"));
        dirs.push(h.join(".cache").join("llama.cpp"));
        dirs.push(h.join(".llama").join("models"));
    }
    if let Some(d) = std::env::var_os("LLAMA_MODELS_DIR") {
        dirs.push(PathBuf::from(d));
    }
    let mut out = Vec::new();
    let mut seen: BTreeMap<String, ()> = BTreeMap::new();
    for dir in dirs {
        let Ok(rd) = std::fs::read_dir(&dir) else {
            continue;
        };
        for entry in rd.flatten() {
            let p = entry.path();
            if !p.is_file() || !p.to_string_lossy().to_lowercase().ends_with(".gguf") {
                continue;
            }
            if seen.insert(p.display().to_string(), ()).is_some() {
                continue;
            }
            let meta = std::fs::metadata(&p).ok();
            let size = meta.as_ref().map(|m| m.len());
            let mtime = meta
                .as_ref()
                .and_then(|m| m.modified().ok())
                .map(rfc3339_from_system_time);
            let name = p
                .file_stem()
                .and_then(|n| n.to_str())
                .unwrap_or("model")
                .to_string();
            out.push(RuntimeModel {
                runtime: "llama.cpp".to_string(),
                name,
                size_bytes: size,
                quantization: quantization_from_name(&p.display().to_string()),
                last_used: mtime,
                path: p.display().to_string(),
            });
        }
    }
    out.sort_by(|a, b| a.name.cmp(&b.name));
    out
}

/// Extract a quantization token (Q4_K_M, Q8_0, ...) from a file name.
fn quantization_from_name(name: &str) -> Option<String> {
    // Patterns: "Q4_K_M", "Q8_0", "q4_0", "IQ4_XS", "Q5_K_S", ...
    let up = name.to_uppercase();
    let bytes: Vec<char> = up.chars().collect();
    let mut i = 0;
    while i + 1 < bytes.len() {
        if bytes[i] == 'Q' && bytes[i + 1].is_ascii_digit() {
            let start = i;
            let mut j = i + 1;
            let mut digits = 0;
            while j < bytes.len() && (bytes[j].is_ascii_digit() || bytes[j] == '_') {
                if bytes[j].is_ascii_digit() {
                    digits += 1;
                }
                j += 1;
            }
            // Suffix letters allowed: Q4_K_M, Q5_K_S, Q6_K, Q3_K_L, ...
            while j < bytes.len()
                && (bytes[j].is_ascii_uppercase() || bytes[j] == '_')
                && j - start < 10
            {
                j += 1;
            }
            // Also match the "IQ4_XS" style (I + Q + digits).
            let token_start = if start > 0 && bytes[start - 1] == 'I' {
                start - 1
            } else {
                start
            };
            let token: String = bytes[token_start..j].iter().collect();
            if digits >= 1
                && token.len() >= 3
                && token
                    .chars()
                    .all(|c| c.is_ascii_uppercase() || c == '_' || c.is_ascii_digit())
            {
                return Some(token);
            }
            i = j;
        } else {
            i += 1;
        }
    }
    None
}

/// Query the LM Studio local API (OpenAI-compatible, default port 1234).
fn detect_lmstudio() -> Vec<RuntimeModel> {
    let url = std::env::var("LM_STUDIO_URL").unwrap_or_else(|_| "http://127.0.0.1:1234/v1/models".to_string());
    let resp = ureq::get(&url)
        .timeout(Duration::from_secs(3))
        .call();
    let Ok(resp) = resp else {
        return Vec::new();
    };
    let Ok(text) = resp.into_string() else {
        return Vec::new();
    };
    parse_lmstudio_list(&text).unwrap_or_default()
}

/// Parse the LM Studio /v1/models response: {"data": [{"id": ...}]}.
fn parse_lmstudio_list(json: &str) -> Option<Vec<RuntimeModel>> {
    #[derive(Deserialize)]
    struct LmModel {
        id: String,
    }
    #[derive(Deserialize)]
    struct List {
        data: Vec<LmModel>,
    }
    let list: List = serde_json::from_str(json).ok()?;
    Some(
        list.data
            .into_iter()
            .map(|m| RuntimeModel {
                runtime: "lm-studio".to_string(),
                name: m.id,
                size_bytes: None,
                quantization: None,
                last_used: None,
                path: String::new(),
            })
            .collect(),
    )
}

// ---------------------------------------------------------------------------
// Prune
// ---------------------------------------------------------------------------

/// Candidates for pruning: models not used within `days` days. The
/// default model of a runtime (any name tagged ":latest" for ollama, and
/// the first model of a runtime when it has no tag) is never a candidate.
pub fn prune_candidates(models: &[RuntimeModel], days: u64) -> Vec<&RuntimeModel> {
    let cutoff = chrono::Utc::now() - chrono::Duration::days(days as i64);
    models
        .iter()
        .filter(|m| !is_default_model(m))
        .filter(|m| {
            let Some(lu) = &m.last_used else {
                // Unknown last use: keep the model (do not guess).
                return false;
            };
            chrono::DateTime::parse_from_rfc3339(lu)
                .map(|t| t.with_timezone(&chrono::Utc) < cutoff)
                .unwrap_or(false)
        })
        .collect()
}

/// Is this model the default of its runtime (protected from pruning)?
fn is_default_model(m: &RuntimeModel) -> bool {
    m.name.ends_with(":latest")
}

/// Remove one model from its runtime. ollama models are removed with
/// `ollama rm`; llama.cpp models are files that are deleted. LM Studio has
/// no deletion endpoint: an error explains why.
pub fn remove_model(m: &RuntimeModel) -> anyhow::Result<()> {
    match m.runtime.as_str() {
        "ollama" => {
            let bin = which::which("ollama")
                .map_err(|_| anyhow::anyhow!("ollama binary not found on PATH"))?;
            let status = std::process::Command::new(&bin)
                .arg("rm")
                .arg(&m.name)
                .stdin(std::process::Stdio::null())
                .stdout(std::process::Stdio::null())
                .stderr(std::process::Stdio::piped())
                .status()?;
            if status.success() {
                Ok(())
            } else {
                Err(anyhow::anyhow!("'ollama rm {}' failed", m.name))
            }
        }
        "llama.cpp" => {
            if m.path.is_empty() {
                return Err(anyhow::anyhow!(
                    "model '{}' has no known file path — nothing to delete",
                    m.name
                ));
            }
            std::fs::remove_file(&m.path)
                .map_err(|e| anyhow::anyhow!("cannot delete {}: {e}", m.path))
        }
        _ => Err(anyhow::anyhow!(
            "LM Studio has no deletion API — remove model '{}' from the LM Studio UI",
            m.name
        )),
    }
}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

/// Run a binary with arguments, capturing stdout, with a hard timeout.
/// Returns None when the binary fails, times out or cannot run.
fn run_capture(bin: &Path, args: &[&str]) -> Option<String> {
    use std::io::Read;
    let mut cmd = std::process::Command::new(bin);
    cmd.args(args)
        .stdin(std::process::Stdio::null())
        .stdout(std::process::Stdio::piped())
        .stderr(std::process::Stdio::null());
    let mut child = cmd.spawn().ok()?;
    let status = match child.wait_timeout(Duration::from_secs(8)) {
        Ok(Some(s)) => s,
        Ok(None) => {
            let _ = child.kill();
            let _ = child.wait();
            return None;
        }
        Err(_) => return None,
    };
    if !status.success() {
        return None;
    }
    let mut out = String::new();
    child.stdout.take()?.read_to_string(&mut out).ok()?;
    Some(out)
}

fn now_secs() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|d| d.as_secs())
        .unwrap_or(0)
}

/// RFC 3339 string of a system time (best effort).
fn rfc3339_from_system_time(t: SystemTime) -> String {
    let dt: chrono::DateTime<chrono::Utc> = t.into();
    dt.to_rfc3339()
}

/// Identity key of a binary: path + modification time + size (probe.rs style).
fn bin_key(bin: &Path) -> String {
    let meta = std::fs::metadata(bin).ok();
    let mtime = meta
        .as_ref()
        .and_then(|m| m.modified().ok())
        .and_then(|t| t.duration_since(UNIX_EPOCH).ok())
        .map(|d| d.as_secs())
        .unwrap_or(0);
    let size = meta.map(|m| m.len()).unwrap_or(0);
    format!("{}|{mtime}|{size}", bin.display())
}

#[derive(Debug, Default, Serialize, Deserialize)]
struct ModelsCache {
    entries: BTreeMap<String, CacheEntry>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
struct CacheEntry {
    key: String,
    fetched_at: u64,
    info: RuntimeInfo,
    models: Vec<RuntimeModel>,
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;

    const OLLAMA_FIXTURE: &str = r#"{
      "models": [
        {
          "name": "llama3.1:8b",
          "size": 4661247078,
          "details": { "quantization_level": "Q4_K_M" },
          "modified_at": "2026-08-10T12:00:00Z"
        },
        {
          "name": "llama3.1:latest",
          "size": 4661247078,
          "details": { "quantization_level": "Q4_K_M" },
          "modified_at": "2026-08-17T12:00:00Z"
        },
        {
          "name": "qwen2.5-coder:7b",
          "size": 4700000000,
          "details": { "quantization_level": "Q8_0" },
          "modified_at": "2026-01-05T09:00:00Z"
        }
      ]
    }"#;

    #[test]
    fn parses_ollama_list() {
        let models = parse_ollama_list(OLLAMA_FIXTURE).unwrap();
        assert_eq!(models.len(), 3);
        assert_eq!(models[0].name, "llama3.1:8b");
        assert_eq!(models[0].runtime, "ollama");
        assert_eq!(models[0].quantization.as_deref(), Some("Q4_K_M"));
        assert_eq!(models[0].last_used.as_deref(), Some("2026-08-10T12:00:00Z"));
        assert_eq!(models[0].size_bytes, Some(4661247078));
    }

    #[test]
    fn parses_ollama_empty_list() {
        let models = parse_ollama_list(r#"{"models": []}"#).unwrap();
        assert!(models.is_empty());
    }

    #[test]
    fn ollama_bad_output_yields_none() {
        assert!(parse_ollama_list("not json").is_none());
    }

    #[test]
    fn parses_lmstudio_list() {
        let models = parse_lmstudio_list(
            r#"{"object":"list","data":[{"id":"llama-3.1-8b","object":"model"},{"id":"codestral"}]}"#,
        )
        .unwrap();
        assert_eq!(models.len(), 2);
        assert_eq!(models[0].runtime, "lm-studio");
        assert_eq!(models[0].name, "llama-3.1-8b");
        assert_eq!(models[0].size_bytes, None);
    }

    #[test]
    fn finds_models_exactly_and_by_short_name() {
        let inv = ModelsInventory {
            runtimes: vec![],
            models: vec![
                RuntimeModel { runtime: "ollama".into(), name: "llama3.1:8b".into(), size_bytes: None, quantization: None, last_used: None, path: String::new() },
                RuntimeModel { runtime: "ollama".into(), name: "llama3.1:latest".into(), size_bytes: None, quantization: None, last_used: None, path: String::new() },
                RuntimeModel { runtime: "lm-studio".into(), name: "codestral".into(), size_bytes: None, quantization: None, last_used: None, path: String::new() },
            ],
        };
        assert_eq!(inv.find("llama3.1:8b").unwrap().name, "llama3.1:8b");
        // Short name that is ambiguous (two llama3.1 variants) is not resolved.
        assert!(inv.find("llama3.1").is_none());
        // A short name matching exactly one model works.
        assert_eq!(inv.find("codestral").unwrap().name, "codestral");
        assert!(inv.find("nope").is_none());
    }

    #[test]
    fn prune_skips_latest_and_recent_models() {
        let models = parse_ollama_list(OLLAMA_FIXTURE).unwrap();
        // qwen was last used in January: candidate; llama3.1:8b (Aug 10)
        // is recent enough with 30 days; :latest is protected.
        let cands = prune_candidates(&models, 30);
        assert_eq!(cands.len(), 1);
        assert_eq!(cands[0].name, "qwen2.5-coder:7b");
        // With 7 days, the Aug 10 model becomes a candidate too.
        let cands = prune_candidates(&models, 7);
        let names: Vec<&str> = cands.iter().map(|m| m.name.as_str()).collect();
        assert!(names.contains(&"llama3.1:8b"));
        assert!(names.contains(&"qwen2.5-coder:7b"));
        assert!(!names.contains(&"llama3.1:latest"));
    }

    #[test]
    fn models_without_last_used_are_kept() {
        let m = RuntimeModel { runtime: "lm-studio".into(), name: "x".into(), size_bytes: None, quantization: None, last_used: None, path: String::new() };
        assert!(prune_candidates(&[m.clone()], 1).is_empty());
        // Unparsable timestamps are also kept.
        let mut m2 = m.clone();
        m2.last_used = Some("garbage".into());
        assert!(prune_candidates(&[m2], 1).is_empty());
    }

    #[test]
    fn quantization_is_read_from_file_names() {
        assert_eq!(quantization_from_name("~/models/llama-3.1.Q4_K_M.gguf").as_deref(), Some("Q4_K_M"));
        assert_eq!(quantization_from_name("model-q8_0.gguf").as_deref(), Some("Q8_0"));
        assert_eq!(quantization_from_name("plain.gguf"), None);
    }

    #[test]
    fn bin_key_is_stable_and_distinct() {
        let dir = tempfile::tempdir().unwrap();
        let a = dir.path().join("a");
        std::fs::write(&a, b"abc").unwrap();
        let k1 = bin_key(&a);
        std::fs::write(&a, b"abcd").unwrap();
        let k2 = bin_key(&a);
        assert_ne!(k1, k2);
    }

    #[test]
    fn cache_roundtrip() {
        let mut cache = ModelsCache::default();
        let info = RuntimeInfo { runtime: "ollama".into(), detected: true, note: None };
        cache.entries.insert(
            "ollama".into(),
            CacheEntry { key: "k".into(), fetched_at: 42, info, models: vec![] },
        );
        let text = serde_json::to_string(&cache).unwrap();
        let back: ModelsCache = serde_json::from_str(&text).unwrap();
        assert_eq!(back.entries["ollama"].fetched_at, 42);
    }
}