flow_models_server/

fetch.rs

//! Engine provisioning for the managed local-model server.
//!
//! On first use we fetch the prebuilt server engine for the host's OS/arch and
//! unpack it into the data dir (`engines/`), so there is zero external setup.
//! This replaces the old Node fetch script.
//!
//! The release coordinates - which distribution, and which artifact per
//! platform - live in `engine_catalog.json` (operator/API data, kept out of
//! source). This module holds only the generic resolve → download → extract →
//! install logic, reporting progress so a long first download is never a silent
//! wait.

use std::io::Write;
use std::path::{Path, PathBuf};

use futures_util::StreamExt;
use serde::Deserialize;

/// A progress update emitted while fetching the engine.
#[derive(Debug, Clone)]
pub struct EngineFetchProgress {
    pub stage: EngineStage,
    /// Bytes done so far (meaningful during `Downloading`).
    pub current: u64,
    /// Total download size when the server reports it.
    pub total: Option<u64>,
    pub message: String,
}

/// Coarse phase of the engine fetch.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum EngineStage {
    Resolving,
    Downloading,
    Extracting,
    Installing,
}

/// Ensure a managed engine binary exists under `dest` (the data-dir
/// `engines/`), fetching the release for this host's OS/arch if absent.
/// Returns the path to the engine binary; reports progress through
/// `on_progress`.
pub async fn ensure_engine(
    dest: &Path,
    on_progress: impl Fn(EngineFetchProgress) + Send + Sync,
) -> Result<PathBuf, String> {
    let binary = dest.join(exe_name());
    if binary.is_file() {
        return Ok(binary);
    }
    std::fs::create_dir_all(dest).map_err(|e| format!("create {}: {e}", dest.display()))?;

    let catalog = load_catalog()?;
    let spec = catalog
        .assets
        .iter()
        .find(|a| a.os == host_os() && a.arch == host_arch())
        .ok_or_else(|| {
            format!(
                "no engine release configured for {}/{}",
                host_os(),
                host_arch()
            )
        })?;

    on_progress(EngineFetchProgress {
        stage: EngineStage::Resolving,
        current: 0,
        total: None,
        message: "Resolving engine release…".into(),
    });

    let client = reqwest::Client::builder()
        .user_agent(catalog.user_agent.clone())
        .build()
        .map_err(|e| format!("http client: {e}"))?;

    let release: Release = client
        .get(&catalog.release_api)
        .header(reqwest::header::ACCEPT, "application/vnd.github+json")
        .send()
        .await
        .and_then(|r| r.error_for_status())
        .map_err(|e| format!("fetch release metadata: {e}"))?
        .json()
        .await
        .map_err(|e| format!("parse release metadata: {e}"))?;

    let ext = spec.archive.ext();
    let asset = release
        .assets
        .iter()
        .find(|a| a.name.contains(&spec.name) && a.name.ends_with(ext))
        .ok_or_else(|| {
            format!(
                "no '{}'*{} asset in release {}",
                spec.name, ext, release.tag_name
            )
        })?;

    // Download (streamed) into a temp file, reporting byte progress.
    let tmp = std::env::temp_dir().join(format!("flow-engine-{}", std::process::id()));
    std::fs::create_dir_all(&tmp).map_err(|e| format!("create temp dir: {e}"))?;
    let archive_path = tmp.join(&asset.name);
    download(
        &client,
        &asset.browser_download_url,
        &archive_path,
        &on_progress,
    )
    .await?;

    on_progress(EngineFetchProgress {
        stage: EngineStage::Extracting,
        current: 0,
        total: None,
        message: "Extracting engine…".into(),
    });

    // Extraction + install is blocking I/O; keep it off the async runtime.
    let dest_owned = dest.to_path_buf();
    let archive = spec.archive;
    let tmp_for_task = tmp.clone();
    let binary = tokio::task::spawn_blocking(move || {
        extract_and_install(&archive_path, &dest_owned, archive, &tmp_for_task)
    })
    .await
    .map_err(|e| format!("extract task: {e}"))??;

    let _ = std::fs::remove_dir_all(&tmp);

    on_progress(EngineFetchProgress {
        stage: EngineStage::Installing,
        current: 0,
        total: None,
        message: "Engine ready".into(),
    });

    Ok(binary)
}

