IdeA/crates/application/src/model_server.rs

//! Use cases for local model servers.

use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::time::Duration;

use domain::events::DomainEvent;
use domain::model_server::{
    LocalModelServerConfig, ModelServerLifecycleStatus, ModelServerReady, ModelServerStatus,
    ModelSource,
};
use domain::ports::{
    EventBus, FileSystem, ManagedProcess, ManagedProcessHandle, ModelServerError, ModelServerProbe,
    ModelServerRegistry, ModelServerRuntime, ProcessStatus, ProfileStore, RemotePath,
};
use domain::{LocalModelServerId, StopPolicy};

use crate::error::AppError;

/// Output of [`ListModelServers::execute`].
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ListModelServersOutput {
    /// Persisted local model-server configs.
    pub servers: Vec<LocalModelServerConfig>,
}

/// Lists local model-server configurations.
pub struct ListModelServers {
    registry: Arc<dyn ModelServerRegistry>,
}

impl ListModelServers {
    /// Builds the use case.
    #[must_use]
    pub fn new(registry: Arc<dyn ModelServerRegistry>) -> Self {
        Self { registry }
    }

    /// Lists configs.
    ///
    /// # Errors
    /// [`AppError::ModelServer`] on registry failure.
    pub async fn execute(&self) -> Result<ListModelServersOutput, AppError> {
        Ok(ListModelServersOutput {
            servers: self.registry.list().await?,
        })
    }
}

/// Input for [`SaveModelServer::execute`].
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SaveModelServerInput {
    /// Config to upsert by id.
    pub config: LocalModelServerConfig,
}

/// Output of [`SaveModelServer::execute`].
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SaveModelServerOutput {
    /// Saved config.
    pub config: LocalModelServerConfig,
}

/// Saves a local model-server configuration.
pub struct SaveModelServer {
    registry: Arc<dyn ModelServerRegistry>,
}

impl SaveModelServer {
    /// Builds the use case.
    #[must_use]
    pub fn new(registry: Arc<dyn ModelServerRegistry>) -> Self {
        Self { registry }
    }

    /// Saves a config.
    ///
    /// # Errors
    /// [`AppError::ModelServer`] on registry failure.
    pub async fn execute(
        &self,
        input: SaveModelServerInput,
    ) -> Result<SaveModelServerOutput, AppError> {
        self.registry.save(input.config.clone()).await?;
        Ok(SaveModelServerOutput {
            config: input.config,
        })
    }
}

/// Input for [`DeleteModelServer::execute`].
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DeleteModelServerInput {
    /// Config id to delete.
    pub server_id: LocalModelServerId,
}

/// Deletes a local model-server config when no OpenCode profile still references it.
pub struct DeleteModelServer {
    registry: Arc<dyn ModelServerRegistry>,
    profiles: Arc<dyn ProfileStore>,
}

impl DeleteModelServer {
    /// Builds the use case.
    #[must_use]
    pub fn new(registry: Arc<dyn ModelServerRegistry>, profiles: Arc<dyn ProfileStore>) -> Self {
        Self { registry, profiles }
    }

    /// Deletes a config after checking profile references.
    ///
    /// # Errors
    /// [`AppError::ModelServer`] with `code=model_server_in_use` when referenced.
    pub async fn execute(&self, input: DeleteModelServerInput) -> Result<(), AppError> {
        let profiles = self.profiles.list().await?;
        if profiles.iter().any(|profile| {
            profile
                .opencode
                .as_ref()
                .and_then(|config| config.local_model_server_id)
                == Some(input.server_id)
        }) {
            return Err(ModelServerError::InUse(input.server_id.to_string()).into());
        }
        self.registry.delete(input.server_id).await?;
        Ok(())
    }
}

/// Input for [`EnsureLocalModelServer::execute`].
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct EnsureLocalModelServerInput {
    /// Referenced local model server.
    pub server_id: LocalModelServerId,
}

/// Output of [`EnsureLocalModelServer::execute`].
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct EnsureLocalModelServerOutput {
    /// Ready server data to inject into OpenCode config.
    pub ready: ModelServerReady,
}

/// Readiness retry policy.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ReadinessPolicy {
    /// Number of probes after spawning.
    pub attempts: usize,
    /// Delay between attempts.
    pub backoff: Duration,
}

impl Default for ReadinessPolicy {
    fn default() -> Self {
        Self {
            attempts: 20,
            backoff: Duration::from_millis(250),
        }
    }
}

#[derive(Debug, Clone)]
struct ActiveServer {
    handle: ManagedProcessHandle,
    port: u16,
    stop_policy: StopPolicy,
}

