codecraft/server/src/provider/mod.rs

//! 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<dyn Provider>,
    pub models: Vec<ModelConfig>,
}

#[derive(Default)]
pub struct Registry {
    pub providers: BTreeMap<String, ProviderEntry>,
    /// Concurrency limits keyed by "provider" and "provider/model".
    limits: BTreeMap<String, Arc<tokio::sync::Semaphore>>,
}

/// Held while a request runs; dropping it frees the slot(s).
pub struct Slot {
    _permits: Vec<tokio::sync::OwnedSemaphorePermit>,
}


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<dyn Provider> = build(client.clone(), pc);
            let mut models: Vec<ModelConfig> = 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<String> {
        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<dyn Provider>, 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<String> {
        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<Value> = 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::<Vec<_>>()
                })
            })
            .collect();
        Value::Array(list)
    }
}