fn exe_name() -> &'static str {
    if cfg!(windows) {
        "llama-server.exe"
    } else {
        "llama-server"
    }
}

fn host_os() -> &'static str {
    if cfg!(target_os = "macos") {
        "macos"
    } else if cfg!(target_os = "windows") {
        "windows"
    } else {
        "linux"
    }
}

fn host_arch() -> &'static str {
    if cfg!(target_arch = "aarch64") {
        "arm64"
    } else {
        "x64"
    }
}

#[derive(Deserialize)]
struct EngineCatalog {
    release_api: String,
    user_agent: String,
    assets: Vec<AssetSpec>,
}

#[derive(Deserialize)]
struct AssetSpec {
    os: String,
    arch: String,
    /// Fragment that identifies this host's artifact within the release.
    name: String,
    archive: ArchiveKind,
}

#[derive(Deserialize, Clone, Copy)]
#[serde(rename_all = "lowercase")]
enum ArchiveKind {
    Targz,
    Zip,
}

impl ArchiveKind {
    fn ext(self) -> &'static str {
        match self {
            ArchiveKind::Targz => ".tar.gz",
            ArchiveKind::Zip => ".zip",
        }
    }
}

fn load_catalog() -> Result<EngineCatalog, String> {
    serde_json::from_str(include_str!("engine_catalog.json"))
        .map_err(|e| format!("engine catalog: {e}"))
}

#[derive(Deserialize)]
struct Release {
    tag_name: String,
    assets: Vec<ReleaseAsset>,
}

#[derive(Deserialize)]
struct ReleaseAsset {
    name: String,
    browser_download_url: String,
}

async fn download(
    client: &reqwest::Client,
    url: &str,
    dest_file: &Path,
    on_progress: &(impl Fn(EngineFetchProgress) + Send + Sync),
) -> Result<(), String> {
    let resp = client
        .get(url)
        .send()
        .await
        .and_then(|r| r.error_for_status())
        .map_err(|e| format!("download: {e}"))?;
    let total = resp.content_length();
    let mut file =
        std::fs::File::create(dest_file).map_err(|e| format!("create {}: {e}", dest_file.display()))?;
    let mut stream = resp.bytes_stream();
    let mut downloaded: u64 = 0;
    let mut last_emit: u64 = 0;
    on_progress(EngineFetchProgress {
        stage: EngineStage::Downloading,
        current: 0,
        total,
        message: "Downloading engine…".into(),
    });
    while let Some(chunk) = stream.next().await {
        let chunk = chunk.map_err(|e| format!("download stream: {e}"))?;
        file.write_all(&chunk).map_err(|e| format!("write: {e}"))?;
        downloaded += chunk.len() as u64;
        // Throttle to ~1 MiB steps so the event channel isn't flooded.
        if downloaded - last_emit >= (1 << 20) {
            last_emit = downloaded;
            on_progress(EngineFetchProgress {
                stage: EngineStage::Downloading,
                current: downloaded,
                total,
                message: "Downloading engine…".into(),
            });
        }
    }
    file.flush().map_err(|e| format!("flush: {e}"))?;
    Ok(())
}

fn extract_and_install(
    archive: &Path,
    dest: &Path,
    kind: ArchiveKind,
    tmp: &Path,
) -> Result<PathBuf, String> {
    let unpack = tmp.join("unpack");
    std::fs::create_dir_all(&unpack).map_err(|e| format!("create unpack dir: {e}"))?;
    match kind {
        ArchiveKind::Targz => {
            let f = std::fs::File::open(archive).map_err(|e| format!("open archive: {e}"))?;
            let gz = flate2::read::GzDecoder::new(f);
            tar::Archive::new(gz)
                .unpack(&unpack)
                .map_err(|e| format!("untar: {e}"))?;
        }
        ArchiveKind::Zip => {
            let f = std::fs::File::open(archive).map_err(|e| format!("open archive: {e}"))?;
            zip::ZipArchive::new(f)
                .map_err(|e| format!("open zip: {e}"))?
                .extract(&unpack)
                .map_err(|e| format!("unzip: {e}"))?;
        }
    }

    let engine_dir =
        find_engine_dir(&unpack).ok_or_else(|| format!("{} not found in archive", exe_name()))?;

    // Copy the engine dir's contents into `dest`, flattening symlinks to real
    // files (the engine resolves sibling shared libraries by name).
    copy_flat(&engine_dir, dest)?;

    let binary = dest.join(exe_name());
    if !binary.is_file() {
        return Err(format!(
            "engine binary missing after install: {}",
            binary.display()
        ));
    }
    make_runnable(dest, &binary);
    Ok(binary)
}

