//! 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 { 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) -> 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) -> 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 { 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, cancel: ModelArtifactCancel, ) -> Result { 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, } #[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 { 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::>() .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::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 { 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 { 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 { 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 { 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>, } 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 { 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 { 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, } #[derive(Debug, Clone, Deserialize)] #[serde(untagged)] enum PersistedModelServersDoc { V2(ModelServersDoc), V1(ModelServersDocV1), } #[derive(Debug, Clone, Deserialize)] #[serde(rename_all = "camelCase")] struct ModelServersDocV1 { servers: Vec, } #[derive(Debug, Clone, Deserialize)] #[serde(rename_all = "camelCase")] struct LocalModelServerConfigV1 { id: LocalModelServerId, kind: LocalModelServerKind, name: String, endpoint: ModelServerEndpoint, model: LocalModelRefV1, #[serde(default)] binary: Option, #[serde(default)] args: Vec, 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, served_name: String, } fn default_stop_policy() -> StopPolicy { StopPolicy::StopOnAppExit } impl From 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, 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, app_data_dir: impl Into) -> 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 { match self.fs.read(&self.path()).await { Ok(bytes) => match serde_json::from_slice::(&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, ModelServerError> { Ok(self .read_doc() .await? .servers .into_iter() .find(|server| &server.id == id)) } async fn list(&self) -> Result, 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 } }