Skip to main content

aurum_core/
engine.rs

1//! Owned engine boundary for library hosts (JOE-1782 / JOE-1654 / JOE-1784 / JOE-1787).
2//!
3//! # Ownership model
4//!
5//! [`AurumEngine`] owns:
6//! * a [`ValidatedConfig`]
7//! * an engine-local [`ResourceGovernor`]
8//! * an engine-local [`Metrics`] sink
9//! * an engine-local STT context pool ([`SttContextPool`])
10//! * an engine-local TTS session pool (when the `tts` feature is enabled)
11//! * lifecycle bookkeeping for explicit shutdown
12//!
13//! # Isolation (JOE-1784)
14//!
15//! Engines do **not** share whisper/TTS residency with each other or with the
16//! process-global pools used by default `LocalWhisperProvider::new` /
17//! `LocalTtsProvider::new`. Shutdown clears **idle** entries in this engine's
18//! pools only.
19//!
20//! Process-global pools remain for CLI and callers that construct providers
21//! without an engine. Call [`crate::providers::local::clear_context_cache`] at
22//! process exit when using those paths with Metal.
23
24use crate::audio::AudioInput;
25use crate::config::{Config, ValidatedConfig};
26use crate::doctor::{run_doctor, DoctorReport};
27use crate::error::{Result, UserError};
28use crate::observability::{Metrics, MetricsSnapshot};
29use crate::providers::local::{LocalWhisperProvider, SttContextPool};
30use crate::providers::{TranscriptionOptions, TranscriptionProvider, TranscriptionResult};
31use crate::runtime::{GovernorConfig, ResourceGovernor};
32use crate::support::{build_support_bundle, SupportBundle};
33use std::sync::atomic::{AtomicBool, Ordering};
34use std::sync::Arc;
35use std::time::Instant;
36
37#[cfg(feature = "tts")]
38use crate::tts::local::{LocalTtsProvider, TtsSessionPool};
39#[cfg(feature = "tts")]
40use crate::tts::provider::{SynthesisOptions, SynthesisProvider, SynthesisResult};
41
42/// Library-facing engine: validated config + owned governor/metrics/model pools.
43pub struct AurumEngine {
44    config: ValidatedConfig,
45    governor: Arc<ResourceGovernor>,
46    metrics: Arc<Metrics>,
47    stt_pool: Arc<SttContextPool>,
48    #[cfg(feature = "tts")]
49    tts_pool: Arc<TtsSessionPool>,
50    closed: AtomicBool,
51}
52
53impl std::fmt::Debug for AurumEngine {
54    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
55        let mut d = f.debug_struct("AurumEngine");
56        d.field("config", &self.config)
57            .field("closed", &self.closed.load(Ordering::SeqCst))
58            .field("metrics", &self.metrics.snapshot())
59            .field("stt_resident", &self.stt_pool.resident_len());
60        #[cfg(feature = "tts")]
61        d.field("tts_resident", &self.tts_pool.resident_len());
62        d.finish_non_exhaustive()
63    }
64}
65
66impl AurumEngine {
67    /// Build from an already-validated config with default governor settings.
68    pub fn new(config: ValidatedConfig) -> Self {
69        Self::with_governor(config, GovernorConfig::default())
70            .expect("default GovernorConfig is always valid")
71    }
72
73    /// Build with an explicit governor profile (mobile/server/custom).
74    ///
75    /// Validates `gov` before construction (JOE-1917 / F-004).
76    pub fn with_governor(config: ValidatedConfig, gov: GovernorConfig) -> Result<Self> {
77        let governor = Arc::new(ResourceGovernor::try_new(gov)?);
78        Ok(Self {
79            config,
80            governor,
81            metrics: Arc::new(Metrics::new()),
82            stt_pool: Arc::new(SttContextPool::new()),
83            #[cfg(feature = "tts")]
84            tts_pool: Arc::new(TtsSessionPool::new()),
85            closed: AtomicBool::new(false),
86        })
87    }
88
89    /// Load config from the default file/env path and validate.
90    pub fn load() -> Result<Self> {
91        Ok(Self::new(ValidatedConfig::load()?))
92    }
93
94    /// Load from an explicit config file path (must exist).
95    pub fn load_from_required(path: &std::path::Path) -> Result<Self> {
96        Ok(Self::new(ValidatedConfig::load_from_required(path)?))
97    }
98
99    /// Validate a raw [`Config`] and wrap it.
100    pub fn from_config(cfg: Config) -> Result<Self> {
101        Ok(Self::new(ValidatedConfig::try_from_config(cfg)?))
102    }
103
104    pub fn config(&self) -> &Config {
105        self.config.as_ref()
106    }
107
108    pub fn validated_config(&self) -> &ValidatedConfig {
109        &self.config
110    }
111
112    pub fn governor(&self) -> &Arc<ResourceGovernor> {
113        &self.governor
114    }
115
116    pub fn metrics(&self) -> &Arc<Metrics> {
117        &self.metrics
118    }
119
120    pub fn stt_pool(&self) -> &Arc<SttContextPool> {
121        &self.stt_pool
122    }
123
124    #[cfg(feature = "tts")]
125    pub fn tts_pool(&self) -> &Arc<TtsSessionPool> {
126        &self.tts_pool
127    }
128
129    pub fn is_closed(&self) -> bool {
130        self.closed.load(Ordering::SeqCst)
131    }
132
133    fn ensure_open(&self) -> Result<()> {
134        if self.is_closed() {
135            return Err(UserError::Other {
136                message: "AurumEngine is closed".into(),
137            }
138            .into());
139        }
140        Ok(())
141    }
142
143    /// Local STT provider bound to this engine's pool and governor (JOE-1784).
144    pub fn local_whisper(&self) -> Result<LocalWhisperProvider> {
145        self.ensure_open()?;
146        Ok(LocalWhisperProvider::with_runtime(
147            self.cache_dir().to_path_buf(),
148            Arc::clone(&self.stt_pool),
149            Arc::clone(&self.governor),
150        )
151        .with_progress(false))
152    }
153
154    /// Local TTS provider bound to this engine's pool and governor (JOE-1784).
155    #[cfg(feature = "tts")]
156    pub fn local_tts(&self) -> Result<LocalTtsProvider> {
157        self.ensure_open()?;
158        Ok(LocalTtsProvider::with_runtime(
159            self.cache_dir().to_path_buf(),
160            Arc::clone(&self.tts_pool),
161            Arc::clone(&self.governor),
162        )
163        .with_progress(false)
164        .with_max_chars(self.config.as_ref().tts_max_chars))
165    }
166
167    /// High-level STT from a prepared [`AudioInput`] (JOE-1787).
168    pub async fn transcribe(
169        &self,
170        input: &AudioInput,
171        options: &TranscriptionOptions,
172    ) -> Result<TranscriptionResult> {
173        self.ensure_open()?;
174        self.metrics.record_start();
175        let start = Instant::now();
176        let provider = self.local_whisper()?.with_progress(false);
177        let out = provider.transcribe(input, options).await;
178        match &out {
179            Ok(_) => self.metrics.record_complete(start.elapsed()),
180            Err(_) => self.metrics.record_failed(),
181        }
182        out
183    }
184
185    /// High-level STT from mono PCM @ whisper sample rate (JOE-1787).
186    pub async fn transcribe_pcm(
187        &self,
188        samples: &[f32],
189        options: &TranscriptionOptions,
190    ) -> Result<TranscriptionResult> {
191        self.ensure_open()?;
192        self.metrics.record_start();
193        let start = Instant::now();
194        let provider = self.local_whisper()?.with_progress(false);
195        let out = provider.transcribe_pcm(samples, options).await;
196        match &out {
197            Ok(_) => self.metrics.record_complete(start.elapsed()),
198            Err(_) => self.metrics.record_failed(),
199        }
200        out
201    }
202
203    /// Preload a local STT model into **this** engine's pool.
204    pub async fn preload_stt(&self, model: &str) -> Result<std::path::PathBuf> {
205        self.ensure_open()?;
206        self.local_whisper()?.preload(model).await
207    }
208
209    /// High-level local TTS synthesis (JOE-1787).
210    #[cfg(feature = "tts")]
211    pub async fn synthesize(
212        &self,
213        text: &str,
214        options: &SynthesisOptions,
215    ) -> Result<SynthesisResult> {
216        self.ensure_open()?;
217        self.metrics.record_start();
218        let start = Instant::now();
219        let provider = self.local_tts()?;
220        let out = provider.synthesize(text, options).await;
221        match &out {
222            Ok(_) => self.metrics.record_complete(start.elapsed()),
223            Err(_) => self.metrics.record_failed(),
224        }
225        out
226    }
227
228    /// Drop idle model residency in **this** engine's pools (JOE-1784).
229    pub fn clear_model_caches(&self) {
230        self.stt_pool.clear();
231        #[cfg(feature = "tts")]
232        self.tts_pool.clear();
233    }
234
235    /// Mark the engine closed and clear idle model caches.
236    ///
237    /// Does not touch process-global pools used by non-engine providers.
238    pub fn shutdown(&self) {
239        self.closed.store(true, Ordering::SeqCst);
240        self.clear_model_caches();
241    }
242
243    /// Read-only doctor report using this engine's config.
244    pub fn doctor(&self) -> DoctorReport {
245        run_doctor(self.config.as_ref())
246    }
247
248    /// Privacy-safe support bundle using this engine's config and **engine** metrics.
249    pub fn support_bundle(&self, user_notes: Option<String>) -> SupportBundle {
250        let mut bundle = build_support_bundle(self.config.as_ref(), user_notes);
251        bundle.metrics = self.metrics.snapshot();
252        bundle.redaction_notes.push(format!(
253            "metrics are engine-local; stt_resident={}{}",
254            self.stt_pool.resident_len(),
255            {
256                #[cfg(feature = "tts")]
257                {
258                    format!(", tts_resident={}", self.tts_pool.resident_len())
259                }
260                #[cfg(not(feature = "tts"))]
261                {
262                    String::new()
263                }
264            }
265        ));
266        bundle
267    }
268
269    pub fn metrics_snapshot(&self) -> MetricsSnapshot {
270        self.metrics.snapshot()
271    }
272
273    /// Cache directory from validated config.
274    pub fn cache_dir(&self) -> &std::path::Path {
275        &self.config.as_ref().cache_dir
276    }
277}
278
279impl Drop for AurumEngine {
280    fn drop(&mut self) {
281        self.closed.store(true, Ordering::SeqCst);
282        self.clear_model_caches();
283    }
284}
285
286#[cfg(test)]
287mod tests {
288    use super::*;
289
290    #[test]
291    fn independent_engines_have_independent_metrics_and_pools() {
292        let a = AurumEngine::load().unwrap();
293        let b = AurumEngine::load().unwrap();
294        a.metrics().record_start();
295        a.metrics()
296            .record_complete(std::time::Duration::from_millis(1));
297        assert_eq!(a.metrics_snapshot().ops_started, 1);
298        assert_eq!(b.metrics_snapshot().ops_started, 0);
299        assert!(!std::ptr::eq(
300            Arc::as_ptr(a.governor()),
301            Arc::as_ptr(b.governor())
302        ));
303        assert!(!std::ptr::eq(
304            Arc::as_ptr(a.stt_pool()),
305            Arc::as_ptr(b.stt_pool())
306        ));
307        #[cfg(feature = "tts")]
308        assert!(!std::ptr::eq(
309            Arc::as_ptr(a.tts_pool()),
310            Arc::as_ptr(b.tts_pool())
311        ));
312        // Process-global default pool is distinct from engine pools.
313        let process = crate::providers::local::process_global_stt_pool();
314        assert!(!std::ptr::eq(
315            Arc::as_ptr(a.stt_pool()),
316            Arc::as_ptr(&process)
317        ));
318    }
319
320    #[test]
321    fn shutdown_flags_closed_and_rejects_local_whisper() {
322        let e = AurumEngine::load().unwrap();
323        assert!(!e.is_closed());
324        e.shutdown();
325        assert!(e.is_closed());
326        assert!(e.local_whisper().is_err());
327    }
328
329    #[test]
330    fn doctor_and_support_bundle_work() {
331        let e = AurumEngine::load().unwrap();
332        let d = e.doctor();
333        assert!(!d.checks.is_empty());
334        let b = e.support_bundle(None);
335        assert_eq!(b.schema_version, crate::support::SUPPORT_BUNDLE_VERSION);
336        let json = b.to_json_pretty().unwrap();
337        assert!(json.contains("engine-local") || json.contains("stt_resident"));
338    }
339
340    #[test]
341    fn local_whisper_uses_engine_pool() {
342        let e = AurumEngine::load().unwrap();
343        let p = e.local_whisper().unwrap();
344        assert!(std::ptr::eq(
345            Arc::as_ptr(p.pool()),
346            Arc::as_ptr(e.stt_pool())
347        ));
348        assert!(std::ptr::eq(
349            Arc::as_ptr(p.governor()),
350            Arc::as_ptr(e.governor())
351        ));
352    }
353}