/// Depth-first search for the directory that holds the engine binary.
fn find_engine_dir(root: &Path) -> Option<PathBuf> {
    let mut stack = vec![root.to_path_buf()];
    while let Some(dir) = stack.pop() {
        let mut subdirs = Vec::new();
        for entry in std::fs::read_dir(&dir).ok()?.flatten() {
            if entry.file_name() == std::ffi::OsStr::new(exe_name()) {
                return Some(dir);
            }
            if entry.path().is_dir() {
                subdirs.push(entry.path());
            }
        }
        stack.extend(subdirs);
    }
    None
}

/// Recursively copy `src` into `dest`, replacing each symlink with a real copy
/// of its resolved target so every name the engine loads exists as a file.
fn copy_flat(src: &Path, dest: &Path) -> Result<(), String> {
    std::fs::create_dir_all(dest).map_err(|e| format!("mkdir {}: {e}", dest.display()))?;
    for entry in std::fs::read_dir(src).map_err(|e| format!("read {}: {e}", src.display()))? {
        let entry = entry.map_err(|e| format!("dir entry: {e}"))?;
        let from = entry.path();
        let to = dest.join(entry.file_name());
        let meta = std::fs::symlink_metadata(&from).map_err(|e| format!("stat: {e}"))?;
        if meta.file_type().is_symlink() {
            let real = std::fs::canonicalize(&from).map_err(|e| format!("resolve symlink: {e}"))?;
            if real.is_file() {
                std::fs::copy(&real, &to).map_err(|e| format!("copy: {e}"))?;
            }
        } else if meta.is_dir() {
            copy_flat(&from, &to)?;
        } else {
            std::fs::copy(&from, &to).map_err(|e| format!("copy: {e}"))?;
        }
    }
    Ok(())
}

/// Make the freshly-installed engine launchable: executable bit on Unix, and on
/// macOS clear the download quarantine + ad-hoc sign so Gatekeeper doesn't kill
/// the unsigned binary on launch.
fn make_runnable(dest: &Path, binary: &Path) {
    #[cfg(unix)]
    {
        use std::os::unix::fs::PermissionsExt;
        if let Ok(meta) = std::fs::metadata(binary) {
            let mut perms = meta.permissions();
            perms.set_mode(0o755);
            let _ = std::fs::set_permissions(binary, perms);
        }
    }
    #[cfg(target_os = "macos")]
    {
        let _ = std::process::Command::new("xattr")
            .args(["-dr", "com.apple.quarantine"])
            .arg(dest)
            .status();
        if let Ok(entries) = std::fs::read_dir(dest) {
            for entry in entries.flatten() {
                let path = entry.path();
                if path.is_file() {
                    let _ = std::process::Command::new("codesign")
                        .args(["--force", "--sign", "-"])
                        .arg(&path)
                        .status();
                }
            }
        }
    }
    // Silence unused-parameter warnings on platforms where a branch compiles out.
    let _ = (dest, binary);
}

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

    #[test]
    fn catalog_parses_and_covers_this_host() {
        let cat = load_catalog().expect("engine_catalog.json parses");
        assert!(
            cat.assets
                .iter()
                .any(|a| a.os == host_os() && a.arch == host_arch()),
            "no engine asset configured for {}/{}",
            host_os(),
            host_arch()
        );
        assert!(cat.release_api.starts_with("http"));
    }

    #[test]
    fn archive_ext_matches_kind() {
        assert_eq!(ArchiveKind::Targz.ext(), ".tar.gz");
        assert_eq!(ArchiveKind::Zip.ext(), ".zip");
    }
}