IdeA/crates/domain/src/model_server.rs

//! Local model-server configuration and status.
//!
//! These value objects describe a local OpenAI-compatible server that can back an
//! OpenCode profile. They are pure persisted configuration: probing, process
//! spawning and filesystem checks live behind ports.

use std::net::Ipv4Addr;

use serde::{Deserialize, Deserializer, Serialize};

use crate::error::DomainError;
use crate::ids::LocalModelServerId;

/// Local model server implementation family.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub enum LocalModelServerKind {
    /// `llama-server` from llama.cpp.
    LlamaCpp,
}

/// HTTP endpoint for an OpenAI-compatible model server.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ModelServerEndpoint {
    /// Normalised base URL, always ending in `/v1`.
    #[serde(rename = "baseURL")]
    pub base_url: String,
    /// TCP port expected for the local server.
    pub port: u16,
}

impl ModelServerEndpoint {
    /// Builds a validated endpoint and normalises its base URL to `/v1`.
    ///
    /// # Errors
    /// Returns [`DomainError`] if the URL is blank, non-HTTP(S), or the port is 0.
    pub fn new(base_url: impl Into<String>, port: u16) -> Result<Self, DomainError> {
        let base_url = normalise_base_url(base_url.into())?;
        if port == 0 {
            return Err(DomainError::Invariant(
                "modelServer.endpoint.port must be in 1..=65535".to_owned(),
            ));
        }
        Ok(Self { base_url, port })
    }
}

fn normalise_base_url(raw: String) -> Result<String, DomainError> {
    let trimmed = raw.trim().trim_end_matches('/').to_owned();
    crate::validation::non_empty(&trimmed, "modelServer.endpoint.baseURL")?;
    if !(trimmed.starts_with("http://") || trimmed.starts_with("https://")) {
        return Err(DomainError::Invariant(
            "modelServer.endpoint.baseURL must start with http:// or https://".to_owned(),
        ));
    }
    if !has_local_host_authority(&trimmed) {
        return Err(DomainError::Invariant(
            "modelServer.endpoint.baseURL must target localhost in V1".to_owned(),
        ));
    }
    if trimmed.ends_with("/v1") {
        Ok(trimmed)
    } else {
        Ok(format!("{trimmed}/v1"))
    }
}

fn has_local_host_authority(url: &str) -> bool {
    let Some(authority_and_path) = url.split_once("://").map(|(_, rest)| rest) else {
        return false;
    };
    let authority = authority_and_path.split('/').next().unwrap_or_default();
    let host = authority
        .rsplit_once('@')
        .map(|(_, host)| host)
        .unwrap_or(authority);
    let host = host
        .strip_prefix('[')
        .and_then(|ipv6| ipv6.split_once(']').map(|(addr, _)| addr))
        .unwrap_or_else(|| host.split(':').next().unwrap_or_default());
    host.eq_ignore_ascii_case("localhost") || host == "::1" || is_ipv4_loopback(host)
}

fn is_ipv4_loopback(host: &str) -> bool {
    host.parse::<Ipv4Addr>()
        .map(|addr| addr.octets()[0] == 127)
        .unwrap_or(false)
}

/// Absolute local path to a model file.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(transparent)]
pub struct ModelPath(pub String);

impl ModelPath {
    /// Builds a validated absolute model path.
    ///
    /// # Errors
    /// Returns [`DomainError`] if the path is blank or relative.
    pub fn new(path: impl Into<String>) -> Result<Self, DomainError> {
        let path = path.into();
        crate::validation::non_empty(&path, "modelServer.model.path")?;
        if !is_absolute_path(&path) {
            return Err(DomainError::Invariant(
                "modelServer.model.path must be absolute".to_owned(),
            ));
        }
        Ok(Self(path))
    }

    /// Returns the path as a string slice.
    #[must_use]
    pub fn as_str(&self) -> &str {
        &self.0
    }
}

/// Explicit source of the served model.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase", tag = "type")]
pub enum ModelSource {
    /// Local `.gguf` file path.
    LocalPath {
        /// Absolute local path.
        path: ModelPath,
    },
    /// Hugging Face llama.cpp `-hf` reference.
    HuggingFace {
        /// `namespace/repo` or `namespace/repo:quant`.
        repo: HfModelRef,
    },
}