/// Ensures a configured local model server is reachable, starting it when allowed.
pub struct EnsureLocalModelServer {
    registry: Arc<dyn ModelServerRegistry>,
    probe: Arc<dyn ModelServerProbe>,
    process: Arc<dyn ManagedProcess>,
    runtime: Arc<dyn ModelServerRuntime>,
    fs: Arc<dyn FileSystem>,
    events: Arc<dyn EventBus>,
    active: Mutex<HashMap<LocalModelServerId, ActiveServer>>,
    readiness: ReadinessPolicy,
}

impl EnsureLocalModelServer {
    /// Builds the use case.
    #[allow(clippy::too_many_arguments)]
    #[must_use]
    pub fn new(
        registry: Arc<dyn ModelServerRegistry>,
        probe: Arc<dyn ModelServerProbe>,
        process: Arc<dyn ManagedProcess>,
        runtime: Arc<dyn ModelServerRuntime>,
        fs: Arc<dyn FileSystem>,
        events: Arc<dyn EventBus>,
    ) -> Self {
        Self {
            registry,
            probe,
            process,
            runtime,
            fs,
            events,
            active: Mutex::new(HashMap::new()),
            readiness: ReadinessPolicy::default(),
        }
    }

    /// Overrides readiness policy, mainly for tests.
    #[must_use]
    pub fn with_readiness_policy(mut self, readiness: ReadinessPolicy) -> Self {
        self.readiness = readiness;
        self
    }

    /// Ensures the server is reachable.
    ///
    /// # Errors
    /// [`AppError::ModelServer`] if the server cannot be prepared.
    pub async fn execute(
        &self,
        input: EnsureLocalModelServerInput,
    ) -> Result<EnsureLocalModelServerOutput, AppError> {
        let config = self
            .registry
            .get(&input.server_id)
            .await?
            .ok_or(ModelServerError::NotConfigured)?;

        self.publish(config.id, ModelServerLifecycleStatus::Probing);
        let initial_probe = match self.probe.probe(&config.endpoint).await {
            Ok(status) => status,
            Err(err) => return self.fail(config.id, err),
        };
        match initial_probe {
            ModelServerStatus::ReadyReused | ModelServerStatus::ReadyStarted => {
                self.publish(
                    config.id,
                    ModelServerLifecycleStatus::Ready { reused: true },
                );
                return Ok(EnsureLocalModelServerOutput {
                    ready: ready(&config, ModelServerStatus::ReadyReused),
                });
            }
            ModelServerStatus::Unreachable => {}
        }

        if !config.auto_start {
            let err = ModelServerError::Probe("server unreachable and autoStart=false".to_owned());
            self.publish_failure(config.id, &err);
            return Err(err.into());
        }

        self.ensure_model_path_accessible(&config).await?;
        self.ensure_no_active_port_collision(&config).await?;

        self.publish(config.id, ModelServerLifecycleStatus::Starting);
        let spec = match self.runtime.build_spawn_spec(&config) {
            Ok(spec) => spec,
            Err(err) => return self.fail(config.id, err),
        };
        let handle = match self.process.spawn(spec).await {
            Ok(handle) => handle,
            Err(err) => return self.fail(config.id, err),
        };
        self.active.lock().unwrap().insert(
            config.id,
            ActiveServer {
                handle: handle.clone(),
                port: config.endpoint.port,
                stop_policy: config.stop_policy,
            },
        );

        for attempt in 0..self.readiness.attempts {
            match self.probe.probe(&config.endpoint).await {
                Err(err) => {
                    self.stop_started_server(config.id, &handle).await;
                    return self.fail(config.id, err);
                }
                Ok(ModelServerStatus::ReadyReused | ModelServerStatus::ReadyStarted) => {
                    self.publish(
                        config.id,
                        ModelServerLifecycleStatus::Ready { reused: false },
                    );
                    return Ok(EnsureLocalModelServerOutput {
                        ready: ready(&config, ModelServerStatus::ReadyStarted),
                    });
                }
                Ok(ModelServerStatus::Unreachable) => {
                    if attempt + 1 < self.readiness.attempts && !self.readiness.backoff.is_zero() {
                        tokio::time::sleep(self.readiness.backoff).await;
                    }
                }
            }
        }

        let err = ModelServerError::Timeout;
        self.stop_started_server(config.id, &handle).await;
        self.fail(config.id, err)
    }

