IdeA/crates/infrastructure/src/model_server/mod.rs

//! Infrastructure adapters for local model servers.

use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::Mutex;
use std::time::Duration;

use async_trait::async_trait;
use futures_util::StreamExt;
use serde::{Deserialize, Serialize};
use tokio::io::AsyncWriteExt;
use tokio::process::{Child, Command};

use domain::model_server::{
    ExecutablePath, HfModelRef, LlamaCppOptions, LocalModelRef, LocalModelServerConfig,
    LocalModelServerKind, ModelPath, ModelServerEndpoint, ModelSource,
};
use domain::ports::{
    FileSystem, ManagedProcess, ManagedProcessHandle, ModelArtifactCancel, ModelArtifactDownloader,
    ModelArtifactProgress, ModelArtifactResolution, ModelServerArgv, ModelServerError,
    ModelServerProbe, ModelServerRegistry, ModelServerRuntime, ProcessStatus, RemotePath,
    SpawnSpec,
};
use domain::{LocalModelServerId, ProjectPath, StopPolicy};

/// HTTP readiness probe for OpenAI-compatible servers.
#[derive(Clone)]
pub struct HttpOpenAiCompatibleProbe {
    client: reqwest::Client,
}

impl Default for HttpOpenAiCompatibleProbe {
    fn default() -> Self {
        Self::new(Duration::from_secs(2))
    }
}

impl HttpOpenAiCompatibleProbe {
    /// Builds the probe with a request timeout.
    #[must_use]
    pub fn new(timeout: Duration) -> Self {
        let client = reqwest::Client::builder()
            .timeout(timeout)
            .build()
            .expect("reqwest client with timeout builds");
        Self { client }
    }
}

#[async_trait]
impl ModelServerProbe for HttpOpenAiCompatibleProbe {
    async fn probe(
        &self,
        endpoint: &ModelServerEndpoint,
    ) -> Result<domain::ModelServerStatus, ModelServerError> {
        let models = format!("{}/models", endpoint.base_url.trim_end_matches('/'));
        if is_ready(self.client.get(models).send().await) {
            return Ok(domain::ModelServerStatus::ReadyReused);
        }
        let root = endpoint.base_url.trim_end_matches("/v1").to_owned();
        if is_ready(self.client.get(root).send().await) {
            return Ok(domain::ModelServerStatus::ReadyReused);
        }
        Ok(domain::ModelServerStatus::Unreachable)
    }
}

fn is_ready(result: Result<reqwest::Response, reqwest::Error>) -> bool {
    result
        .map(|response| response.status().is_success())
        .unwrap_or(false)
}

/// Hugging Face artifact resolver/downloader backed by an IdeA-owned cache dir.
#[derive(Clone)]
pub struct HfModelArtifactDownloader {
    cache_dir: PathBuf,
    client: reqwest::Client,
}

impl HfModelArtifactDownloader {
    /// Builds a downloader that stores artifacts under `cache_dir`.
    #[must_use]
    pub fn new(cache_dir: impl Into<PathBuf>) -> Self {
        Self {
            cache_dir: cache_dir.into(),
            client: reqwest::Client::new(),
        }
    }

    /// Deterministic cache path for a Hugging Face model reference.
    #[must_use]
    pub fn cache_path_for(&self, repo: &HfModelRef) -> PathBuf {
        let (base, quant) = split_hf_ref(repo);
        let file_stem = quant.unwrap_or("model");
        self.cache_dir
            .join(base.replace('/', "--"))
            .join(format!("{}.gguf", safe_cache_component(file_stem)))
    }

    async fn resolve_remote_filename(&self, repo: &HfModelRef) -> Result<String, ModelServerError> {
        let (base, quant) = split_hf_ref(repo);
        let url = format!("https://huggingface.co/api/models/{base}");
        let response = self
            .client
            .get(url)
            .send()
            .await
            .map_err(|e| ModelServerError::Probe(e.to_string()))?;
        if !response.status().is_success() {
            return Err(ModelServerError::Probe(format!(
                "huggingface model metadata returned {}",
                response.status()
            )));
        }
        let metadata: HfModelMetadata = response
            .json()
            .await
            .map_err(|e| ModelServerError::Probe(e.to_string()))?;
        select_gguf_file(&metadata.siblings, quant).ok_or_else(|| {
            ModelServerError::PathNotAccessible(format!(
                "no matching .gguf artifact found for {}",
                repo.as_str()
            ))
        })
    }
}

