cortiq-gateway 0.2.43

//! Local CMF model server lifecycle.
//!
//! Closes the loop with <https://github.com/infosave2007/cmf>: install the
//! `cortiq` CLI from crates.io, keep it up to date, and run one `cortiq serve`
//! per configured local model as a managed child process. Each served model
//! exposes an OpenAI-compatible API, so the registry registers it as an ordinary
//! provider and the gateway routes to it like any other backend.

use crate::config::{CmfCfg, CmfServer};
use std::process::Stdio;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use tokio::process::{Child, Command};

/// Observable state of one managed local model server.
#[derive(Clone, Debug, Default, serde::Serialize)]
pub struct CmfServerStatus {
    /// Pool id it is registered under.
    pub id: String,
    /// Path of the `.cmf` being served.
    pub model: String,
    /// Port the server binds to.
    pub port: u16,
    /// OpenAI base URL the local server is registered under.
    pub base_url: String,
    /// The child `cortiq serve` process has been spawned.
    pub running: bool,
    /// `/healthz` on the local server answered 200.
    pub healthy: bool,
    /// Last fatal error for this server, if any.
    pub last_error: Option<String>,
}

/// Observable state of the managed local CMF layer (surfaced to the admin API).
#[derive(Clone, Debug, Default, serde::Serialize)]
pub struct CmfStatus {
    /// `manage_server` is on and at least one local model is configured.
    pub enabled: bool,
    /// Version reported by `cortiq --version` (None = not installed).
    pub installed_version: Option<String>,
    /// Newest `cortiq-cli` on crates.io (populated when `auto_update` runs).
    pub latest_version: Option<String>,
    /// Per-model server status.
    pub servers: Vec<CmfServerStatus>,
    /// Recent lifecycle log lines (bounded).
    pub log: Vec<String>,
}

/// A spawned server: its pool id and the child process handle.
struct Managed {
    id: String,
    child: Child,
}

/// Holds the managed child processes and their live status.
pub struct CmfRuntime {
    status: Mutex<CmfStatus>,
    children: Mutex<Vec<Managed>>,
    /// Server ids whose healthz handshake advertised tool support — lets the
    /// pipeline route tools traffic to local models with no manual caps.
    pub detected_tools: Mutex<std::collections::HashSet<String>>,
}

impl CmfRuntime {
    pub fn new() -> Arc<Self> {
        Arc::new(Self {
            status: Mutex::new(CmfStatus::default()),
            children: Mutex::new(Vec::new()),
            detected_tools: Mutex::new(std::collections::HashSet::new()),
        })
    }

    pub fn status(&self) -> CmfStatus {
        self.status.lock().unwrap().clone()
    }

    fn log(&self, msg: impl Into<String>) {
        let msg = msg.into();
        tracing::info!(target: "cmf", "{msg}");
        let mut s = self.status.lock().unwrap();
        s.log.push(msg);
        let n = s.log.len();
        if n > 80 {
            s.log.drain(0..n - 80);
        }
    }

    /// A global (install/update) failure — not tied to a single server.
    fn fail(&self, msg: impl Into<String>) {
        let msg = msg.into();
        tracing::warn!(target: "cmf", "{msg}");
        self.status.lock().unwrap().log.push(format!("✗ {msg}"));
    }

    /// Record a fatal error against a specific server (and the log).
    fn fail_server(&self, id: &str, msg: impl Into<String>) {
        let msg = msg.into();
        tracing::warn!(target: "cmf", "[{id}] {msg}");
        let mut s = self.status.lock().unwrap();
        if let Some(sv) = s.servers.iter_mut().find(|x| x.id == id) {
            sv.last_error = Some(msg.clone());
        }
        s.log.push(format!("✗ [{id}] {msg}"));
    }

