Lumotia/crates/transcription/src/registry.rs

//! `EngineRegistry` is the catalogue of available providers, built
//! once at app boot and dependency-injected into the orchestrator.
//!
//! Registration is push-style: the boot code constructs each provider
//! (a `LocalProviderAdapter` wrapping a `LocalEngine`, or a cloud
//! provider implementing `TranscriptionProvider` directly) and calls
//! `register`. The default provider is set at construction time and
//! used when a profile does not specify one.
//!
//! The registry is read-only after boot; mutation requires a fresh
//! registry instance. Providers are held behind `Arc<dyn
//! TranscriptionProvider>` so the orchestrator can clone the handle
//! cheaply across tokio tasks.

use std::collections::HashMap;
use std::sync::Arc;

use lumotia_cloud_providers::{ProviderId, TranscriptionProvider};

/// Catalogue of providers known to the orchestrator.
pub struct EngineRegistry {
    providers: HashMap<ProviderId, Arc<dyn TranscriptionProvider>>,
    default_provider: ProviderId,
}

impl EngineRegistry {
    /// Construct an empty registry with the given default provider id.
    /// The default need not exist at construction time; callers are
    /// expected to register it before the registry is queried.
    pub fn new(default_provider: ProviderId) -> Self {
        Self {
            providers: HashMap::new(),
            default_provider,
        }
    }

    /// Register a provider. If a provider with the same id is already
    /// registered, it is replaced. The provider's id is read once via
    /// `provider.provider_id()` and used as the registry key.
    pub fn register(&mut self, provider: Arc<dyn TranscriptionProvider>) {
        let id = provider.provider_id();
        self.providers.insert(id, provider);
    }

    /// Fetch a provider by id. Returns `None` when the id is not in the
    /// registry; the orchestrator surfaces this as a clear error so the
    /// UI can prompt the user to pick a different engine.
    pub fn get(&self, id: &ProviderId) -> Option<Arc<dyn TranscriptionProvider>> {
        self.providers.get(id).cloned()
    }

    /// Fetch the default provider. Returns `None` if the default has
    /// not been registered.
    pub fn default(&self) -> Option<Arc<dyn TranscriptionProvider>> {
        self.providers.get(&self.default_provider).cloned()
    }

    /// The id of the default provider. Always returns; the provider
    /// itself may not be registered.
    pub fn default_id(&self) -> &ProviderId {
        &self.default_provider
    }

    /// All registered provider ids. The order is unspecified; the UI
    /// is responsible for any sorting it wants to present.
    pub fn ids(&self) -> Vec<ProviderId> {
        self.providers.keys().cloned().collect()
    }

    /// Number of registered providers.
    pub fn len(&self) -> usize {
        self.providers.len()
    }

    pub fn is_empty(&self) -> bool {
        self.providers.is_empty()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    use async_trait::async_trait;
    use lumotia_cloud_providers::{
        CostClass, EngineProfile, ProviderCapabilities, ProviderKind, ProviderTranscript,
    };
    use lumotia_core::error::Result;
    use lumotia_core::types::{AudioSamples, Transcript, TranscriptionOptions};

    struct DummyProvider {
        id: ProviderId,
    }

    #[async_trait]
    impl TranscriptionProvider for DummyProvider {
        fn provider_id(&self) -> ProviderId {
            self.id.clone()
        }
        fn kind(&self) -> ProviderKind {
            ProviderKind::Local
        }
        fn capabilities(&self) -> ProviderCapabilities {
            ProviderCapabilities {
                sample_rate: 16_000,
                channels: 1,
                initial_prompt_supported: false,
                language_hint_supported: true,
                streaming_supported: false,
                cost_class: CostClass::Free,
            }
        }
        async fn prepare(&self, _profile: &EngineProfile) -> Result<()> {
            Ok(())
        }
        async fn transcribe(
            &self,
            audio: AudioSamples,
            _options: TranscriptionOptions,
        ) -> Result<ProviderTranscript> {
            Ok(ProviderTranscript {
                transcript: Transcript::new(Vec::new(), "en".to_string(), audio.duration_secs()),
                inference_ms: 0,
            })
        }
    }

    fn dummy(id: &str) -> Arc<dyn TranscriptionProvider> {
        Arc::new(DummyProvider {
            id: ProviderId::new(id),
        })
    }

    #[test]
    fn empty_registry_returns_none_for_default() {
        let registry = EngineRegistry::new(ProviderId::new("local-whisper"));
        assert!(registry.default().is_none());
        assert_eq!(registry.default_id().as_str(), "local-whisper");
        assert!(registry.is_empty());
    }

    #[test]
    fn register_and_get_round_trip() {
        let mut registry = EngineRegistry::new(ProviderId::new("local-whisper"));
        registry.register(dummy("local-whisper"));
        registry.register(dummy("local-parakeet"));

        assert_eq!(registry.len(), 2);
        assert!(registry.get(&ProviderId::new("local-whisper")).is_some());
        assert!(registry.get(&ProviderId::new("local-parakeet")).is_some());
        assert!(registry.get(&ProviderId::new("nonexistent")).is_none());
    }

    #[test]
    fn default_resolves_after_registration() {
        let mut registry = EngineRegistry::new(ProviderId::new("local-whisper"));
        registry.register(dummy("local-whisper"));
        let default = registry.default().expect("default provider registered");
        assert_eq!(default.provider_id().as_str(), "local-whisper");
    }

    #[test]
    fn re_register_replaces() {
        let mut registry = EngineRegistry::new(ProviderId::new("local-whisper"));
        registry.register(dummy("local-whisper"));
        registry.register(dummy("local-whisper"));
        assert_eq!(registry.len(), 1);
    }

    #[test]
    fn ids_lists_all_registered() {
        let mut registry = EngineRegistry::new(ProviderId::new("local-whisper"));
        registry.register(dummy("local-whisper"));
        registry.register(dummy("local-parakeet"));
        let mut ids: Vec<String> = registry
            .ids()
            .into_iter()
            .map(|id| id.as_str().to_string())
            .collect();
        ids.sort();
        assert_eq!(ids, vec!["local-parakeet", "local-whisper"]);
    }
}