/// Hugging Face model reference accepted by `llama-server -hf`.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(transparent)]
pub struct HfModelRef(pub String);

impl HfModelRef {
    /// Builds a validated Hugging Face model reference.
    ///
    /// # Errors
    /// Returns [`DomainError`] if the reference is blank, URL-like, contains
    /// spaces, or does not match `namespace/repo[:quant]`.
    pub fn new(repo: impl Into<String>) -> Result<Self, DomainError> {
        let repo = repo.into();
        crate::validation::non_empty(&repo, "modelServer.model.source.repo")?;
        if repo.chars().any(char::is_whitespace) {
            return Err(DomainError::Invariant(
                "modelServer.model.source.repo must not contain spaces".to_owned(),
            ));
        }
        if repo.starts_with("http://") || repo.starts_with("https://") {
            return Err(DomainError::Invariant(
                "modelServer.model.source.repo must not be an URL in V1".to_owned(),
            ));
        }
        if repo.matches(':').count() > 1 {
            return Err(DomainError::Invariant(
                "modelServer.model.source.repo must contain at most one ':'".to_owned(),
            ));
        }
        let (base, quant) = repo
            .split_once(':')
            .map_or((repo.as_str(), None), |(base, quant)| (base, Some(quant)));
        let Some((namespace, name)) = base.split_once('/') else {
            return Err(DomainError::Invariant(
                "modelServer.model.source.repo must be namespace/repo or namespace/repo:quant"
                    .to_owned(),
            ));
        };
        if namespace.is_empty() || name.is_empty() || base.matches('/').count() != 1 {
            return Err(DomainError::Invariant(
                "modelServer.model.source.repo must be namespace/repo or namespace/repo:quant"
                    .to_owned(),
            ));
        }
        if quant.is_some_and(str::is_empty) {
            return Err(DomainError::Invariant(
                "modelServer.model.source.repo quant must not be empty".to_owned(),
            ));
        }
        Ok(Self(repo))
    }

    /// Returns the reference as a string slice.
    #[must_use]
    pub fn as_str(&self) -> &str {
        &self.0
    }
}

/// Absolute local path to an executable.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(transparent)]
pub struct ExecutablePath(pub String);

impl ExecutablePath {
    /// Builds a validated executable path.
    ///
    /// Absolute paths are accepted directly. Bare command names are accepted so
    /// infrastructure can resolve them through `PATH` without shell interpolation.
    ///
    /// # Errors
    /// Returns [`DomainError`] if the path is blank.
    pub fn new(path: impl Into<String>) -> Result<Self, DomainError> {
        let path = path.into();
        crate::validation::non_empty(&path, "modelServer.binary")?;
        Ok(Self(path))
    }

    /// Returns the path as a string slice.
    #[must_use]
    pub fn as_str(&self) -> &str {
        &self.0
    }
}

fn is_absolute_path(path: &str) -> bool {
    path.starts_with('/')
        || path.starts_with('\\')
        || (path.len() >= 3
            && path.as_bytes()[1] == b':'
            && path.as_bytes()[0].is_ascii_alphabetic()
            && matches!(path.as_bytes()[2], b'/' | b'\\'))
}

/// Model served by a local model server.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct LocalModelRef {
    /// Stable model id in IdeA configuration.
    pub id: String,
    /// Human display label.
    pub label: String,
    /// Explicit source. Required for auto-start.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub source: Option<ModelSource>,
    /// Name exposed by the server and used in `OpenCodeConfig.model`.
    pub served_name: String,
}

impl LocalModelRef {
    /// Builds a validated local model reference.
    ///
    /// # Errors
    /// Returns [`DomainError`] if `id`, `label` or `served_name` is blank.
    pub fn new(
        id: impl Into<String>,
        label: impl Into<String>,
        source: Option<ModelSource>,
        served_name: impl Into<String>,
    ) -> Result<Self, DomainError> {
        let id = id.into();
        let label = label.into();
        let served_name = served_name.into();
        crate::validation::non_empty(&id, "modelServer.model.id")?;
        crate::validation::non_empty(&label, "modelServer.model.label")?;
        crate::validation::non_empty(&served_name, "modelServer.model.servedName")?;
        Ok(Self {
            id,
            label,
            source,
            served_name,
        })
    }
}

