//! Model providers (implementations in codecraft-core) and the registry of //! configured providers with their concurrency limits. pub use codecraft_core::provider::*; use crate::config::{Config, ModelConfig}; use anyhow::{Result, anyhow}; use serde_json::Value; use std::collections::BTreeMap; use std::sync::Arc; pub struct ProviderEntry { pub name: String, pub provider: Arc, pub models: Vec, } #[derive(Default)] pub struct Registry { pub providers: BTreeMap, /// Concurrency limits keyed by "provider" and "provider/model". limits: BTreeMap>, } /// Held while a request runs; dropping it frees the slot(s). pub struct Slot { _permits: Vec, } impl Registry { /// Build providers from config. Providers with an empty model list are queried /// for their models; failures are logged and leave the list empty. pub async fn build(cfg: &Config) -> Self { let client = http_client(); let mut providers = BTreeMap::new(); let mut limits = BTreeMap::new(); for (name, pc) in &cfg.providers { if !pc.enabled { continue; } let provider: Arc = build(client.clone(), pc); let mut models: Vec = pc.models.iter().map(|m| m.resolved()).collect(); if let Some(n) = pc.max_concurrency.filter(|n| *n > 0) { limits.insert(name.clone(), Arc::new(tokio::sync::Semaphore::new(n))); } for m in &models { if let Some(n) = m.max_concurrency.filter(|n| *n > 0) { limits.insert(format!("{name}/{}", m.id), Arc::new(tokio::sync::Semaphore::new(n))); } } if models.is_empty() { match tokio::time::timeout(std::time::Duration::from_secs(15), provider.list_models()).await { Ok(Ok(m)) => models = m, Ok(Err(e)) => tracing::warn!("provider {name}: model discovery failed: {e:#}"), Err(_) => tracing::warn!("provider {name}: model discovery timed out"), } } providers.insert( name.clone(), ProviderEntry { name: name.clone(), provider, models }, ); } // Claude Code (the CLI's own agent) appears as a model list when it is installed. if !providers.contains_key(crate::claude_code::PROVIDER) && crate::claude_code::binary().is_some() { providers.insert( crate::claude_code::PROVIDER.to_string(), ProviderEntry { name: crate::claude_code::PROVIDER.to_string(), provider: Arc::new(crate::claude_code::Marker), models: crate::claude_code::models() }, ); } Self { providers, limits } } /// First model of any provider but `skip`. pub fn first_model_excluding(&self, skip: &str) -> Option { self.providers.values().filter(|e| e.name != skip).find_map(|e| e.models.first().map(|m| format!("{}/{}", e.name, m.id))) } /// Wait for a free slot on the model and its provider. Always acquired in the /// same order (model, then provider), so concurrent requests cannot deadlock. pub async fn acquire(&self, spec: &str) -> Slot { let provider = spec.split_once('/').map(|(p, _)| p).unwrap_or(spec); let mut permits = Vec::new(); for key in [spec, provider] { if let Some(sem) = self.limits.get(key) { if let Ok(p) = sem.clone().acquire_owned().await { permits.push(p); } } } Slot { _permits: permits } } /// True if a request to this model would have to wait. pub fn is_saturated(&self, spec: &str) -> bool { let provider = spec.split_once('/').map(|(p, _)| p).unwrap_or(spec); [spec, provider].iter().any(|k| self.limits.get(*k).is_some_and(|s| s.available_permits() == 0)) } /// Resolve "provider/model". Unknown models of a known provider are allowed /// with default settings so users can type any model id. pub fn resolve(&self, spec: &str) -> Result<(Arc, ModelConfig)> { let (pname, mid) = spec .split_once('/') .ok_or_else(|| anyhow!("model must be 'provider/model', got '{spec}'"))?; let entry = self .providers .get(pname) .ok_or_else(|| anyhow!("unknown or disabled provider '{pname}'"))?; let model = entry .models .iter() .find(|m| m.id == mid) .cloned() .unwrap_or_else(|| ModelConfig { id: mid.to_string(), ..Default::default() }); Ok((entry.provider.clone(), model)) } pub fn first_model(&self) -> Option { self.providers .values() .find_map(|e| e.models.first().map(|m| format!("{}/{}", e.name, m.id))) } pub fn models_json(&self) -> Value { let list: Vec = self .providers .values() .map(|e| { serde_json::json!({ "provider": e.name, "models": e.models.iter().map(|m| serde_json::json!({ "id": m.id, "name": m.name.clone().unwrap_or_else(|| m.id.clone()), "context": m.context, "max_output": m.max_output, "reasoning": m.reasoning, "vision": m.vision, "effort": m.effort, "input_price": m.input_price, "output_price": m.output_price, "cache_read_price": m.cache_read_price, "cache_write_price": m.cache_write_price, "max_concurrency": m.max_concurrency, })).collect::>() }) }) .collect(); Value::Array(list) } }