bamboo_server/tools/
model_catalog.rs

1//! Server-side `ModelCatalogPort`: lists models per configured provider so the
2//! parent agent can pin a child session to an explicit model
3//! (`SubAgent` tool `action=list_models` / `create.model`).
4
5use std::sync::Arc;
6use std::time::Duration;
7
8use async_trait::async_trait;
9
10use bamboo_engine::session_app::child_session::{ModelCatalogPort, ProviderModelList};
11use bamboo_llm::ProviderRegistry;
12
13/// How long one provider's `list_models` call may take before we report it as
14/// timed out (the listing is best-effort; a slow provider must not stall the
15/// tool call).
16const PER_PROVIDER_TIMEOUT: Duration = Duration::from_secs(8);
17
18/// Model catalog backed by the live provider registry.
19pub struct RegistryModelCatalog {
20    registry: Arc<ProviderRegistry>,
21}
22
23impl RegistryModelCatalog {
24    pub fn new(registry: Arc<ProviderRegistry>) -> Self {
25        Self { registry }
26    }
27}
28
29#[async_trait]
30impl ModelCatalogPort for RegistryModelCatalog {
31    async fn list_models(&self) -> Vec<ProviderModelList> {
32        let mut names = self.registry.provider_names();
33        names.sort();
34
35        // Query all providers concurrently; tolerate individual failures.
36        let futures = names.into_iter().map(|name| {
37            let provider = self.registry.get(&name);
38            async move {
39                let Some(provider) = provider else {
40                    return ProviderModelList {
41                        provider: name,
42                        models: Vec::new(),
43                        error: Some("provider not initialized".to_string()),
44                    };
45                };
46                match tokio::time::timeout(PER_PROVIDER_TIMEOUT, provider.list_models()).await {
47                    Ok(Ok(mut models)) => {
48                        models.sort();
49                        ProviderModelList {
50                            provider: name,
51                            models,
52                            error: None,
53                        }
54                    }
55                    Ok(Err(e)) => ProviderModelList {
56                        provider: name,
57                        models: Vec::new(),
58                        error: Some(e.to_string()),
59                    },
60                    Err(_) => ProviderModelList {
61                        provider: name,
62                        models: Vec::new(),
63                        error: Some(format!(
64                            "timed out after {}s",
65                            PER_PROVIDER_TIMEOUT.as_secs()
66                        )),
67                    },
68                }
69            }
70        });
71
72        futures::future::join_all(futures).await
73    }
74
75    fn default_provider(&self) -> String {
76        self.registry.default_provider_name()
77    }
78}
bamboo_server/tools/model_catalog.rs

bamboo_server/tools/
model_catalog.rs
