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 serde::{Deserialize, Serialize};
use tokio::process::{Child, Command};

use domain::model_server::{LocalModelServerConfig, LocalModelServerKind, ModelServerEndpoint};
use domain::ports::{
    FileSystem, ManagedProcess, ManagedProcessHandle, ModelServerError, ModelServerProbe,
    ModelServerRegistry, ModelServerRuntime, ProcessStatus, RemotePath, SpawnSpec,
};
use domain::{LocalModelServerId, ProjectPath};

/// 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)
}

/// 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_spawn_spec(
        &self,
        config: &LocalModelServerConfig,
    ) -> Result<SpawnSpec, ModelServerError> {
        if config.kind != LocalModelServerKind::LlamaCpp {
            return Err(ModelServerError::Invalid(
                "only LlamaCpp local model servers are supported".to_owned(),
            ));
        }
        let model_path =
            config.model.path.as_ref().ok_or_else(|| {
                ModelServerError::PathNotAccessible("model.path missing".to_owned())
            })?;
        let command = resolve_binary(
            config
                .binary
                .as_ref()
                .map(|binary| binary.as_str())
                .unwrap_or("llama-server"),
        )?;
        let mut args = vec![
            "--model".to_owned(),
            model_path.as_str().to_owned(),
            "--port".to_owned(),
            config.endpoint.port.to_string(),
        ];
        args.extend(config.args.clone());
        Ok(SpawnSpec {
            command,
            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 = 1;

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

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) => serde_json::from_slice(&bytes)
                .map_err(|e| ModelServerError::Store(format!("serialization failed: {e}"))),
            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> {
        let mut doc = self.read_doc().await?;
        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
    }
}