#[async_trait]
impl ModelArtifactDownloader for HfModelArtifactDownloader {
    async fn resolve_hf_model(
        &self,
        repo: &HfModelRef,
        progress: std::sync::Arc<dyn Fn(ModelArtifactProgress) + Send + Sync>,
        cancel: ModelArtifactCancel,
    ) -> Result<ModelArtifactResolution, ModelServerError> {
        if cancel.is_cancelled() {
            return Err(ModelServerError::Cancelled);
        }
        let path = self.cache_path_for(repo);
        if path.is_file() {
            return Ok(ModelArtifactResolution {
                path: model_path_from_pathbuf(path)?,
                cache_hit: true,
            });
        }

        let filename = self.resolve_remote_filename(repo).await?;
        if cancel.is_cancelled() {
            return Err(ModelServerError::Cancelled);
        }
        let (base, _) = split_hf_ref(repo);
        let url = format!(
            "https://huggingface.co/{base}/resolve/main/{}",
            url_path_segment(&filename)
        );
        let response = self
            .client
            .get(url)
            .send()
            .await
            .map_err(|e| ModelServerError::Probe(e.to_string()))?;
        if !response.status().is_success() {
            return Err(ModelServerError::Probe(format!(
                "huggingface artifact download returned {}",
                response.status()
            )));
        }
        if let Some(parent) = path.parent() {
            tokio::fs::create_dir_all(parent)
                .await
                .map_err(|e| ModelServerError::Store(e.to_string()))?;
        }
        let tmp_path = path.with_extension("gguf.part");
        let mut file = tokio::fs::File::create(&tmp_path)
            .await
            .map_err(|e| ModelServerError::Store(e.to_string()))?;
        let total = response.content_length();
        let mut downloaded = 0_u64;
        let mut stream = response.bytes_stream();
        while let Some(chunk) = stream.next().await {
            if cancel.is_cancelled() {
                let _ = tokio::fs::remove_file(&tmp_path).await;
                return Err(ModelServerError::Cancelled);
            }
            let chunk = chunk.map_err(|e| ModelServerError::Probe(e.to_string()))?;
            file.write_all(&chunk)
                .await
                .map_err(|e| ModelServerError::Store(e.to_string()))?;
            downloaded += chunk.len() as u64;
            progress(ModelArtifactProgress {
                downloaded_bytes: Some(downloaded),
                total_bytes: total,
                source: Some(repo.as_str().to_owned()),
            });
        }
        file.flush()
            .await
            .map_err(|e| ModelServerError::Store(e.to_string()))?;
        drop(file);
        tokio::fs::rename(&tmp_path, &path)
            .await
            .map_err(|e| ModelServerError::Store(e.to_string()))?;
        Ok(ModelArtifactResolution {
            path: model_path_from_pathbuf(path)?,
            cache_hit: false,
        })
    }
}

#[derive(Debug, Deserialize)]
struct HfModelMetadata {
    #[serde(default)]
    siblings: Vec<HfSibling>,
}

#[derive(Debug, Deserialize)]
struct HfSibling {
    #[serde(rename = "rfilename")]
    filename: String,
}

fn split_hf_ref(repo: &HfModelRef) -> (&str, Option<&str>) {
    repo.as_str()
        .split_once(':')
        .map_or((repo.as_str(), None), |(base, quant)| (base, Some(quant)))
}

fn select_gguf_file(siblings: &[HfSibling], quant: Option<&str>) -> Option<String> {
    let mut ggufs: Vec<&str> = siblings
        .iter()
        .map(|sibling| sibling.filename.as_str())
        .filter(|filename| filename.ends_with(".gguf"))
        .collect();
    ggufs.sort_unstable();
    let Some(quant) = quant else {
        return ggufs.first().map(|filename| (*filename).to_owned());
    };
    let quant = quant.to_ascii_lowercase();
    ggufs
        .into_iter()
        .find(|filename| filename.to_ascii_lowercase().contains(&quant))
        .map(str::to_owned)
}

