bamboo-llm 2026.10.21

//! Aggregates model lists from all configured providers into a unified [`ProviderCatalog`].

use std::sync::Arc;
use std::time::Duration;

use bamboo_domain::provider_catalog::{
    ModelCapabilities, ModelSource, ProviderCatalog, ProviderDescriptor, ProviderModelDescriptor,
};
use bamboo_domain::ProviderModelRef;

use crate::provider_registry::ProviderRegistry;

/// Service that aggregates model metadata across all registered providers.
pub struct ModelCatalogService {
    registry: Arc<ProviderRegistry>,
}

impl ModelCatalogService {
    pub fn new(registry: Arc<ProviderRegistry>) -> Self {
        Self { registry }
    }

    /// Runtime candidates come only from admission configuration. Upstream
    /// metadata enriches these IDs within a bounded best-effort lookup; a
    /// failed lookup never removes an admitted custom model.
    pub async fn get_catalog(&self) -> ProviderCatalog {
        let mut providers = Vec::new();
        let mut models = Vec::new();
        let mut metadata = self.registry.provider_metadata();
        metadata.sort_by(|left, right| left.id.cmp(&right.id));
        let lookups = metadata.into_iter().map(|meta| {
            let (provider, admitted) = self.registry.runtime_provider_snapshot(&meta.id);
            async move {
                let info = match provider.as_ref().filter(|_| !admitted.is_empty()) {
                    Some(provider) => {
                        tokio::time::timeout(Duration::from_secs(3), provider.list_model_info())
                            .await
                            .ok()
                            .and_then(Result::ok)
                            .unwrap_or_default()
                    }
                    None => Vec::new(),
                };
                (meta, provider, admitted, info)
            }
        });
        for (meta, provider, admitted, info) in futures::future::join_all(lookups).await {
            let authenticated = provider.is_some();
            providers.push(ProviderDescriptor {
                id: meta.id.clone(),
                display_name: meta.display_name.clone(),
                enabled: true,
                authenticated,
            });

            if authenticated {
                let info: std::collections::HashMap<_, _> = info
                    .into_iter()
                    .map(|info| (info.id.clone(), info))
                    .collect();
                for id in admitted {
                    let upstream = info.get(&id);
                    models.push(ProviderModelDescriptor {
                        reference: ProviderModelRef::new(&meta.id, &id),
                        display_name: id.clone(),
                        provider_display_name: meta.display_name.clone(),
                        capabilities: ModelCapabilities {
                            supports_vision: provider.as_ref().unwrap().supports_vision(&id).await,
                            max_context_tokens: upstream.and_then(|info| info.max_context_tokens),
                            max_output_tokens: upstream.and_then(|info| info.max_output_tokens),
                            ..ModelCapabilities::default()
                        },
                        source: Some(if upstream.is_some() {
                            ModelSource::Upstream
                        } else {
                            ModelSource::Manual
                        }),
                        discovered_at: None,
                    });
                }
            }
        }

        ProviderCatalog {
            providers,
            models,
            updated_at: Some(chrono::Utc::now().to_rfc3339()),
        }
    }

    /// Discover upstream models for selection in settings. Discovery never
    /// mutates the admitted runtime list.
    pub async fn list_models_for_provider(
        &self,
        provider_name: &str,
    ) -> Result<Vec<ProviderModelDescriptor>, String> {
        let provider = self
            .registry
            .get(provider_name)
            .ok_or_else(|| format!("Provider '{}' not found", provider_name))?;
        let provider_display_name = self
            .registry
            .get_metadata(provider_name)
            .map(|meta| meta.display_name)
            .unwrap_or_else(|| display_name_for_provider(provider_name));

        let info_list = provider
            .list_model_info()
            .await
            .map_err(|e| e.to_string())?;

        let mut models = Vec::new();
        for info in info_list {
            models.push(ProviderModelDescriptor {
                reference: ProviderModelRef::new(provider_name, &info.id),
                display_name: info.id.clone(),
                provider_display_name: provider_display_name.clone(),
                capabilities: ModelCapabilities {
                    supports_vision: provider.supports_vision(&info.id).await,
                    max_context_tokens: info.max_context_tokens,
                    max_output_tokens: info.max_output_tokens,
                    ..ModelCapabilities::default()
                },
                source: Some(ModelSource::Upstream),
                discovered_at: None,
            });
        }
        Ok(models)
    }

