1use crate::audio::AudioInput;
33use crate::config::{Config, ValidatedConfig};
34use crate::doctor::{run_doctor, DoctorReport};
35use crate::error::{Result, UserError};
36use crate::observability::{Metrics, MetricsSnapshot};
37use crate::provider_platform::{
38 ProviderBuildContext, ProviderId, ProviderRegistry, ProviderResolveOptions,
39};
40use crate::providers::local::{LocalWhisperProvider, SttContextPool};
41use crate::providers::{
42 OpenRouterSttMode, TranscriptionOptions, TranscriptionProvider, TranscriptionResult,
43};
44use crate::runtime::{GovernorConfig, ResourceGovernor};
45use crate::support::{build_support_bundle, SupportBundle};
46use std::sync::atomic::{AtomicBool, Ordering};
47use std::sync::Arc;
48use std::time::Instant;
49
50#[cfg(feature = "tts")]
51use crate::tts::local::{LocalTtsProvider, TtsSessionPool};
52#[cfg(feature = "tts")]
53use crate::tts::provider::{SynthesisOptions, SynthesisProvider, SynthesisResult};
54
55pub struct AurumEngine {
57 config: ValidatedConfig,
58 governor: Arc<ResourceGovernor>,
59 metrics: Arc<Metrics>,
60 stt_pool: Arc<SttContextPool>,
61 #[cfg(feature = "tts")]
62 tts_pool: Arc<TtsSessionPool>,
63 registry: Arc<ProviderRegistry>,
64 closed: AtomicBool,
65}
66
67impl std::fmt::Debug for AurumEngine {
68 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
69 let mut d = f.debug_struct("AurumEngine");
70 d.field("config", &self.config)
71 .field("closed", &self.closed.load(Ordering::SeqCst))
72 .field("metrics", &self.metrics.snapshot())
73 .field("stt_resident", &self.stt_pool.resident_len())
74 .field("registry", &*self.registry);
75 #[cfg(feature = "tts")]
76 d.field("tts_resident", &self.tts_pool.resident_len());
77 d.finish_non_exhaustive()
78 }
79}
80
81impl AurumEngine {
82 pub fn new(config: ValidatedConfig) -> Self {
84 Self::with_governor(config, GovernorConfig::default())
85 .expect("default GovernorConfig is always valid")
86 }
87
88 pub fn with_governor(config: ValidatedConfig, gov: GovernorConfig) -> Result<Self> {
92 let registry = ProviderRegistry::builtin()
93 .expect("builtin provider registry must construct (compile-time product factories)");
94 Self::with_governor_and_registry(config, gov, registry)
95 }
96
97 pub fn with_governor_and_registry(
99 config: ValidatedConfig,
100 gov: GovernorConfig,
101 registry: ProviderRegistry,
102 ) -> Result<Self> {
103 let governor = Arc::new(ResourceGovernor::try_new(gov)?);
104 Ok(Self {
105 config,
106 governor,
107 metrics: Arc::new(Metrics::new()),
108 stt_pool: Arc::new(SttContextPool::new()),
109 #[cfg(feature = "tts")]
110 tts_pool: Arc::new(TtsSessionPool::new()),
111 registry: Arc::new(registry),
112 closed: AtomicBool::new(false),
113 })
114 }
115
116 pub fn load() -> Result<Self> {
118 Ok(Self::new(ValidatedConfig::load()?))
119 }
120
121 pub fn load_from_required(path: &std::path::Path) -> Result<Self> {
123 Ok(Self::new(ValidatedConfig::load_from_required(path)?))
124 }
125
126 pub fn from_config(cfg: Config) -> Result<Self> {
128 Ok(Self::new(ValidatedConfig::try_from_config(cfg)?))
129 }
130
131 pub fn config(&self) -> &Config {
132 self.config.as_ref()
133 }
134
135 pub fn validated_config(&self) -> &ValidatedConfig {
136 &self.config
137 }
138
139 pub fn governor(&self) -> &Arc<ResourceGovernor> {
140 &self.governor
141 }
142
143 pub fn metrics(&self) -> &Arc<Metrics> {
144 &self.metrics
145 }
146
147 pub fn stt_pool(&self) -> &Arc<SttContextPool> {
148 &self.stt_pool
149 }
150
151 #[cfg(feature = "tts")]
152 pub fn tts_pool(&self) -> &Arc<TtsSessionPool> {
153 &self.tts_pool
154 }
155
156 pub fn registry(&self) -> &ProviderRegistry {
158 &self.registry
159 }
160
161 pub fn is_closed(&self) -> bool {
162 self.closed.load(Ordering::SeqCst)
163 }
164
165 fn ensure_open(&self) -> Result<()> {
166 if self.is_closed() {
167 return Err(UserError::Other {
168 message: "AurumEngine is closed".into(),
169 }
170 .into());
171 }
172 Ok(())
173 }
174
175 pub fn build_context_for(&self, id: &ProviderId) -> Result<ProviderBuildContext> {
179 self.build_context_for_with(id, ProviderResolveOptions::default())
180 }
181
182 pub fn build_context_for_with(
184 &self,
185 id: &ProviderId,
186 opts: ProviderResolveOptions,
187 ) -> Result<ProviderBuildContext> {
188 self.ensure_open()?;
189 let cfg = self.config.as_ref();
190 let local_only = opts.local_only.unwrap_or(cfg.local_only);
191 let stt_mode = match opts.stt_mode {
192 Some(m) => m,
193 None => OpenRouterSttMode::parse(&cfg.openrouter_stt_mode)?,
194 };
195
196 let mut ctx = ProviderBuildContext::new(self.cache_dir().to_path_buf())
197 .with_local_only(local_only)
198 .with_api_key(cfg.provider_secret(id))
199 .with_show_progress(opts.show_progress)
200 .with_stt_mode(stt_mode)
201 .with_tts_max_chars(Some(cfg.tts_max_chars))
202 .with_stt_pool(Arc::clone(&self.stt_pool))
203 .with_governor(Arc::clone(&self.governor))
204 .with_metrics(Arc::clone(&self.metrics));
205
206 #[cfg(feature = "tts")]
207 {
208 ctx = ctx.with_tts_pool(Arc::clone(&self.tts_pool));
209 }
210
211 match id.as_str() {
213 "openrouter" => {
214 ctx = ctx
215 .with_base_url(Some(cfg.openrouter_base_url.clone()))
216 .with_allow_custom_endpoint(cfg.openrouter_allow_custom_endpoint)
217 .with_use_system_proxy(cfg.openrouter_use_system_proxy);
218 }
219 "openai" => {
220 if let Some(url) = cfg.providers.openai.base_url.clone() {
221 ctx = ctx.with_base_url(Some(url));
222 }
223 }
224 "elevenlabs" => {
225 if let Some(url) = cfg.providers.elevenlabs.base_url.clone() {
226 ctx = ctx.with_base_url(Some(url));
227 }
228 }
229 "xai" => {
230 if let Some(url) = cfg.providers.xai.base_url.clone() {
231 ctx = ctx.with_base_url(Some(url));
232 }
233 }
234 _ => {}
235 }
236
237 Ok(ctx)
238 }
239
240 pub fn stt_provider(&self, id: &ProviderId) -> Result<Arc<dyn TranscriptionProvider>> {
242 self.stt_provider_with(id, ProviderResolveOptions::default())
243 }
244
245 pub fn stt_provider_with(
246 &self,
247 id: &ProviderId,
248 opts: ProviderResolveOptions,
249 ) -> Result<Arc<dyn TranscriptionProvider>> {
250 self.ensure_open()?;
251 let factory = self.registry.stt_factory(id)?;
252 let ctx = self.build_context_for_with(id, opts)?;
253 factory.build(&ctx)
254 }
255
256 #[cfg(feature = "tts")]
258 pub fn tts_provider(&self, id: &ProviderId) -> Result<Arc<dyn SynthesisProvider>> {
259 self.tts_provider_with(id, ProviderResolveOptions::default())
260 }
261
262 #[cfg(feature = "tts")]
263 pub fn tts_provider_with(
264 &self,
265 id: &ProviderId,
266 opts: ProviderResolveOptions,
267 ) -> Result<Arc<dyn SynthesisProvider>> {
268 self.ensure_open()?;
269 let factory = self.registry.tts_factory(id)?;
270 let ctx = self.build_context_for_with(id, opts)?;
271 factory.build(&ctx)
272 }
273
274 pub fn stt_provider_id(&self) -> Result<ProviderId> {
276 ProviderId::parse(&self.config.as_ref().provider)
277 }
278
279 #[cfg(feature = "tts")]
281 pub fn tts_provider_id(&self) -> Result<ProviderId> {
282 ProviderId::parse(&self.config.as_ref().tts_provider)
283 }
284
285 pub fn local_whisper(&self) -> Result<LocalWhisperProvider> {
289 self.ensure_open()?;
290 Ok(LocalWhisperProvider::with_runtime(
291 self.cache_dir().to_path_buf(),
292 Arc::clone(&self.stt_pool),
293 Arc::clone(&self.governor),
294 )
295 .with_progress(false))
296 }
297
298 #[cfg(feature = "tts")]
300 pub fn local_tts(&self) -> Result<LocalTtsProvider> {
301 self.ensure_open()?;
302 Ok(LocalTtsProvider::with_runtime(
303 self.cache_dir().to_path_buf(),
304 Arc::clone(&self.tts_pool),
305 Arc::clone(&self.governor),
306 )
307 .with_progress(false)
308 .with_max_chars(self.config.as_ref().tts_max_chars))
309 }
310
311 pub async fn transcribe(
313 &self,
314 input: &AudioInput,
315 options: &TranscriptionOptions,
316 ) -> Result<TranscriptionResult> {
317 self.ensure_open()?;
318 self.metrics.record_start();
319 let start = Instant::now();
320 let id = self.stt_provider_id()?;
321 let provider = self.stt_provider(&id)?;
322 let out = provider.transcribe(input, options).await;
323 match &out {
324 Ok(_) => self.metrics.record_complete(start.elapsed()),
325 Err(_) => self.metrics.record_failed(),
326 }
327 out
328 }
329
330 pub async fn transcribe_pcm(
335 &self,
336 samples: &[f32],
337 options: &TranscriptionOptions,
338 ) -> Result<TranscriptionResult> {
339 let input = AudioInput::from_pcm_slice(samples, crate::audio::WHISPER_SAMPLE_RATE)?;
340 self.transcribe(&input, options).await
341 }
342
343 pub async fn preload_stt(&self, model: &str) -> Result<std::path::PathBuf> {
345 self.ensure_open()?;
346 self.local_whisper()?.preload(model).await
347 }
348
349 #[cfg(feature = "tts")]
351 pub async fn synthesize(
352 &self,
353 text: &str,
354 options: &SynthesisOptions,
355 ) -> Result<SynthesisResult> {
356 self.ensure_open()?;
357 self.metrics.record_start();
358 let start = Instant::now();
359 let id = self.tts_provider_id()?;
360 let provider = self.tts_provider(&id)?;
361 let out = provider.synthesize(text, options).await;
362 match &out {
363 Ok(_) => self.metrics.record_complete(start.elapsed()),
364 Err(_) => self.metrics.record_failed(),
365 }
366 out
367 }
368
369 pub fn clear_model_caches(&self) {
371 self.stt_pool.clear();
372 #[cfg(feature = "tts")]
373 self.tts_pool.clear();
374 }
375
376 pub fn shutdown(&self) {
380 self.closed.store(true, Ordering::SeqCst);
381 self.clear_model_caches();
382 }
383
384 pub fn doctor(&self) -> DoctorReport {
386 run_doctor(self.config.as_ref())
387 }
388
389 pub fn support_bundle(&self, user_notes: Option<String>) -> SupportBundle {
391 let mut bundle = build_support_bundle(self.config.as_ref(), user_notes);
392 bundle.metrics = self.metrics.snapshot();
393 bundle.redaction_notes.push(format!(
394 "metrics are engine-local; stt_resident={}{}",
395 self.stt_pool.resident_len(),
396 {
397 #[cfg(feature = "tts")]
398 {
399 format!(", tts_resident={}", self.tts_pool.resident_len())
400 }
401 #[cfg(not(feature = "tts"))]
402 {
403 String::new()
404 }
405 }
406 ));
407 bundle
408 }
409
410 pub fn metrics_snapshot(&self) -> MetricsSnapshot {
411 self.metrics.snapshot()
412 }
413
414 pub fn cache_dir(&self) -> &std::path::Path {
416 &self.config.as_ref().cache_dir
417 }
418}
419
420impl Drop for AurumEngine {
421 fn drop(&mut self) {
422 self.closed.store(true, Ordering::SeqCst);
423 self.clear_model_caches();
424 }
425}
426
427#[cfg(test)]
428mod tests {
429 use super::*;
430 use crate::provider_platform::preflight_stt_with_registry;
431
432 #[test]
433 fn independent_engines_have_independent_metrics_and_pools() {
434 let a = AurumEngine::load().unwrap();
435 let b = AurumEngine::load().unwrap();
436 a.metrics().record_start();
437 a.metrics()
438 .record_complete(std::time::Duration::from_millis(1));
439 assert_eq!(a.metrics_snapshot().ops_started, 1);
440 assert_eq!(b.metrics_snapshot().ops_started, 0);
441 assert!(!std::ptr::eq(
442 Arc::as_ptr(a.governor()),
443 Arc::as_ptr(b.governor())
444 ));
445 assert!(!std::ptr::eq(
446 Arc::as_ptr(a.stt_pool()),
447 Arc::as_ptr(b.stt_pool())
448 ));
449 #[cfg(feature = "tts")]
450 assert!(!std::ptr::eq(
451 Arc::as_ptr(a.tts_pool()),
452 Arc::as_ptr(b.tts_pool())
453 ));
454 let process = crate::providers::local::process_global_stt_pool();
456 assert!(!std::ptr::eq(
457 Arc::as_ptr(a.stt_pool()),
458 Arc::as_ptr(&process)
459 ));
460 }
461
462 #[test]
463 fn shutdown_flags_closed_and_rejects_local_whisper() {
464 let e = AurumEngine::load().unwrap();
465 assert!(!e.is_closed());
466 e.shutdown();
467 assert!(e.is_closed());
468 assert!(e.local_whisper().is_err());
469 assert!(e.stt_provider(&ProviderId::local()).is_err());
470 }
471
472 #[test]
473 fn doctor_and_support_bundle_work() {
474 let e = AurumEngine::load().unwrap();
475 let d = e.doctor();
476 assert!(!d.checks.is_empty());
477 let b = e.support_bundle(None);
478 assert_eq!(b.schema_version, crate::support::SUPPORT_BUNDLE_VERSION);
479 let json = b.to_json_pretty().unwrap();
480 assert!(json.contains("engine-local") || json.contains("stt_resident"));
481 }
482
483 #[test]
484 fn local_whisper_uses_engine_pool() {
485 let e = AurumEngine::load().unwrap();
486 let p = e.local_whisper().unwrap();
487 assert!(std::ptr::eq(
488 Arc::as_ptr(p.pool()),
489 Arc::as_ptr(e.stt_pool())
490 ));
491 assert!(std::ptr::eq(
492 Arc::as_ptr(p.governor()),
493 Arc::as_ptr(e.governor())
494 ));
495 }
496
497 #[test]
498 fn registry_stt_local_builds() {
499 let e = AurumEngine::load().unwrap();
500 let p = e.stt_provider(&ProviderId::local()).unwrap();
501 assert_eq!(p.name(), "local");
502 }
503
504 #[test]
505 fn registry_unknown_stt_fails_closed() {
506 let e = AurumEngine::load().unwrap();
507 let err = match e.stt_provider(&ProviderId::must("elevenlabs")) {
509 Ok(_) => panic!("expected unknown STT factory error"),
510 Err(e) => e,
511 };
512 assert!(err.to_string().contains("elevenlabs") || err.to_string().contains("provider"));
513 }
514
515 #[test]
516 fn openai_stt_builds_with_key() {
517 let mut cfg = Config::load().unwrap();
518 cfg.providers.openai.api_key = Some(crate::secret::SecretString::new("sk-test-openai-key"));
519 let e = AurumEngine::from_config(cfg).unwrap();
520 let p = e.stt_provider(&ProviderId::must("openai")).unwrap();
521 assert_eq!(p.name(), "openai");
522 }
523
524 #[test]
525 fn openrouter_local_only_rejected() {
526 let mut cfg = Config::load().unwrap();
527 cfg.local_only = true;
528 let e = AurumEngine::from_config(cfg).unwrap();
529 let err = match e.stt_provider_with(
530 &ProviderId::openrouter(),
531 ProviderResolveOptions {
532 local_only: Some(true),
533 ..Default::default()
534 },
535 ) {
536 Ok(_) => panic!("expected local_only rejection"),
537 Err(e) => e,
538 };
539 assert!(
540 err.to_string().contains("local_only")
541 || err.to_string().contains("network")
542 || err.to_string().contains("remote")
543 );
544 }
545
546 #[test]
547 fn openrouter_missing_key_fails() {
548 let mut cfg = Config::load().unwrap();
549 cfg.openrouter_api_key = None;
550 let e = AurumEngine::from_config(cfg).unwrap();
551 let err = match e.stt_provider(&ProviderId::openrouter()) {
552 Ok(_) => panic!("expected missing key"),
553 Err(e) => e,
554 };
555 let s = err.to_string().to_ascii_lowercase();
556 assert!(
557 s.contains("api") || s.contains("key") || s.contains("auth"),
558 "unexpected: {s}"
559 );
560 }
561
562 #[test]
563 fn build_context_scopes_secret_to_id() {
564 let mut cfg = Config::load().unwrap();
565 cfg.openrouter_api_key = Some(crate::secret::SecretString::new("sk-or-test-secret"));
566 let e = AurumEngine::from_config(cfg).unwrap();
567 let local_ctx = e.build_context_for(&ProviderId::local()).unwrap();
568 assert!(!local_ctx.has_api_key());
569 let or_ctx = e.build_context_for(&ProviderId::openrouter()).unwrap();
570 assert!(or_ctx.has_api_key());
571 let dbg = format!("{or_ctx:?}");
572 assert!(!dbg.contains("sk-or-test"));
573 }
574
575 #[test]
576 fn preflight_openrouter_local_only() {
577 let e = AurumEngine::load().unwrap();
578 let err = preflight_stt_with_registry(
579 e.registry(),
580 &ProviderId::openrouter(),
581 "openai/whisper-large-v3",
582 false,
583 true,
584 OpenRouterSttMode::Auto,
585 )
586 .unwrap_err();
587 assert!(err.to_string().contains("network") || err.to_string().contains("local"));
588 }
589
590 #[cfg(feature = "tts")]
591 #[test]
592 fn registry_tts_local_builds() {
593 let e = AurumEngine::load().unwrap();
594 let p = e.tts_provider(&ProviderId::local()).unwrap();
595 assert_eq!(p.name(), "local");
596 }
597
598 #[test]
599 fn shutdown_rejects_stt_provider() {
600 let e = AurumEngine::load().unwrap();
601 e.shutdown();
602 assert!(e.stt_provider(&ProviderId::local()).is_err());
603 assert!(e.build_context_for(&ProviderId::local()).is_err());
604 }
605}