    /// Seed the per-server status slots from the effective server list.
    fn init_servers(&self, servers: &[CmfServer], host: &str) {
        let mut s = self.status.lock().unwrap();
        s.enabled = true;
        s.servers = servers
            .iter()
            .map(|sv| CmfServerStatus {
                id: sv.id.clone(),
                model: sv.model.clone(),
                port: sv.port,
                base_url: format!("http://{host}:{}/v1", sv.port),
                running: false,
                healthy: false,
                last_error: None,
            })
            .collect();
    }

    fn set_running(&self, id: &str, running: bool) {
        if let Some(sv) = self
            .status
            .lock()
            .unwrap()
            .servers
            .iter_mut()
            .find(|x| x.id == id)
        {
            sv.running = running;
        }
    }

    fn set_healthy(&self, id: &str, healthy: bool) {
        if let Some(sv) = self
            .status
            .lock()
            .unwrap()
            .servers
            .iter_mut()
            .find(|x| x.id == id)
        {
            sv.healthy = healthy;
        }
    }

    fn add_child(&self, id: String, child: Child) {
        self.children.lock().unwrap().push(Managed {
            id: id.clone(),
            child,
        });
        self.set_running(&id, true);
    }

    fn is_running(&self, id: &str) -> bool {
        self.status
            .lock()
            .unwrap()
            .servers
            .iter()
            .any(|s| s.id == id && s.running)
    }

    /// Stop all managed servers (used on shutdown / reconfigure).
    pub async fn stop(&self) {
        let mut kids = std::mem::take(&mut *self.children.lock().unwrap());
        for m in kids.iter_mut() {
            let _ = m.child.start_kill();
        }
        let mut s = self.status.lock().unwrap();
        for sv in s.servers.iter_mut() {
            sv.running = false;
            sv.healthy = false;
        }
    }

    /// Stop and forget a single managed server (used when a model is deleted).
    pub async fn stop_one(&self, id: &str) {
        let removed = {
            let mut kids = self.children.lock().unwrap();
            kids.iter()
                .position(|m| m.id == id)
                .map(|pos| kids.remove(pos))
        };
        if let Some(mut m) = removed {
            let _ = m.child.start_kill();
        }
        self.status.lock().unwrap().servers.retain(|sv| sv.id != id);
    }
}

impl Default for CmfRuntime {
    fn default() -> Self {
        Self {
            status: Mutex::new(CmfStatus::default()),
            children: Mutex::new(Vec::new()),
            detected_tools: Mutex::new(std::collections::HashSet::new()),
        }
    }
}

/// crates.io `newest_version` for any crate.
async fn latest_version_of(krate: &str) -> Option<String> {
    let client = reqwest::Client::builder()
        .timeout(Duration::from_secs(10))
        .build()
        .ok()?;
    let body = client
        .get(format!("https://crates.io/api/v1/crates/{krate}"))
        .header("User-Agent", "cortiq-gateway")
        .send()
        .await
        .ok()?
        .text()
        .await
        .ok()?;
    let v: serde_json::Value = serde_json::from_str(&body).ok()?;
    v["crate"]["newest_version"].as_str().map(|s| s.to_string())
}

/// crates.io `newest_version` for `cortiq-cli`.
async fn latest_crates_version() -> Option<String> {
    latest_version_of("cortiq-cli").await
}

/// Newest published gateway version — the "update available" footer check.
pub(crate) async fn latest_gateway_version() -> Option<String> {
    latest_version_of("cortiq-gateway").await
}

/// `<bin> --version` → the trailing version token (e.g. "cortiq 0.1.2" → "0.1.2").
pub(crate) fn installed_version(bin: &str) -> Option<String> {
    let out = std::process::Command::new(bin)
        .arg("--version")
        .output()
        .ok()?;
    if !out.status.success() {
        return None;
    }
    let s = String::from_utf8_lossy(&out.stdout);
    s.split_whitespace().last().map(|s| s.trim().to_string())
}

fn parse_ver(s: &str) -> (u32, u32, u32) {
    let mut it = s.split('.').map(|x| x.trim().parse::<u32>().unwrap_or(0));
    (
        it.next().unwrap_or(0),
        it.next().unwrap_or(0),
        it.next().unwrap_or(0),
    )
}