fn safe_cache_component(raw: &str) -> String {
    raw.chars()
        .map(|ch| match ch {
            'A'..='Z' | 'a'..='z' | '0'..='9' | '.' | '_' | '-' => ch,
            _ => '_',
        })
        .collect()
}

fn url_path_segment(raw: &str) -> String {
    raw.split('/')
        .map(url_component)
        .collect::<Vec<_>>()
        .join("/")
}

fn url_component(raw: &str) -> String {
    raw.bytes()
        .flat_map(|byte| match byte {
            b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'.' | b'_' | b'-' => {
                vec![byte as char]
            }
            other => format!("%{other:02X}").chars().collect(),
        })
        .collect()
}

fn model_path_from_pathbuf(path: PathBuf) -> Result<ModelPath, ModelServerError> {
    ModelPath::new(path.to_string_lossy().into_owned())
        .map_err(|e| ModelServerError::Invalid(e.to_string()))
}

/// Builds `llama-server` argv without shell interpolation.
#[derive(Debug, Default, Clone, Copy)]
pub struct LlamaCppRuntime;

impl LlamaCppRuntime {
    /// Creates the runtime.
    #[must_use]
    pub const fn new() -> Self {
        Self
    }
}

impl ModelServerRuntime for LlamaCppRuntime {
    fn build_argv(
        &self,
        config: &LocalModelServerConfig,
    ) -> Result<ModelServerArgv, ModelServerError> {
        if config.kind != LocalModelServerKind::LlamaCpp {
            return Err(ModelServerError::Invalid(
                "only LlamaCpp local model servers are supported".to_owned(),
            ));
        }
        let command = resolve_binary(
            config
                .binary
                .as_ref()
                .map(|binary| binary.as_str())
                .unwrap_or("llama-server"),
        )?;
        let mut args = Vec::new();
        match config
            .model
            .source
            .as_ref()
            .ok_or_else(|| ModelServerError::PathNotAccessible("model.source missing".to_owned()))?
        {
            ModelSource::LocalPath { path } => {
                args.push("--model".to_owned());
                args.push(path.as_str().to_owned());
            }
            ModelSource::HuggingFace { repo } => {
                args.push("-hf".to_owned());
                args.push(repo.as_str().to_owned());
            }
        }
        args.extend([
            "--port".to_owned(),
            config.endpoint.port.to_string(),
            "--host".to_owned(),
            config.options.host.clone(),
        ]);
        if let Some(gpu_layers) = config.options.gpu_layers {
            args.extend(["-ngl".to_owned(), gpu_layers.to_string()]);
        }
        if let Some(context_size) = config.options.context_size {
            args.extend(["-c".to_owned(), context_size.to_string()]);
        }
        if config.options.jinja {
            args.push("--jinja".to_owned());
        }
        args.extend(config.args.clone());
        Ok(ModelServerArgv { command, args })
    }

    fn build_spawn_spec(
        &self,
        config: &LocalModelServerConfig,
    ) -> Result<SpawnSpec, ModelServerError> {
        let argv = self.build_argv(config)?;
        Ok(SpawnSpec {
            command: argv.command,
            args: argv.args,
            cwd: ProjectPath::new("/").map_err(|e| ModelServerError::Invalid(e.to_string()))?,
            env: Vec::new(),
            context_plan: None,
            sandbox: None,
        })
    }
}

fn resolve_binary(raw: &str) -> Result<String, ModelServerError> {
    if is_path_like(raw) {
        let path = Path::new(raw);
        if path.is_file() {
            return Ok(raw.to_owned());
        }
        return Err(ModelServerError::PathNotAccessible(raw.to_owned()));
    }
    if let Some(path) = find_in_path(raw) {
        return Ok(path.to_string_lossy().into_owned());
    }
    Err(ModelServerError::PathNotAccessible(format!(
        "{raw} not found in PATH"
    )))
}

fn is_path_like(raw: &str) -> bool {
    raw.contains('/') || raw.contains('\\') || Path::new(raw).is_absolute()
}

