imagegen_bridge_runtime/
registry.rs

1//! Immutable provider registration and discovery.
2
3use std::{collections::BTreeMap, sync::Arc};
4
5use imagegen_bridge_core::{
6    BridgeError, ErrorCode, ImageProvider, ProviderCapabilities, ProviderDescriptor,
7};
8use serde::{Deserialize, Serialize};
9
10const MAX_PROVIDER_NAME_BYTES: usize = 64;
11
12/// Immutable set of providers used by one runtime instance.
13#[derive(Clone)]
14pub struct ProviderRegistry {
15    providers: Arc<BTreeMap<String, Arc<dyn ImageProvider>>>,
16    default: Arc<str>,
17}
18
19impl std::fmt::Debug for ProviderRegistry {
20    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
21        formatter
22            .debug_struct("ProviderRegistry")
23            .field("providers", &self.providers.keys().collect::<Vec<_>>())
24            .field("default", &self.default)
25            .finish()
26    }
27}
28
29impl ProviderRegistry {
30    /// Creates a validated registry and deterministic default provider.
31    pub fn new(
32        providers: impl IntoIterator<Item = Arc<dyn ImageProvider>>,
33        default: impl Into<String>,
34    ) -> Result<Self, BridgeError> {
35        let mut registered = BTreeMap::new();
36        for provider in providers {
37            let descriptor = provider.descriptor();
38            validate_provider_name(&descriptor.name)?;
39            if registered
40                .insert(descriptor.name.clone(), provider)
41                .is_some()
42            {
43                return Err(registry_error(format!(
44                    "provider '{}' is registered more than once",
45                    descriptor.name
46                )));
47            }
48        }
49        if registered.is_empty() {
50            return Err(registry_error("at least one provider must be registered"));
51        }
52        let default = default.into();
53        validate_provider_name(&default)?;
54        if !registered.contains_key(&default) {
55            return Err(registry_error(format!(
56                "default provider '{default}' is not registered"
57            )));
58        }
59        Ok(Self {
60            providers: Arc::new(registered),
61            default: Arc::from(default),
62        })
63    }
64
65    /// Returns the configured default provider name.
66    #[must_use]
67    pub fn default_name(&self) -> &str {
68        &self.default
69    }
70
71    /// Resolves an explicit provider or the deterministic default.
72    pub fn resolve(&self, requested: Option<&str>) -> Result<Arc<dyn ImageProvider>, BridgeError> {
73        let name = requested.unwrap_or(&self.default);
74        validate_provider_name(name)?;
75        self.providers.get(name).cloned().ok_or_else(|| {
76            BridgeError::new(
77                ErrorCode::Configuration,
78                "requested provider is not registered",
79            )
80            .with_detail("provider", name)
81        })
82    }
83
84    /// Returns provider descriptors in stable name order.
85    #[must_use]
86    pub fn descriptors(&self) -> Vec<ProviderDescriptor> {
87        self.providers
88            .values()
89            .map(|provider| provider.descriptor())
90            .collect()
91    }
92
93    /// Returns dynamic capabilities for one selected provider/model.
94    pub async fn capabilities(
95        &self,
96        provider: Option<&str>,
97        model: Option<&str>,
98    ) -> Result<ProviderCapabilities, BridgeError> {
99        let provider = self.resolve(provider)?;
100        let descriptor = provider.descriptor();
101        let capabilities = provider.capabilities(model).await?;
102        if capabilities.provider != descriptor.name {
103            return Err(BridgeError::new(
104                ErrorCode::Protocol,
105                "provider capability identity does not match its registry identity",
106            )
107            .with_provider(&descriptor.name));
108        }
109        capabilities.validate()?;
110        Ok(capabilities)
111    }
112
113    /// Performs non-generating readiness checks in stable name order.
114    pub async fn readiness(&self) -> Vec<ProviderReadiness> {
115        let mut checks = Vec::with_capacity(self.providers.len());
116        for (name, provider) in self.providers.iter() {
117            let status = match provider.check_ready().await {
118                Ok(()) => ProviderReadinessStatus::Ready,
119                Err(error) => ProviderReadinessStatus::NotReady { error },
120            };
121            checks.push(ProviderReadiness {
122                provider: name.clone(),
123                status,
124            });
125        }
126        checks
127    }
128
129    pub(crate) fn entries(&self) -> impl Iterator<Item = (&str, Arc<dyn ImageProvider>)> + '_ {
130        self.providers
131            .iter()
132            .map(|(name, provider)| (name.as_str(), Arc::clone(provider)))
133    }
134}
135
136/// Redaction-safe readiness result for one provider.
137#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
138#[serde(deny_unknown_fields)]
139pub struct ProviderReadiness {
140    /// Stable provider name.
141    pub provider: String,
142    /// Readiness result without credentials or prompt content.
143    #[serde(flatten)]
144    pub status: ProviderReadinessStatus,
145}
146
147/// Ready or safely classified non-ready state.
148#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
149#[serde(tag = "status", rename_all = "snake_case", deny_unknown_fields)]
150pub enum ProviderReadinessStatus {
151    /// Provider authentication/configuration is ready.
152    Ready,
153    /// Provider cannot currently accept work.
154    NotReady {
155        /// Stable redaction-safe cause.
156        error: BridgeError,
157    },
158}
159
160fn validate_provider_name(name: &str) -> Result<(), BridgeError> {
161    let valid = !name.is_empty()
162        && name.len() <= MAX_PROVIDER_NAME_BYTES
163        && name.as_bytes()[0].is_ascii_lowercase()
164        && name
165            .bytes()
166            .all(|byte| byte.is_ascii_lowercase() || byte.is_ascii_digit() || byte == b'-');
167    if valid {
168        Ok(())
169    } else {
170        Err(registry_error(
171            "provider names must be lowercase ASCII identifiers",
172        ))
173    }
174}
175
176fn registry_error(message: impl Into<String>) -> BridgeError {
177    BridgeError::new(ErrorCode::Configuration, message)
178}
179
180#[cfg(test)]
181mod tests {
182    #![allow(clippy::unwrap_used)]
183
184    use async_trait::async_trait;
185    use imagegen_bridge_core::{ImageRequest, ImageResponse, ProviderContext, ProviderEventStream};
186
187    use super::*;
188
189    struct NamedProvider(&'static str);
190
191    #[async_trait]
192    impl ImageProvider for NamedProvider {
193        fn descriptor(&self) -> ProviderDescriptor {
194            ProviderDescriptor {
195                name: self.0.to_owned(),
196                display_name: self.0.to_owned(),
197                version: "test".to_owned(),
198                experimental: false,
199                models: vec!["test-image".to_owned()],
200            }
201        }
202
203        async fn capabilities(
204            &self,
205            _model: Option<&str>,
206        ) -> Result<ProviderCapabilities, BridgeError> {
207            unreachable!()
208        }
209
210        async fn execute(
211            &self,
212            _request: ImageRequest,
213            _context: ProviderContext,
214        ) -> Result<ImageResponse, BridgeError> {
215            unreachable!()
216        }
217
218        async fn execute_stream(
219            &self,
220            _request: ImageRequest,
221            _context: ProviderContext,
222        ) -> Result<ProviderEventStream, BridgeError> {
223            unreachable!()
224        }
225
226        async fn check_ready(&self) -> Result<(), BridgeError> {
227            Ok(())
228        }
229    }
230
231    fn provider(name: &'static str) -> Arc<dyn ImageProvider> {
232        Arc::new(NamedProvider(name))
233    }
234
235    #[test]
236    fn resolves_default_and_explicit_providers_deterministically() {
237        let registry =
238            ProviderRegistry::new([provider("second"), provider("first")], "first").unwrap();
239        assert_eq!(registry.resolve(None).unwrap().descriptor().name, "first");
240        assert_eq!(
241            registry.resolve(Some("second")).unwrap().descriptor().name,
242            "second"
243        );
244        assert_eq!(
245            registry
246                .descriptors()
247                .into_iter()
248                .map(|value| value.name)
249                .collect::<Vec<_>>(),
250            ["first", "second"]
251        );
252    }
253
254    #[test]
255    fn rejects_collisions_invalid_names_and_missing_defaults() {
256        assert!(ProviderRegistry::new([provider("same"), provider("same")], "same").is_err());
257        assert!(ProviderRegistry::new([provider("Bad_Name")], "Bad_Name").is_err());
258        assert!(ProviderRegistry::new([provider("only")], "missing").is_err());
259    }
260}
imagegen_bridge_runtime/registry.rs

imagegen_bridge_runtime/
registry.rs