/// Structured llama.cpp server options owned by IdeA.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct LlamaCppOptions {
    /// Host passed to `--host`.
    pub host: String,
    /// GPU layers passed to `-ngl`.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub gpu_layers: Option<u32>,
    /// Context size passed to `-c`.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub context_size: Option<u32>,
    /// Enables llama.cpp Jinja templates.
    pub jinja: bool,
}

/// Minimum configurable server warmup deadline in seconds.
pub const MIN_WARMUP_DEADLINE_SECS: u64 = 30;

/// Maximum configurable server warmup deadline in seconds.
pub const MAX_WARMUP_DEADLINE_SECS: u64 = 1_800;

impl Default for LlamaCppOptions {
    fn default() -> Self {
        Self {
            host: "127.0.0.1".to_owned(),
            gpu_layers: None,
            context_size: None,
            jinja: false,
        }
    }
}

impl LlamaCppOptions {
    /// Builds validated llama.cpp options.
    ///
    /// # Errors
    /// Returns [`DomainError`] if host is blank/contains spaces or numeric
    /// options are zero.
    pub fn new(
        host: impl Into<String>,
        gpu_layers: Option<u32>,
        context_size: Option<u32>,
        jinja: bool,
    ) -> Result<Self, DomainError> {
        let options = Self {
            host: host.into(),
            gpu_layers,
            context_size,
            jinja,
        };
        options.validate()?;
        Ok(options)
    }

    /// Validates option invariants.
    ///
    /// # Errors
    /// Returns [`DomainError`] on invalid options.
    pub fn validate(&self) -> Result<(), DomainError> {
        crate::validation::non_empty(&self.host, "modelServer.options.host")?;
        if self.host.chars().any(char::is_whitespace) {
            return Err(DomainError::Invariant(
                "modelServer.options.host must not contain spaces".to_owned(),
            ));
        }
        if self.gpu_layers == Some(0) {
            return Err(DomainError::Invariant(
                "modelServer.options.gpuLayers must be > 0".to_owned(),
            ));
        }
        if self.context_size == Some(0) {
            return Err(DomainError::Invariant(
                "modelServer.options.contextSize must be > 0".to_owned(),
            ));
        }
        Ok(())
    }
}

/// Validates free-form runtime arguments do not duplicate structured options.
///
/// # Errors
/// Returns [`DomainError`] if an argument is a reserved exact flag or
/// `--flag=value` form.
pub fn validate_free_args(args: &[String]) -> Result<(), DomainError> {
    const RESERVED: &[&str] = &[
        "--model",
        "-m",
        "-hf",
        "--port",
        "--host",
        "-ngl",
        "--gpu-layers",
        "-c",
        "--ctx-size",
        "--jinja",
    ];
    for arg in args {
        if RESERVED.contains(&arg.as_str()) {
            return Err(DomainError::Invariant(format!(
                "modelServer.args contains reserved flag {arg}"
            )));
        }
        if let Some((flag, _)) = arg.split_once('=') {
            if RESERVED.contains(&flag) {
                return Err(DomainError::Invariant(format!(
                    "modelServer.args contains reserved flag {flag}"
                )));
            }
        }
    }
    Ok(())
}

fn validate_warmup_deadline_secs(value: Option<u64>) -> Result<(), DomainError> {
    let Some(value) = value else {
        return Ok(());
    };
    if !(MIN_WARMUP_DEADLINE_SECS..=MAX_WARMUP_DEADLINE_SECS).contains(&value) {
        return Err(DomainError::Invariant(format!(
            "modelServer.warmupDeadlineSecs must be between {MIN_WARMUP_DEADLINE_SECS} and {MAX_WARMUP_DEADLINE_SECS}"
        )));
    }
    Ok(())
}

fn deserialize_warmup_deadline_secs<'de, D>(deserializer: D) -> Result<Option<u64>, D::Error>
where
    D: Deserializer<'de>,
{
    let value = Option::<u64>::deserialize(deserializer)?;
    validate_warmup_deadline_secs(value).map_err(serde::de::Error::custom)?;
    Ok(value)
}

/// Stop policy for a managed local model server.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub enum StopPolicy {
    /// Leave the process running.
    KeepAlive,
    /// Stop when no profile uses it. Not used aggressively in V1.
    StopWhenUnused,
    /// Stop on application shutdown.
    StopOnAppExit,
}