fn find_in_path(command: &str) -> Option<PathBuf> {
    std::env::var_os("PATH").and_then(|path| {
        std::env::split_paths(&path)
            .map(|dir| dir.join(command))
            .find(|candidate| candidate.is_file())
    })
}

/// Local child-process manager for long-lived model servers.
#[derive(Default)]
pub struct LocalManagedProcess {
    children: Mutex<HashMap<String, Child>>,
}

impl LocalManagedProcess {
    /// Creates an empty process manager.
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }
}

#[async_trait]
impl ManagedProcess for LocalManagedProcess {
    async fn spawn(&self, spec: SpawnSpec) -> Result<ManagedProcessHandle, ModelServerError> {
        let mut command = Command::new(&spec.command);
        command.args(&spec.args);
        if spec.cwd.as_str() != "/" {
            command.current_dir(spec.cwd.as_str());
        }
        for (key, value) in &spec.env {
            command.env(key, value);
        }
        let child = command
            .spawn()
            .map_err(|e| ModelServerError::Process(format!("{}: {e}", spec.command)))?;
        let id = uuid::Uuid::new_v4().to_string();
        self.children.lock().unwrap().insert(id.clone(), child);
        Ok(ManagedProcessHandle { id })
    }

    async fn kill(&self, handle: &ManagedProcessHandle) -> Result<(), ModelServerError> {
        let Some(mut child) = self.children.lock().unwrap().remove(&handle.id) else {
            return Ok(());
        };
        child
            .start_kill()
            .map_err(|e| ModelServerError::Process(e.to_string()))
    }

    async fn status(
        &self,
        handle: &ManagedProcessHandle,
    ) -> Result<ProcessStatus, ModelServerError> {
        let mut children = self.children.lock().unwrap();
        let Some(child) = children.get_mut(&handle.id) else {
            return Ok(ProcessStatus::Unknown);
        };
        match child
            .try_wait()
            .map_err(|e| ModelServerError::Process(e.to_string()))?
        {
            Some(status) => Ok(ProcessStatus::Exited {
                code: status.code(),
            }),
            None => Ok(ProcessStatus::Running),
        }
    }
}

/// File name of the global local-model-server registry.
const MODEL_SERVERS_FILE: &str = "model-servers.json";
const MODEL_SERVERS_VERSION: u32 = 2;

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
struct ModelServersDoc {
    version: u32,
    servers: Vec<LocalModelServerConfig>,
}

#[derive(Debug, Clone, Deserialize)]
#[serde(untagged)]
enum PersistedModelServersDoc {
    V2(ModelServersDoc),
    V1(ModelServersDocV1),
}

#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
struct ModelServersDocV1 {
    servers: Vec<LocalModelServerConfigV1>,
}

#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
struct LocalModelServerConfigV1 {
    id: LocalModelServerId,
    kind: LocalModelServerKind,
    name: String,
    endpoint: ModelServerEndpoint,
    model: LocalModelRefV1,
    #[serde(default)]
    binary: Option<ExecutablePath>,
    #[serde(default)]
    args: Vec<String>,
    auto_start: bool,
    #[serde(default = "default_stop_policy")]
    stop_policy: StopPolicy,
}

#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
struct LocalModelRefV1 {
    id: String,
    label: String,
    #[serde(default)]
    path: Option<ModelPath>,
    served_name: String,
}

fn default_stop_policy() -> StopPolicy {
    StopPolicy::StopOnAppExit
}

impl From<ModelServersDocV1> for ModelServersDoc {
    fn from(doc: ModelServersDocV1) -> Self {
        Self {
            version: MODEL_SERVERS_VERSION,
            servers: doc
                .servers
                .into_iter()
                .map(|server| LocalModelServerConfig {
                    id: server.id,
                    kind: server.kind,
                    name: server.name,
                    endpoint: server.endpoint,
                    model: LocalModelRef {
                        id: server.model.id,
                        label: server.model.label,
                        source: server
                            .model
                            .path
                            .map(|path| ModelSource::LocalPath { path }),
                        served_name: server.model.served_name,
                    },
                    binary: server.binary,
                    options: LlamaCppOptions::default(),
                    args: server.args,
                    auto_start: server.auto_start,
                    stop_policy: server.stop_policy,
                    warmup_deadline_secs: None,
                })
                .collect(),
        }
    }
}