pub(crate) fn version_lt(a: &str, b: &str) -> bool {
    parse_ver(a) < parse_ver(b)
}

/// `cargo install cortiq-cli [--force]` — blocking compile, run off-thread.
async fn cargo_install(force: bool) -> Result<(), String> {
    let mut cmd = Command::new("cargo");
    cmd.args(["install", "cortiq-cli", "--locked"]);
    if force {
        cmd.arg("--force");
    }
    let out = cmd.output().await.map_err(|e| e.to_string())?;
    if out.status.success() {
        Ok(())
    } else {
        Err(String::from_utf8_lossy(&out.stderr).trim().to_string())
    }
}

/// Install/upgrade `cortiq-cli` from crates.io on demand (the admin "Install
/// cortiq" button). Runs in the background; progress + result land in the CMF
/// status log, and `installed_version` is refreshed on success.
pub async fn install_now(rt: Arc<CmfRuntime>, bin: String) {
    rt.log("installing cortiq-cli from crates.io (this compiles — a few minutes)…");
    match cargo_install(true).await {
        Ok(()) => {
            let v = installed_version(&bin);
            rt.status.lock().unwrap().installed_version = v.clone();
            rt.log(format!(
                "installed cortiq-cli {}",
                v.as_deref().unwrap_or("?")
            ));
        }
        Err(e) => rt.fail(format!("cargo install cortiq-cli failed: {e}")),
    }
}

/// Spawn `cortiq serve <model> --host <host> --port <port>` for one server; the
/// child is killed when its handle is dropped (i.e. when the gateway shuts down).
fn spawn_serve(server: &CmfServer, cfg: &CmfCfg) -> Result<Child, String> {
    let mut cmd = Command::new(&cfg.cortiq_bin);
    cmd.arg("serve")
        .arg(&server.model)
        .arg("--host")
        .arg(&cfg.local_host)
        .arg("--port")
        .arg(server.port.to_string());
    // Parallelise decode matvecs across cores (CMF_THREADS). 0/1 = leave the
    // runtime's own default; N>1 gave ~2.8x on a 15B CPU model.
    if server.threads > 1 {
        cmd.env("CMF_THREADS", server.threads.to_string());
    }
    if server.gpu {
        cmd.env("CMF_GPU", "1");
    }
    // O(1) Nyström attention: layers spec + sub-parameters.
    if let Some(o1) = &server.o1 {
        let s = o1.trim();
        if !s.is_empty() {
            // Pass even "off" so it explicitly overrides a file converter hint.
            cmd.arg("--o1").arg(s);
        }
    }
    if let Some(m) = server.o1_m {
        cmd.arg("--o1-m").arg(m.to_string());
    }
    if let Some(w) = server.o1_window {
        cmd.arg("--o1-window").arg(w.to_string());
    }
    if let Some(sink) = server.o1_sink {
        cmd.arg("--o1-sink").arg(sink.to_string());
    }
    // MTP (Multi-Token Prediction) speculative decoding: the engine reads
    // CMF_MTP env; "0" disables.  There is no CLI flag for this.
    if server.skip_mtp {
        cmd.env("CMF_MTP", "0");
    }
    cmd.stdout(Stdio::null())
        .stderr(Stdio::null())
        .kill_on_drop(true)
        .spawn()
        .map_err(|e| e.to_string())
}

/// Poll `/healthz` until the local server answers 200 (or the deadline
/// passes). Newer `cortiq serve` advertises the loaded model's capabilities
/// in the healthz body ({"capabilities":{"tools":true}}) — returned so the
/// caller can register them; older servers return `None` there.
async fn wait_healthy(host: &str, port: u16, tries: u32) -> (bool, Option<bool>) {
    let url = format!("http://{host}:{port}/healthz");
    let client = reqwest::Client::new();
    for _ in 0..tries {
        if let Ok(r) = client
            .get(&url)
            .timeout(Duration::from_secs(2))
            .send()
            .await
        {
            if r.status().is_success() {
                let tools = r
                    .json::<serde_json::Value>()
                    .await
                    .ok()
                    .and_then(|v| v["capabilities"]["tools"].as_bool());
                return (true, tools);
            }
        }
        tokio::time::sleep(Duration::from_millis(500)).await;
    }
    (false, None)
}