    /// Fetch model lists from all registered providers.
    ///
    /// Returns a result per provider, preserving individual success/failure status
    /// so callers can see partial results.
    pub async fn fetch_models_for_all_providers(&self) -> Vec<ProviderFetchResult> {
        let names = self.registry.provider_names();
        let mut results = Vec::with_capacity(names.len());

        for name in names {
            let result = self.list_models_for_provider(&name).await;
            results.push(match result {
                Ok(models) => ProviderFetchResult {
                    provider: name,
                    models: Some(models),
                    error: None,
                },
                Err(e) => ProviderFetchResult {
                    provider: name,
                    models: None,
                    error: Some(e),
                },
            });
        }

        results
    }
}

/// Result of fetching models from a single provider during a bulk fetch.
#[derive(Debug, serde::Serialize)]
pub struct ProviderFetchResult {
    pub provider: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub models: Option<Vec<ProviderModelDescriptor>>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub error: Option<String>,
}

fn display_name_for_provider(id: &str) -> String {
    match id {
        "openai" => "OpenAI".to_string(),
        "anthropic" => "Anthropic".to_string(),
        "gemini" => "Gemini".to_string(),
        "copilot" => "GitHub Copilot".to_string(),
        other => other.to_string(),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::provider::{LLMError, LLMProvider, LLMStream, ProviderModelInfo, Result};
    use async_trait::async_trait;
    use bamboo_domain::Message;
    use bamboo_domain::ToolSchema;
    use futures::stream;
    use std::collections::HashMap;
    use std::sync::Arc;

    struct MockProvider {
        models: Vec<ProviderModelInfo>,
    }

    #[async_trait]
    impl LLMProvider for MockProvider {
        async fn chat_stream(
            &self,
            _messages: &[Message],
            _tools: &[ToolSchema],
            _max_output_tokens: Option<u32>,
            _model: &str,
        ) -> Result<LLMStream> {
            Ok(Box::pin(stream::empty()))
        }

        async fn list_model_info(&self) -> Result<Vec<ProviderModelInfo>> {
            Ok(self.models.clone())
        }

        async fn list_models(&self) -> Result<Vec<String>> {
            Ok(self.models.iter().map(|m| m.id.clone()).collect())
        }
    }

    struct FailingProvider;

    #[async_trait]
    impl LLMProvider for FailingProvider {
        async fn chat_stream(
            &self,
            _messages: &[Message],
            _tools: &[ToolSchema],
            _max_output_tokens: Option<u32>,
            _model: &str,
        ) -> Result<LLMStream> {
            Err(LLMError::Api("fail".to_string()))
        }

        async fn list_model_info(&self) -> Result<Vec<ProviderModelInfo>> {
            Err(LLMError::Api("list failed".to_string()))
        }
    }

    fn make_registry(
        providers: Vec<(&str, Vec<ProviderModelInfo>)>,
        default: &str,
    ) -> Arc<ProviderRegistry> {
        let mut map = HashMap::new();
        let mut admitted = HashMap::new();
        for (name, models) in providers {
            admitted.insert(
                name.to_string(),
                models.iter().map(|model| model.id.clone()).collect(),
            );
            map.insert(
                name.to_string(),
                Arc::new(MockProvider { models }) as Arc<dyn LLMProvider>,
            );
        }
        let registry = Arc::new(ProviderRegistry::new(map, default.to_string()));
        registry.set_runtime_models(admitted);
        registry
    }

    // ---- display_name_for_provider ----

    #[test]
    fn display_name_openai() {
        assert_eq!(display_name_for_provider("openai"), "OpenAI");
    }

    #[test]
    fn display_name_anthropic() {
        assert_eq!(display_name_for_provider("anthropic"), "Anthropic");
    }

    #[test]
    fn display_name_gemini() {
        assert_eq!(display_name_for_provider("gemini"), "Gemini");
    }

    #[test]
    fn display_name_copilot() {
        assert_eq!(display_name_for_provider("copilot"), "GitHub Copilot");
    }

    #[test]
    fn display_name_unknown_passthrough() {
        assert_eq!(display_name_for_provider("custom-llm"), "custom-llm");
    }

    // ---- get_catalog ----

    #[tokio::test]
    async fn catalog_includes_all_providers() {
        let registry = make_registry(
            vec![
                ("openai", vec![ProviderModelInfo::from_id("gpt-4o")]),
                ("anthropic", vec![ProviderModelInfo::from_id("claude-3")]),
            ],
            "openai",
        );
        let service = ModelCatalogService::new(registry);
        let catalog = service.get_catalog().await;

        assert_eq!(catalog.providers.len(), 2);
        let provider_ids: Vec<&str> = catalog.providers.iter().map(|p| p.id.as_str()).collect();
        assert!(provider_ids.contains(&"openai"));
        assert!(provider_ids.contains(&"anthropic"));
    }

    #[tokio::test]
    async fn catalog_aggregates_models_from_all_providers() {
        let registry = make_registry(
            vec![
                (
                    "openai",
                    vec![
                        ProviderModelInfo::from_id("gpt-4o"),
                        ProviderModelInfo::from_id("gpt-4-turbo"),
                    ],
                ),
                ("anthropic", vec![ProviderModelInfo::from_id("claude-3")]),
            ],
            "openai",
        );
        let service = ModelCatalogService::new(registry);
        let catalog = service.get_catalog().await;

        assert_eq!(catalog.models.len(), 3);
    }

    #[tokio::test]
    async fn catalog_model_refs_include_provider() {
        let registry = make_registry(
            vec![("openai", vec![ProviderModelInfo::from_id("gpt-4o")])],
            "openai",
        );
        let service = ModelCatalogService::new(registry);
        let catalog = service.get_catalog().await;

        let model = &catalog.models[0];
        assert_eq!(model.reference.provider, "openai");
        assert_eq!(model.reference.model, "gpt-4o");
        assert_eq!(model.provider_display_name, "OpenAI");
    }

    #[tokio::test]
    async fn catalog_preserves_dynamic_model_token_limits() {
        let registry = make_registry(
            vec![(
                "copilot",
                vec![ProviderModelInfo {
                    id: "dynamic-model".to_string(),
                    max_context_tokens: Some(264_000),
                    max_output_tokens: Some(64_000),
                }],
            )],
            "copilot",
        );
        let service = ModelCatalogService::new(registry);
        let models = service.get_catalog().await.models;

        assert_eq!(models.len(), 1);
        assert_eq!(models[0].capabilities.max_context_tokens, Some(264_000));
        assert_eq!(models[0].capabilities.max_output_tokens, Some(64_000));
    }

    #[tokio::test]
    async fn catalog_providers_marked_authenticated() {
        let registry = make_registry(
            vec![("openai", vec![ProviderModelInfo::from_id("gpt-4o")])],
            "openai",
        );
        let service = ModelCatalogService::new(registry);
        let catalog = service.get_catalog().await;

        assert!(catalog.providers[0].authenticated);
    }

    #[tokio::test]
    async fn catalog_empty_registry() {
        let registry = make_registry(vec![], "openai");
        let service = ModelCatalogService::new(registry);
        let catalog = service.get_catalog().await;

        assert!(catalog.providers.is_empty());
        assert!(catalog.models.is_empty());
        assert!(catalog.updated_at.is_some());
    }

    #[tokio::test]
    async fn catalog_has_updated_at_timestamp() {
        let registry = make_registry(vec![], "openai");
        let service = ModelCatalogService::new(registry);
        let catalog = service.get_catalog().await;

        assert!(catalog.updated_at.is_some());
        assert!(!catalog.updated_at.unwrap().is_empty());
    }

    #[tokio::test]
    async fn catalog_includes_admitted_custom_model_even_when_discovery_fails() {
        let mut map = HashMap::new();
        map.insert(
            "openai".to_string(),
            Arc::new(MockProvider {
                models: vec![ProviderModelInfo::from_id("gpt-4o")],
            }) as Arc<dyn LLMProvider>,
        );
        map.insert(
            "broken".to_string(),
            Arc::new(FailingProvider) as Arc<dyn LLMProvider>,
        );
        let registry = Arc::new(ProviderRegistry::new(map, "openai".to_string()));
        registry.set_runtime_models(HashMap::from([
            ("openai".into(), vec!["gpt-4o".into()]),
            ("broken".into(), vec!["custom/id".into()]),
        ]));
        let service = ModelCatalogService::new(registry);
        let catalog = service.get_catalog().await;

        assert_eq!(catalog.providers.len(), 2);
        assert_eq!(catalog.models.len(), 2);
        assert!(catalog
            .models
            .iter()
            .any(|model| model.reference.model == "custom/id"));
    }

    #[tokio::test]
    async fn discovery_never_admits_its_relay_catalog_to_runtime() {
        let registry = make_registry(
            vec![(
                "relay",
                (0..1000)
                    .map(|n| ProviderModelInfo::from_id(format!("relay-{n}")))
                    .collect(),
            )],
            "relay",
        );
        registry.set_runtime_models(HashMap::from([("relay".into(), vec!["custom/id".into()])]));
        let service = ModelCatalogService::new(registry);
        assert_eq!(
            service
                .list_models_for_provider("relay")
                .await
                .unwrap()
                .len(),
            1000
        );
        let catalog = service.get_catalog().await;
        assert_eq!(catalog.models.len(), 1);
        assert_eq!(catalog.models[0].reference.model, "custom/id");
    }

    // ---- list_models_for_provider ----

    #[tokio::test]
    async fn list_models_for_single_provider() {
        let registry = make_registry(
            vec![
                ("openai", vec![ProviderModelInfo::from_id("gpt-4o")]),
                ("anthropic", vec![ProviderModelInfo::from_id("claude-3")]),
            ],
            "openai",
        );
        let service = ModelCatalogService::new(registry);
        let models = service.list_models_for_provider("openai").await.unwrap();

        assert_eq!(models.len(), 1);
        assert_eq!(models[0].reference.provider, "openai");
        assert_eq!(models[0].reference.model, "gpt-4o");
    }

    #[tokio::test]
    async fn list_models_for_unknown_provider_returns_error() {
        let registry = make_registry(vec![], "openai");
        let service = ModelCatalogService::new(registry);
        let result = service.list_models_for_provider("nonexistent").await;

        assert!(result.is_err());
        assert!(result.unwrap_err().contains("not found"));
    }

    // ---- fetch_models_for_all_providers ----

    #[tokio::test]
    async fn fetch_all_returns_result_for_each_provider() {
        let registry = make_registry(
            vec![
                ("openai", vec![ProviderModelInfo::from_id("gpt-4o")]),
                ("anthropic", vec![ProviderModelInfo::from_id("claude-3")]),
            ],
            "openai",
        );
        let service = ModelCatalogService::new(registry);
        let results = service.fetch_models_for_all_providers().await;

        assert_eq!(results.len(), 2);
        assert!(results
            .iter()
            .all(|r| r.models.is_some() && r.error.is_none()));
    }

    #[tokio::test]
    async fn fetch_all_preserves_individual_failures() {
        let mut map = HashMap::new();
        map.insert(
            "openai".to_string(),
            Arc::new(MockProvider {
                models: vec![ProviderModelInfo::from_id("gpt-4o")],
            }) as Arc<dyn LLMProvider>,
        );
        map.insert(
            "broken".to_string(),
            Arc::new(FailingProvider) as Arc<dyn LLMProvider>,
        );
        let registry = Arc::new(ProviderRegistry::new(map, "openai".to_string()));
        let service = ModelCatalogService::new(registry);
        let results = service.fetch_models_for_all_providers().await;

        assert_eq!(results.len(), 2);
        let ok = results.iter().find(|r| r.provider == "openai").unwrap();
        assert!(ok.models.is_some());
        assert!(ok.error.is_none());

        let fail = results.iter().find(|r| r.provider == "broken").unwrap();
        assert!(fail.models.is_none());
        assert!(fail.error.is_some());
    }

    #[tokio::test]
    async fn fetch_all_empty_registry() {
        let registry = make_registry(vec![], "openai");
        let service = ModelCatalogService::new(registry);
        let results = service.fetch_models_for_all_providers().await;

        assert!(results.is_empty());
    }
}