impl Default for ModelServersDoc {
    fn default() -> Self {
        Self {
            version: MODEL_SERVERS_VERSION,
            servers: Vec::new(),
        }
    }
}

/// Filesystem-backed global model-server registry.
#[derive(Clone)]
pub struct FsModelServerRegistry {
    fs: std::sync::Arc<dyn FileSystem>,
    app_data_dir: String,
}

impl FsModelServerRegistry {
    /// Builds the registry from a filesystem and global app-data dir.
    #[must_use]
    pub fn new(fs: std::sync::Arc<dyn FileSystem>, app_data_dir: impl Into<String>) -> Self {
        Self {
            fs,
            app_data_dir: app_data_dir.into(),
        }
    }

    fn path(&self) -> RemotePath {
        let base = self.app_data_dir.trim_end_matches(['/', '\\']);
        RemotePath::new(format!("{base}/{MODEL_SERVERS_FILE}"))
    }

    async fn read_doc(&self) -> Result<ModelServersDoc, ModelServerError> {
        match self.fs.read(&self.path()).await {
            Ok(bytes) => match serde_json::from_slice::<PersistedModelServersDoc>(&bytes)
                .map_err(|e| ModelServerError::Store(format!("serialization failed: {e}")))?
            {
                PersistedModelServersDoc::V2(doc) => Ok(ModelServersDoc {
                    version: MODEL_SERVERS_VERSION,
                    servers: doc.servers,
                }),
                PersistedModelServersDoc::V1(doc) => Ok(doc.into()),
            },
            Err(domain::ports::FsError::NotFound(_)) => Ok(ModelServersDoc::default()),
            Err(domain::ports::FsError::PermissionDenied(p)) => {
                Err(ModelServerError::PermissionDenied(p))
            }
            Err(e) => Err(ModelServerError::Store(e.to_string())),
        }
    }

    async fn write_doc(&self, doc: &ModelServersDoc) -> Result<(), ModelServerError> {
        let dir = RemotePath::new(self.app_data_dir.trim_end_matches(['/', '\\']).to_owned());
        self.fs.create_dir_all(&dir).await.map_err(|e| match e {
            domain::ports::FsError::PermissionDenied(p) => ModelServerError::PermissionDenied(p),
            other => ModelServerError::Store(other.to_string()),
        })?;
        let bytes = serde_json::to_vec_pretty(doc)
            .map_err(|e| ModelServerError::Store(format!("serialization failed: {e}")))?;
        self.fs
            .write(&self.path(), &bytes)
            .await
            .map_err(|e| match e {
                domain::ports::FsError::PermissionDenied(p) => {
                    ModelServerError::PermissionDenied(p)
                }
                other => ModelServerError::Store(other.to_string()),
            })
    }
}

#[async_trait]
impl ModelServerRegistry for FsModelServerRegistry {
    async fn get(
        &self,
        id: &LocalModelServerId,
    ) -> Result<Option<LocalModelServerConfig>, ModelServerError> {
        Ok(self
            .read_doc()
            .await?
            .servers
            .into_iter()
            .find(|server| &server.id == id))
    }

    async fn list(&self) -> Result<Vec<LocalModelServerConfig>, ModelServerError> {
        Ok(self.read_doc().await?.servers)
    }

    async fn save(&self, config: LocalModelServerConfig) -> Result<(), ModelServerError> {
        config
            .options
            .validate()
            .map_err(|e| ModelServerError::Invalid(e.to_string()))?;
        domain::model_server::validate_free_args(&config.args)
            .map_err(|e| ModelServerError::Invalid(e.to_string()))?;
        let mut doc = self.read_doc().await?;
        doc.version = MODEL_SERVERS_VERSION;
        if let Some(slot) = doc.servers.iter_mut().find(|server| server.id == config.id) {
            *slot = config;
        } else {
            doc.servers.push(config);
        }
        self.write_doc(&doc).await
    }

    async fn delete(&self, id: LocalModelServerId) -> Result<(), ModelServerError> {
        let mut doc = self.read_doc().await?;
        doc.servers.retain(|server| server.id != id);
        self.write_doc(&doc).await
    }
}