/// Orchestrate the whole lifecycle: ensure installed → maybe update → spawn each
/// configured local model → wait healthy. Safe to call in a background task; it
/// never panics.
pub async fn manage(rt: Arc<CmfRuntime>, cfg: CmfCfg) {
    let servers = cfg.effective_servers();
    if servers.is_empty() {
        return;
    }
    rt.init_servers(&servers, &cfg.local_host);

    // 1. Ensure the `cortiq` binary is available (once for all servers).
    let mut ver = installed_version(&cfg.cortiq_bin);
    if ver.is_none() {
        if cfg.auto_install {
            rt.log("cortiq not found — installing cortiq-cli from crates.io (this compiles, may take a few minutes)…");
            if let Err(e) = cargo_install(false).await {
                rt.fail(format!("cargo install cortiq-cli failed: {e}"));
                return;
            }
            ver = installed_version(&cfg.cortiq_bin);
            if ver.is_some() {
                rt.log("installed cortiq-cli");
            }
        } else {
            rt.fail(
                "cortiq binary not found (set [cmf].auto_install = true, or install cortiq-cli)",
            );
            return;
        }
    }
    rt.status.lock().unwrap().installed_version = ver.clone();

    // 2. Optionally update to the newest crates.io release.
    if cfg.auto_update {
        if let Some(latest) = latest_crates_version().await {
            rt.status.lock().unwrap().latest_version = Some(latest.clone());
            if ver
                .as_deref()
                .map(|c| version_lt(c, &latest))
                .unwrap_or(false)
            {
                rt.log(format!(
                    "updating cortiq-cli {} → {latest}…",
                    ver.as_deref().unwrap_or("?")
                ));
                if let Err(e) = cargo_install(true).await {
                    rt.log(format!("update failed (keeping current): {e}"));
                } else {
                    ver = installed_version(&cfg.cortiq_bin);
                    rt.status.lock().unwrap().installed_version = ver.clone();
                    rt.log(format!("updated to {}", ver.as_deref().unwrap_or("?")));
                }
            }
        }
    }

    // 3. Spawn each server (fast).
    for server in &servers {
        if !std::path::Path::new(&server.model).exists() {
            rt.fail_server(
                &server.id,
                format!("model file not found: {}", server.model),
            );
            continue;
        }
        rt.log(format!(
            "starting: cortiq serve {} --host {} --port {}",
            server.model, cfg.local_host, server.port
        ));
        match spawn_serve(server, &cfg) {
            Ok(child) => rt.add_child(server.id.clone(), child),
            Err(e) => rt.fail_server(&server.id, format!("failed to spawn: {e}")),
        }
    }

    // 4. Wait until each spawned server is serving (model load can be slow).
    for server in &servers {
        if !rt.is_running(&server.id) {
            continue; // failed to spawn / file missing
        }
        let (ok, tools) = wait_healthy(&cfg.local_host, server.port, 240).await;
        if ok {
            rt.set_healthy(&server.id, true);
            // capability handshake: the server said whether its model's chat
            // template can drive tool calling — no manual caps needed
            if tools == Some(true) {
                rt.detected_tools.lock().unwrap().insert(server.id.clone());
                tracing::info!(id = %server.id, "local server advertises tool support");
            } else {
                rt.detected_tools.lock().unwrap().remove(&server.id);
            }
            rt.log(format!(
                "local CMF server ready — '{}' at http://{}:{}/v1",
                server.id, cfg.local_host, server.port
            ));
        } else {
            rt.fail_server(&server.id, "did not become healthy in time");
        }
    }
}