/// Persisted configuration for a local model server.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct LocalModelServerConfig {
    /// Stable identity.
    pub id: LocalModelServerId,
    /// Server implementation family.
    pub kind: LocalModelServerKind,
    /// Display name.
    pub name: String,
    /// OpenAI-compatible endpoint.
    pub endpoint: ModelServerEndpoint,
    /// Served model.
    pub model: LocalModelRef,
    /// Optional executable path or bare command name.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub binary: Option<ExecutablePath>,
    /// Structured llama.cpp options.
    pub options: LlamaCppOptions,
    /// Extra argv passed to the runtime.
    #[serde(default)]
    pub args: Vec<String>,
    /// Whether IdeA may start the server lazily.
    pub auto_start: bool,
    /// Process stop policy.
    #[serde(default = "default_stop_policy")]
    pub stop_policy: StopPolicy,
    /// Optional readiness warmup deadline override in seconds.
    ///
    /// Absent means the application-level default applies.
    #[serde(
        default,
        deserialize_with = "deserialize_warmup_deadline_secs",
        skip_serializing_if = "Option::is_none"
    )]
    pub warmup_deadline_secs: Option<u64>,
}

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

impl LocalModelServerConfig {
    /// Builds a validated local server configuration.
    ///
    /// # Errors
    /// Returns [`DomainError`] if names are blank or auto-start invariants are not
    /// satisfied by static data.
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        id: LocalModelServerId,
        kind: LocalModelServerKind,
        name: impl Into<String>,
        endpoint: ModelServerEndpoint,
        model: LocalModelRef,
        binary: Option<ExecutablePath>,
        options: LlamaCppOptions,
        args: Vec<String>,
        auto_start: bool,
        stop_policy: StopPolicy,
    ) -> Result<Self, DomainError> {
        let name = name.into();
        crate::validation::non_empty(&name, "modelServer.name")?;
        options.validate()?;
        validate_free_args(&args)?;
        if auto_start {
            if kind != LocalModelServerKind::LlamaCpp {
                return Err(DomainError::Invariant(
                    "modelServer.autoStart requires kind=LlamaCpp".to_owned(),
                ));
            }
            if model.source.is_none() {
                return Err(DomainError::Invariant(
                    "modelServer.autoStart requires model.source".to_owned(),
                ));
            }
        }
        Ok(Self {
            id,
            kind,
            name,
            endpoint,
            model,
            binary,
            options,
            args,
            auto_start,
            stop_policy,
            warmup_deadline_secs: None,
        })
    }

    /// Returns a copy with a validated readiness warmup deadline override.
    ///
    /// # Errors
    /// Returns [`DomainError`] if the value is outside the supported
    /// `[30, 1800]` seconds range.
    pub fn with_warmup_deadline_secs(
        mut self,
        warmup_deadline_secs: Option<u64>,
    ) -> Result<Self, DomainError> {
        validate_warmup_deadline_secs(warmup_deadline_secs)?;
        self.warmup_deadline_secs = warmup_deadline_secs;
        Ok(self)
    }
}

/// Observed model-server status.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ModelServerStatus {
    /// The server is reachable and was already running.
    ReadyReused,
    /// The server is reachable after IdeA started it.
    ReadyStarted,
    /// The server is not reachable.
    Unreachable,
}

/// Successful ensure result used by agent launch.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ModelServerReady {
    /// Base URL to use in OpenCode config.
    pub base_url: String,
    /// Served model name to use in OpenCode config.
    pub model: String,
    /// Whether the server was reused or started.
    pub status: ModelServerStatus,
}