    /// Stops active servers whose policy is [`StopPolicy::StopOnAppExit`].
    ///
    /// # Errors
    /// Returns the first process error after attempting every eligible stop.
    pub async fn stop_on_app_exit(&self) -> Result<(), AppError> {
        let entries: Vec<(LocalModelServerId, ActiveServer)> = self
            .active
            .lock()
            .unwrap()
            .iter()
            .map(|(id, active)| (*id, active.clone()))
            .collect();
        let mut first_error: Option<ModelServerError> = None;
        for (id, active) in entries {
            if active.stop_policy != StopPolicy::StopOnAppExit {
                continue;
            }
            if let Err(err) = self.process.kill(&active.handle).await {
                if first_error.is_none() {
                    first_error = Some(err);
                }
            } else {
                self.active.lock().unwrap().remove(&id);
            }
        }
        if let Some(err) = first_error {
            Err(err.into())
        } else {
            Ok(())
        }
    }

    async fn ensure_model_path_accessible(
        &self,
        config: &LocalModelServerConfig,
    ) -> Result<(), AppError> {
        let Some(ModelSource::LocalPath { path }) = config.model.source.as_ref() else {
            return Ok(());
        };
        if path.as_str().is_empty() {
            let err = ModelServerError::PathNotAccessible("model.source.path missing".to_owned());
            self.publish_failure(config.id, &err);
            return Err(err.into());
        }
        match self
            .fs
            .exists(&RemotePath::new(path.as_str().to_owned()))
            .await
        {
            Ok(true) => Ok(()),
            Ok(false) => {
                let err = ModelServerError::PathNotAccessible(path.as_str().to_owned());
                self.publish_failure(config.id, &err);
                Err(err.into())
            }
            Err(domain::ports::FsError::PermissionDenied(p)) => {
                let err = ModelServerError::PermissionDenied(p);
                self.publish_failure(config.id, &err);
                Err(err.into())
            }
            Err(err) => {
                let err = ModelServerError::PathNotAccessible(err.to_string());
                self.publish_failure(config.id, &err);
                Err(err.into())
            }
        }
    }

    async fn ensure_no_active_port_collision(
        &self,
        config: &LocalModelServerConfig,
    ) -> Result<(), AppError> {
        let active: Vec<(LocalModelServerId, ActiveServer)> = self
            .active
            .lock()
            .unwrap()
            .iter()
            .map(|(id, active)| (*id, active.clone()))
            .collect();
        for (id, active) in active {
            if active.port != config.endpoint.port {
                continue;
            }
            match self.process.status(&active.handle).await {
                Ok(ProcessStatus::Running) => {
                    let err = ModelServerError::PortOccupied(config.endpoint.port);
                    self.publish_failure(config.id, &err);
                    return Err(err.into());
                }
                Ok(ProcessStatus::Exited { .. } | ProcessStatus::Unknown) => {
                    self.active.lock().unwrap().remove(&id);
                }
                Err(err) => return self.fail(config.id, err),
            }
        }
        Ok(())
    }

    async fn stop_started_server(
        &self,
        server_id: LocalModelServerId,
        handle: &ManagedProcessHandle,
    ) {
        let _ = self.process.kill(handle).await;
        self.active.lock().unwrap().remove(&server_id);
    }

    fn publish(&self, server_id: LocalModelServerId, status: ModelServerLifecycleStatus) {
        self.events
            .publish(DomainEvent::ModelServerStatusChanged { server_id, status });
    }

    fn fail<T>(&self, server_id: LocalModelServerId, err: ModelServerError) -> Result<T, AppError> {
        self.publish_failure(server_id, &err);
        Err(err.into())
    }

    fn publish_failure(&self, server_id: LocalModelServerId, err: &ModelServerError) {
        self.publish(
            server_id,
            ModelServerLifecycleStatus::Failed {
                code: model_server_error_code(err).to_owned(),
                message: err.to_string(),
            },
        );
    }
}

fn ready(config: &LocalModelServerConfig, status: ModelServerStatus) -> ModelServerReady {
    ModelServerReady {
        base_url: config.endpoint.base_url.clone(),
        model: config.model.served_name.clone(),
        status,
    }
}

/// Stable model-server error code for event/DTO mapping.
#[must_use]
pub fn model_server_error_code(err: &ModelServerError) -> &'static str {
    match err {
        ModelServerError::NotConfigured => "not_configured",
        ModelServerError::Invalid(_) => "invalid",
        ModelServerError::PermissionDenied(_) => "permission_denied",
        ModelServerError::PathNotAccessible(_) => "path_not_accessible",
        ModelServerError::PortOccupied(_) => "port_occupied",
        ModelServerError::InUse(_) => "model_server_in_use",
        ModelServerError::Probe(_) => "probe",
        ModelServerError::Process(_) => "process",
        ModelServerError::Store(_) => "store",
        ModelServerError::Timeout => "timeout",
    }
}