/// Lifecycle status emitted for presentation.
#[derive(Debug, Clone, PartialEq)]
pub enum ModelServerLifecycleStatus {
    /// No local model server is configured for the profile.
    NotConfigured,
    /// A readiness probe is in progress.
    Probing,
    /// IdeA is starting the server.
    Starting,
    /// The server process is alive but still downloading/preparing a remote model.
    Downloading {
        /// Downloaded bytes when known.
        downloaded_bytes: Option<u64>,
        /// Total bytes when known.
        total_bytes: Option<u64>,
        /// Completion percentage when known.
        percent: Option<f32>,
        /// Remote model source, for example a Hugging Face repository id.
        source: Option<String>,
    },
    /// The server is ready.
    Ready {
        /// `true` when an already-running server was reused.
        reused: bool,
    },
    /// The server failed to become ready.
    Failed {
        /// Stable error code.
        code: String,
        /// Human-readable message.
        message: String,
    },
}

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

    #[test]
    fn endpoint_normalises_to_v1() {
        let endpoint = ModelServerEndpoint::new("http://localhost:8080", 8080).unwrap();
        assert_eq!(endpoint.base_url, "http://localhost:8080/v1");
        let endpoint = ModelServerEndpoint::new("http://localhost:8080/v1/", 8080).unwrap();
        assert_eq!(endpoint.base_url, "http://localhost:8080/v1");
    }

    #[test]
    fn endpoint_rejects_remote_host() {
        let err = ModelServerEndpoint::new("http://example.com:8080", 8080).unwrap_err();
        assert!(err.to_string().contains("localhost"));
        let err = ModelServerEndpoint::new("http://127.0.0.1.evil.com:8080", 8080).unwrap_err();
        assert!(err.to_string().contains("localhost"));
    }

    #[test]
    fn auto_start_requires_model_source() {
        let endpoint = ModelServerEndpoint::new("http://localhost:8080/v1", 8080).unwrap();
        let model = LocalModelRef::new("qwen", "Qwen", None, "qwen").unwrap();
        let err = LocalModelServerConfig::new(
            LocalModelServerId::from_uuid(uuid::Uuid::nil()),
            LocalModelServerKind::LlamaCpp,
            "llama.cpp",
            endpoint,
            model,
            None,
            LlamaCppOptions::default(),
            Vec::new(),
            true,
            StopPolicy::StopOnAppExit,
        )
        .unwrap_err();
        assert!(err.to_string().contains("model.source"));
    }

    #[test]
    fn hf_model_ref_accepts_repo_and_optional_quant() {
        assert_eq!(
            HfModelRef::new("Qwen/Qwen3-Coder:Q4_K_M").unwrap().as_str(),
            "Qwen/Qwen3-Coder:Q4_K_M"
        );
        assert!(HfModelRef::new("Qwen/Qwen3 Coder").is_err());
        assert!(HfModelRef::new("https://huggingface.co/Qwen/Qwen3").is_err());
        assert!(HfModelRef::new("Qwen").is_err());
        assert!(HfModelRef::new("a/b:c:d").is_err());
    }

    #[test]
    fn warmup_deadline_secs_is_optional_and_validated() {
        let endpoint = ModelServerEndpoint::new("http://localhost:8080/v1", 8080).unwrap();
        let model = LocalModelRef::new(
            "qwen",
            "Qwen",
            Some(ModelSource::LocalPath {
                path: ModelPath::new("/models/qwen.gguf").unwrap(),
            }),
            "qwen",
        )
        .unwrap();
        let config = LocalModelServerConfig::new(
            LocalModelServerId::from_uuid(uuid::Uuid::nil()),
            LocalModelServerKind::LlamaCpp,
            "llama.cpp",
            endpoint,
            model,
            None,
            LlamaCppOptions::default(),
            Vec::new(),
            true,
            StopPolicy::StopOnAppExit,
        )
        .unwrap();

        assert_eq!(config.warmup_deadline_secs, None);
        assert_eq!(
            config
                .clone()
                .with_warmup_deadline_secs(Some(MIN_WARMUP_DEADLINE_SECS))
                .unwrap()
                .warmup_deadline_secs,
            Some(MIN_WARMUP_DEADLINE_SECS)
        );
        assert!(config
            .clone()
            .with_warmup_deadline_secs(Some(MIN_WARMUP_DEADLINE_SECS - 1))
            .is_err());
        assert!(config
            .with_warmup_deadline_secs(Some(MAX_WARMUP_DEADLINE_SECS + 1))
            .is_err());
    }

    #[test]
    fn validate_free_args_rejects_reserved_flags_token_and_equals() {
        for flag in [
            "--model",
            "-m",
            "-hf",
            "--port",
            "--host",
            "-ngl",
            "--gpu-layers",
            "-c",
            "--ctx-size",
            "--jinja",
        ] {
            assert!(validate_free_args(&[flag.to_owned()]).is_err(), "{flag}");
            assert!(
                validate_free_args(&[format!("{flag}=value")]).is_err(),
                "{flag}=value"
            );
        }
        validate_free_args(&["--threads".to_owned(), "8".to_owned()]).unwrap();
    }
}