Skip to main content

aurum_core/tts/
local.rs

1//! Local ONNX KittenTTS provider (MIT binary path — no GPL phonemizer).
2//!
3//! Session loading is singleflight-coalesced (JOE-1597). Concurrent synthesis is
4//! bounded by the process [`ResourceGovernor`] (JOE-1596/1600). Caller timeouts
5//! never abandon untracked native work: the permit stays held until the ONNX
6//! job returns (honest soft-deadline semantics).
7
8use super::adapter::{
9    TrustMode, ADAPTER_FAKE_SINE_V1, ADAPTER_KITTEN_ONNX_V1, ADAPTER_KOKORO_ONNX_V0,
10};
11use super::catalogue::{
12    ensure_voice_pack, lookup_model, onnx_path, resolve_voice_for_model, validate_speaking_rate,
13    voices_path,
14};
15use super::chunk::{prepare_tts_chunks_with, PhonemeCodec, TtsChunk, CHUNK_PAUSE_MS};
16use super::conformance::synthesize_fake_sine_ms;
17use super::npz::load_voices_npz;
18use super::pack::load_pack_dir;
19use super::pcm_post::{
20    duration_ms_from_pcm, trim_trailing_silence, validate_raw_pcm, TailTrimPolicy, PEAK_LIMIT,
21};
22use super::provider::{BackendKind, SynthesisOptions, SynthesisProvider, SynthesisResult};
23use super::validate::{
24    normalize_tts_language, prepare_text, resolve_sample_rate, DEFAULT_MAX_CHARS,
25};
26use super::wav::peak_guard_f32_to_i16;
27use crate::error::{ProviderError, Result, UserError};
28use crate::runtime::{
29    LoadKey, ModelRegistry, OpContext, RegistryConfig, RegistryPin, ResidencyWeight,
30    ResourceGovernor, Singleflight,
31};
32use async_trait::async_trait;
33use once_cell::sync::Lazy;
34use ort::session::Session;
35use ort::value::Tensor;
36use std::collections::HashMap;
37use std::path::{Path, PathBuf};
38use std::sync::{Arc, Mutex};
39use std::time::Duration;
40
41/// TTS pack residency + singleflight loader (JOE-1784).
42///
43/// Own one per [`crate::AurumEngine`], or share via [`process_global_tts_pool`].
44pub struct TtsSessionPool {
45    flight: Singleflight<LoadedPack>,
46    registry: ModelRegistry<LoadedPack>,
47}
48
49impl Default for TtsSessionPool {
50    fn default() -> Self {
51        Self::new()
52    }
53}
54
55impl TtsSessionPool {
56    pub fn new() -> Self {
57        Self {
58            flight: Singleflight::default(),
59            registry: ModelRegistry::new(RegistryConfig::default()),
60        }
61    }
62
63    pub fn resident_len(&self) -> usize {
64        self.registry.len()
65    }
66
67    fn weight_for(onnx: &Path) -> ResidencyWeight {
68        let disk = std::fs::metadata(onnx).map(|m| m.len()).unwrap_or(0);
69        // ONNX + voices + runtime overhead headroom.
70        let bytes = disk.saturating_mul(2).max(32 * 1024 * 1024);
71        ResidencyWeight { bytes }
72    }
73
74    /// Load or reuse a pack and return an **active registry pin** for the full
75    /// synthesis operation (JOE-1646).
76    fn get_or_load_pin<F>(
77        &self,
78        key: LoadKey,
79        weight: ResidencyWeight,
80        gov: &Arc<ResourceGovernor>,
81        loader: F,
82    ) -> Result<RegistryPin<LoadedPack>>
83    where
84        F: FnOnce() -> Result<LoadedPack>,
85    {
86        if weight.bytes > self.registry.config().max_resident_bytes {
87            return Err(ProviderError::Overload {
88                reason: format!(
89                    "TTS pack weight {} exceeds residency budget {}",
90                    weight.bytes,
91                    self.registry.config().max_resident_bytes
92                ),
93            }
94            .into());
95        }
96        loop {
97            if let Some(pin) = self.registry.get_and_pin(&key) {
98                return Ok(pin);
99            }
100            match self.flight.begin_or_wait_guard(key.clone()) {
101                Ok(None) => continue,
102                Err(message) => {
103                    return Err(ProviderError::ModelLoad {
104                        model: key.id.clone(),
105                        reason: message,
106                    }
107                    .into());
108                }
109                Ok(Some(leader)) => {
110                    let gov = Arc::clone(gov);
111                    let load_result = (|| -> Result<Arc<LoadedPack>> {
112                        let _permit = gov.acquire(crate::runtime::PermitKind::ModelLoad, None)?;
113                        Ok(Arc::new(loader()?))
114                    })();
115                    match load_result {
116                        Ok(pack) => {
117                            match self.registry.insert_and_pin(
118                                key.clone(),
119                                Arc::clone(&pack),
120                                weight,
121                            ) {
122                                Ok(pin) => {
123                                    leader.success();
124                                    return Ok(pin);
125                                }
126                                Err(e) => {
127                                    leader.fail(e.to_string());
128                                    drop(pack);
129                                    return Err(e);
130                                }
131                            }
132                        }
133                        Err(e) => {
134                            leader.fail(e.to_string());
135                            return Err(e);
136                        }
137                    }
138                }
139            }
140        }
141    }
142
143    /// Drop idle sessions only. Active registry pins retain their entry and
144    /// weight so a concurrent clear cannot force a second same-key load
145    /// (JOE-1646 third-pass residual).
146    pub fn clear(&self) {
147        let _ = self.registry.clear_idle();
148        // Do not wipe in-flight singleflight Loading slots (would strand leaders).
149        // Failed TTL slots can stay until natural expiry.
150    }
151}
152
153static PROCESS_TTS_POOL: Lazy<Arc<TtsSessionPool>> = Lazy::new(|| Arc::new(TtsSessionPool::new()));
154
155/// Process-global TTS pool used by default providers / CLI (JOE-1784).
156pub fn process_global_tts_pool() -> Arc<TtsSessionPool> {
157    Arc::clone(&PROCESS_TTS_POOL)
158}
159
160/// On-device KittenTTS via ONNX Runtime + misaki-rs G2P (no espeak / GPL).
161///
162/// Pool size is intentionally 1 session per model (serial inference under a
163/// mutex). Throughput concurrency is governed by ResourceGovernor TTS permits
164/// rather than unbounded parallel ONNX sessions (memory cost of multi-session
165/// pools is not free for mobile/desktop defaults).
166///
167/// Loaded packs participate in a weighted residency registry (JOE-1646 /
168/// JOE-1784). Default providers share the process-global pool; engines own
169/// isolated pools.
170pub struct LocalTtsProvider {
171    cache_dir: PathBuf,
172    show_progress: bool,
173    local_only: bool,
174    max_chars: usize,
175    pool: Arc<TtsSessionPool>,
176    governor: Arc<ResourceGovernor>,
177}
178
179struct LoadedPack {
180    session: Mutex<Session>,
181    voices: HashMap<String, super::npz::VoiceMatrix>,
182    sample_rate_hz: u32,
183    /// Catalogue / pack phoneme capacity (min of matrix rows and catalogue).
184    max_phoneme_tokens: usize,
185    /// Optional speed priors from config.json (internal key → multiplier).
186    speed_priors: HashMap<String, f32>,
187}
188
189impl LocalTtsProvider {
190    pub fn new(cache_dir: PathBuf) -> Self {
191        Self::with_runtime(
192            cache_dir,
193            process_global_tts_pool(),
194            ResourceGovernor::process_global(),
195        )
196    }
197
198    /// Construct with an explicit TTS pool and governor (JOE-1784 / engine path).
199    pub fn with_runtime(
200        cache_dir: PathBuf,
201        pool: Arc<TtsSessionPool>,
202        governor: Arc<ResourceGovernor>,
203    ) -> Self {
204        Self {
205            cache_dir,
206            show_progress: false,
207            local_only: false,
208            max_chars: DEFAULT_MAX_CHARS,
209            pool,
210            governor,
211        }
212    }
213
214    pub fn with_progress(mut self, v: bool) -> Self {
215        self.show_progress = v;
216        self
217    }
218
219    pub fn with_local_only(mut self, v: bool) -> Self {
220        self.local_only = v;
221        self
222    }
223
224    pub fn with_max_chars(mut self, n: usize) -> Self {
225        self.max_chars = n.max(1);
226        self
227    }
228
229    pub fn pool(&self) -> &Arc<TtsSessionPool> {
230        &self.pool
231    }
232
233    /// Drop **idle** loaded ONNX sessions in **this provider's** pool.
234    ///
235    /// Active synthesis leases keep their registry entry and residency weight
236    /// until the pin drops — a clear during an in-flight synth cannot make the
237    /// same key reloadable as a second native session (JOE-1646).
238    ///
239    /// Does not delete cached files under the TTS cache directory.
240    pub fn clear_sessions(&self) {
241        self.pool.clear();
242    }
243
244    async fn ensure_loaded_pin(
245        &self,
246        model: &str,
247        local_only: bool,
248    ) -> Result<RegistryPin<LoadedPack>> {
249        let key = LoadKey::tts(model, self.cache_dir.join(model).display().to_string());
250        if let Some(pin) = self.pool.registry.get_and_pin(&key) {
251            return Ok(pin);
252        }
253
254        let info = lookup_model(model)?;
255        let _pack_dir = ensure_voice_pack(
256            &self.cache_dir,
257            model,
258            self.show_progress,
259            local_only || self.local_only,
260        )
261        .await?;
262
263        let onnx = onnx_path(&self.cache_dir, info);
264        let voices_file = voices_path(&self.cache_dir, info);
265        let speed_priors = load_speed_priors(&self.cache_dir, info);
266        let sample_rate = info.sample_rate_hz;
267        let catalogue_max = info.max_phoneme_tokens;
268        let key_for_load = key.clone();
269        let model_id = model.to_string();
270        let weight = TtsSessionPool::weight_for(&onnx);
271        let pool = Arc::clone(&self.pool);
272        let gov = Arc::clone(&self.governor);
273
274        let pin = tokio::task::spawn_blocking(move || {
275            pool.get_or_load_pin(key_for_load, weight, &gov, || {
276                load_pack(
277                    &onnx,
278                    &voices_file,
279                    sample_rate,
280                    speed_priors,
281                    catalogue_max,
282                    &model_id,
283                )
284            })
285        })
286        .await
287        .map_err(|e| crate::error::TranscriptionError::internal(format!("TTS load join: {e}")))??;
288
289        Ok(pin)
290    }
291
292    /// Load a verified (or explicitly unverified) local model pack for a supported
293    /// adapter (JOE-1619). Cache key is isolated from built-in catalogue identities.
294    async fn ensure_loaded_from_pack_pin(
295        &self,
296        pack_dir: &Path,
297        allow_unverified: bool,
298    ) -> Result<(RegistryPin<LoadedPack>, super::adapter::ModelPackManifest)> {
299        let (root, manifest) = load_pack_dir(pack_dir, allow_unverified)?;
300        if manifest.adapter_id != ADAPTER_KITTEN_ONNX_V1
301            && manifest.adapter_id != ADAPTER_KOKORO_ONNX_V0
302        {
303            return Err(UserError::InvalidConfig {
304                reason: format!(
305                    "local pack override synthesis currently supports adapters \
306                     '{ADAPTER_KITTEN_ONNX_V1}' and '{ADAPTER_KOKORO_ONNX_V0}' \
307                     (got '{}'); use `aurum tts inspect` / conformance for others",
308                    manifest.adapter_id
309                ),
310            }
311            .into());
312        }
313        let onnx_name = manifest
314            .artifact("onnx")
315            .map(|a| a.filename.as_str())
316            .ok_or_else(|| UserError::InvalidConfig {
317                reason: "pack missing onnx artifact".into(),
318            })?;
319        let voices_name = manifest
320            .artifact("voices")
321            .map(|a| a.filename.as_str())
322            .ok_or_else(|| UserError::InvalidConfig {
323                reason: "pack missing voices artifact".into(),
324            })?;
325        let onnx = root.join(onnx_name);
326        let voices_file = root.join(voices_name);
327        let sample_rate = manifest.sample_rate_hz;
328        let catalogue_max = manifest.max_phoneme_tokens;
329        let model_id = manifest.model_id.clone();
330        // Isolate local packs from built-in cache keys.
331        let key = LoadKey::tts(
332            format!("local-pack:{}", manifest.model_id),
333            root.display().to_string(),
334        );
335        if let Some(pin) = self.pool.registry.get_and_pin(&key) {
336            return Ok((pin, manifest));
337        }
338        let key_for_load = key.clone();
339        let speed_priors = load_speed_priors_from_path(
340            &root.join(
341                manifest
342                    .artifact("config")
343                    .map(|a| a.filename.as_str())
344                    .unwrap_or("config.json"),
345            ),
346        );
347        let weight = TtsSessionPool::weight_for(&onnx);
348        let pool = Arc::clone(&self.pool);
349        let gov = Arc::clone(&self.governor);
350
351        let pin = tokio::task::spawn_blocking(move || {
352            pool.get_or_load_pin(key_for_load, weight, &gov, || {
353                load_pack(
354                    &onnx,
355                    &voices_file,
356                    sample_rate,
357                    speed_priors,
358                    catalogue_max,
359                    &model_id,
360                )
361            })
362        })
363        .await
364        .map_err(|e| {
365            crate::error::TranscriptionError::internal(format!("TTS pack load join: {e}"))
366        })??;
367
368        Ok((pin, manifest))
369    }
370}
371
372#[allow(clippy::too_many_arguments)]
373fn synthesize_with_pack(
374    pack: &LoadedPack,
375    text: &str,
376    opts: &SynthesisOptions,
377    voice_internal: &str,
378    voice_canonical: &str,
379    model_canonical: &str,
380    text_chars: usize,
381    op: &OpContext,
382    adapter: &str,
383    trust: TrustMode,
384    provenance: &str,
385) -> Result<SynthesisResult> {
386    op.check()?;
387
388    let rate = validate_speaking_rate(opts.speaking_rate)?;
389    let sample_rate = resolve_sample_rate(opts.sample_rate_hz, pack.sample_rate_hz)?;
390
391    let voice_mat = pack
392        .voices
393        .get(voice_internal)
394        .ok_or_else(|| UserError::Other {
395            message: format!(
396                "voice embedding '{voice_internal}' missing from pack; available: {:?}",
397                pack.voices.keys().collect::<Vec<_>>()
398            ),
399        })?;
400
401    // Voice matrix has one style embedding per supported sequence length.
402    // Reserve one row because the token vector includes start/end pads.
403    let pack_max = voice_mat.nrows.saturating_sub(1);
404    let max_tokens = pack_max.min(pack.max_phoneme_tokens);
405    if max_tokens <= 2 {
406        return Err(ProviderError::ModelLoad {
407            model: model_canonical.to_string(),
408            reason: "voice embedding has no usable sequence rows".into(),
409        }
410        .into());
411    }
412
413    let codec = phoneme_codec_for_adapter(adapter);
414    let chunks = prepare_tts_chunks_with(text, max_tokens, codec)?;
415    let chunk_count = chunks.len();
416    let effective_speed = rate
417        * pack
418            .speed_priors
419            .get(voice_internal)
420            .copied()
421            .unwrap_or(1.0);
422    let pause_samples = (sample_rate as u64)
423        .saturating_mul(CHUNK_PAUSE_MS)
424        .checked_div(1_000)
425        .unwrap_or(0) as usize;
426
427    let mut pcm_f32: Vec<f32> = Vec::new();
428    let trim_policy = TailTrimPolicy::default();
429
430    for (index, chunk) in chunks.iter().enumerate() {
431        op.check()?;
432        let chunk_audio =
433            synthesize_chunk_f32(pack, voice_mat, chunk, effective_speed).map_err(|err| {
434                ProviderError::Other {
435                    message: format!(
436                        "TTS chunk {}/{} failed near {:?}: {err}",
437                        index + 1,
438                        chunks.len(),
439                        chunk.text.chars().take(80).collect::<String>()
440                    ),
441                }
442            })?;
443        let trimmed = trim_trailing_silence(&chunk_audio, trim_policy);
444        if index > 0 {
445            pcm_f32.resize(pcm_f32.len().saturating_add(pause_samples), 0.0);
446        }
447        pcm_f32.extend_from_slice(trimmed);
448    }
449
450    validate_raw_pcm(&pcm_f32, sample_rate)?;
451    let pcm = peak_guard_f32_to_i16(&pcm_f32, PEAK_LIMIT);
452    if pcm.is_empty() {
453        return Err(ProviderError::Other {
454            message: "synthesis produced empty audio after validation".into(),
455        }
456        .into());
457    }
458
459    let duration_ms = duration_ms_from_pcm(pcm.len(), sample_rate);
460    let language = normalize_tts_language(&opts.language).unwrap_or_else(|_| "en".into());
461
462    Ok(SynthesisResult {
463        pcm_i16_mono: pcm,
464        sample_rate_hz: sample_rate,
465        channels: 1,
466        backend_kind: BackendKind::Local,
467        provider: "local".into(),
468        model: model_canonical.to_string(),
469        voice: voice_canonical.to_string(),
470        language,
471        duration_ms,
472        text_chars,
473        text_truncated: false,
474        chunk_count,
475        synthesized_chars: text_chars,
476        adapter: Some(adapter.into()),
477        trust: Some(trust.as_str().into()),
478        provenance: Some(provenance.into()),
479    })
480}
481
482fn phoneme_codec_for_adapter(adapter: &str) -> PhonemeCodec {
483    if adapter == ADAPTER_KOKORO_ONNX_V0 {
484        PhonemeCodec::Kokoro
485    } else {
486        PhonemeCodec::Kitten
487    }
488}
489
490fn synthesize_chunk_f32(
491    pack: &LoadedPack,
492    voice_mat: &super::npz::VoiceMatrix,
493    chunk: &TtsChunk,
494    effective_speed: f32,
495) -> Result<Vec<f32>> {
496    let seq_len = chunk.ids.len();
497    // Kitten and Kokoro both index style embeddings by phoneme sequence length.
498    let style = voice_mat.style_row(seq_len).to_vec();
499    let style_dim = style.len();
500    let t_ids = Tensor::<i64>::from_array(([1usize, seq_len], chunk.ids.clone())).map_err(|e| {
501        ProviderError::Other {
502            message: format!("tokens/input_ids tensor: {e}"),
503        }
504    })?;
505    let t_style = Tensor::<f32>::from_array(([1usize, style_dim], style)).map_err(|e| {
506        ProviderError::Other {
507            message: format!("style tensor: {e}"),
508        }
509    })?;
510    let t_speed = Tensor::<f32>::from_array(([1usize], vec![effective_speed])).map_err(|e| {
511        ProviderError::Other {
512            message: format!("speed tensor: {e}"),
513        }
514    })?;
515
516    let mut session = pack
517        .session
518        .lock()
519        .map_err(|_| crate::error::TranscriptionError::internal("ORT session mutex poisoned"))?;
520    // Positional order matches both Kitten and Kokoro graphs: tokens, style, speed.
521    let outputs = session
522        .run(ort::inputs![t_ids, t_style, t_speed])
523        .map_err(|e| ProviderError::Other {
524            message: format!("ONNX inference failed: {e}"),
525        })?;
526    let (_shape, audio_data) =
527        outputs[0]
528            .try_extract_tensor::<f32>()
529            .map_err(|e| ProviderError::Other {
530                message: format!("extract audio tensor: {e}"),
531            })?;
532    Ok(audio_data.to_vec())
533}
534
535fn load_pack(
536    onnx: &Path,
537    voices_file: &Path,
538    sample_rate_hz: u32,
539    speed_priors: HashMap<String, f32>,
540    catalogue_max_tokens: usize,
541    model_label: &str,
542) -> Result<LoadedPack> {
543    let session = Session::builder()
544        .map_err(|e| ProviderError::ModelLoad {
545            model: model_label.to_string(),
546            reason: format!("ORT session builder: {e}"),
547        })?
548        .commit_from_file(onnx)
549        .map_err(|e| ProviderError::ModelLoad {
550            model: model_label.to_string(),
551            reason: format!("load ONNX: {e}"),
552        })?;
553    let voices = load_voices_npz(voices_file)?;
554    // Derive capacity from the first voice matrix (all Kitten voices share shape).
555    let matrix_max = voices
556        .values()
557        .map(|m| m.nrows.saturating_sub(1))
558        .min()
559        .unwrap_or(0);
560    let max_phoneme_tokens = if matrix_max == 0 {
561        catalogue_max_tokens
562    } else {
563        matrix_max.min(catalogue_max_tokens)
564    };
565    Ok(LoadedPack {
566        session: Mutex::new(session),
567        voices,
568        sample_rate_hz,
569        max_phoneme_tokens,
570        speed_priors,
571    })
572}
573
574fn load_speed_priors(
575    cache_dir: &Path,
576    info: &super::catalogue::TtsModelInfo,
577) -> HashMap<String, f32> {
578    load_speed_priors_from_path(&super::catalogue::config_path(cache_dir, info))
579}
580
581fn load_speed_priors_from_path(path: &Path) -> HashMap<String, f32> {
582    let Ok(bytes) = std::fs::read(path) else {
583        return HashMap::new();
584    };
585    let Ok(v) = serde_json::from_slice::<serde_json::Value>(&bytes) else {
586        return HashMap::new();
587    };
588    let mut out = HashMap::new();
589    if let Some(map) = v.get("speed_priors").and_then(|x| x.as_object()) {
590        for (k, val) in map {
591            if let Some(f) = val.as_f64() {
592                out.insert(k.clone(), f as f32);
593            }
594        }
595    }
596    out
597}
598
599/// Synthesize via the fake-sine adapter (conformance / custom packs; no ONNX).
600fn synthesize_fake_adapter(
601    opts: &SynthesisOptions,
602    model_id: &str,
603    trust: TrustMode,
604    provenance: &str,
605    text_chars: usize,
606) -> Result<SynthesisResult> {
607    // Map text length to a short bounded duration (deterministic, no network).
608    let duration_ms = ((text_chars as u64).saturating_mul(40)).clamp(50, 2_000);
609    let pcm_f32 = synthesize_fake_sine_ms(duration_ms).map_err(|e| ProviderError::Other {
610        message: format!("fake-sine synth: {e}"),
611    })?;
612    validate_raw_pcm(&pcm_f32, 24_000)?;
613    let pcm = peak_guard_f32_to_i16(&pcm_f32, PEAK_LIMIT);
614    let sample_rate = resolve_sample_rate(opts.sample_rate_hz, 24_000)?;
615    let language = normalize_tts_language(&opts.language).unwrap_or_else(|_| "en".into());
616    let voice = if opts.voice.trim().is_empty() {
617        "Tone".into()
618    } else {
619        opts.voice.clone()
620    };
621    let out_duration_ms = duration_ms_from_pcm(pcm.len(), sample_rate);
622    Ok(SynthesisResult {
623        pcm_i16_mono: pcm,
624        sample_rate_hz: sample_rate,
625        channels: 1,
626        backend_kind: BackendKind::Local,
627        provider: "local".into(),
628        model: model_id.into(),
629        voice,
630        language,
631        duration_ms: out_duration_ms,
632        text_chars,
633        text_truncated: false,
634        chunk_count: 1,
635        synthesized_chars: text_chars,
636        adapter: Some(ADAPTER_FAKE_SINE_V1.into()),
637        trust: Some(trust.as_str().into()),
638        provenance: Some(provenance.into()),
639    })
640}
641
642/// Resolve a voice for a local pack: prefer matching catalogue model voices, else pack manifest.
643fn resolve_pack_voice(
644    manifest: &super::adapter::ModelPackManifest,
645    requested: &str,
646) -> Result<(String, String)> {
647    let req = requested.trim();
648    let default_voice = if manifest.adapter_id == ADAPTER_KOKORO_ONNX_V0 {
649        super::catalogue::KOKORO_DEFAULT_VOICE
650    } else {
651        super::catalogue::DEFAULT_TTS_VOICE
652    };
653    let req = if req.is_empty() { default_voice } else { req };
654    if let Ok((_, v)) = resolve_voice_for_model(&manifest.model_id, req) {
655        return Ok((v.id.to_string(), v.internal_key.to_string()));
656    }
657    if let Ok((_, v)) = resolve_voice_for_model(super::catalogue::DEFAULT_TTS_MODEL, req) {
658        return Ok((v.id.to_string(), v.internal_key.to_string()));
659    }
660    if let Some(v) = manifest
661        .voices
662        .iter()
663        .find(|v| v.id.eq_ignore_ascii_case(req) || v.internal_key.eq_ignore_ascii_case(req))
664    {
665        return Ok((v.id.clone(), v.internal_key.clone()));
666    }
667    let available: Vec<_> = manifest.voices.iter().map(|v| v.id.as_str()).collect();
668    Err(UserError::Other {
669        message: format!(
670            "voice '{req}' not found in pack '{}'; available: {available:?}",
671            manifest.model_id
672        ),
673    }
674    .into())
675}
676
677#[async_trait]
678impl SynthesisProvider for LocalTtsProvider {
679    fn name(&self) -> &'static str {
680        "local"
681    }
682
683    async fn synthesize(&self, text: &str, opts: &SynthesisOptions) -> Result<SynthesisResult> {
684        let prepared = prepare_text(text, self.max_chars)?;
685        let mut opts = opts.clone();
686        opts.language = normalize_tts_language(&opts.language)?;
687
688        // Local pack override path (JOE-1619): never hits network, never shadows
689        // built-in cache identity. Bare ONNX is rejected by load_pack_dir.
690        if let Some(pack_dir) = opts.pack_dir.clone() {
691            let (root, manifest) = load_pack_dir(&pack_dir, opts.allow_unverified)?;
692            let _ = root;
693            if manifest.adapter_id == ADAPTER_FAKE_SINE_V1 {
694                let trust = manifest.trust;
695                // Prefer pack identity over the built-in default when the caller
696                // did not name a distinct custom model id.
697                let model_id = if opts.model.trim().is_empty()
698                    || opts.model == super::catalogue::DEFAULT_TTS_MODEL
699                {
700                    manifest.model_id.clone()
701                } else {
702                    opts.model.clone()
703                };
704                if opts.voice.trim().is_empty() || opts.voice == super::catalogue::DEFAULT_TTS_VOICE
705                {
706                    if let Some(v) = manifest.voices.first() {
707                        opts.voice = v.id.clone();
708                    } else {
709                        opts.voice = "Tone".into();
710                    }
711                }
712                return synthesize_fake_adapter(
713                    &opts,
714                    &model_id,
715                    trust,
716                    "local_pack",
717                    prepared.text_chars,
718                );
719            }
720            if manifest.adapter_id != ADAPTER_KITTEN_ONNX_V1
721                && manifest.adapter_id != ADAPTER_KOKORO_ONNX_V0
722            {
723                return Err(UserError::UnsupportedCapability {
724                    provider: "tts".into(),
725                    model: manifest.model_id,
726                    reason: format!(
727                        "adapter '{}' is not enabled for local pack synthesis",
728                        manifest.adapter_id
729                    ),
730                    hint:
731                        "use kitten-onnx-v1, kokoro-onnx-v0, or fake-sine-v1 packs; see `aurum tts adapters`"
732                            .into(),
733                }
734                .into());
735            }
736            // Prefer pack-declared model id for honesty when caller left default.
737            let model_canonical = if opts.model == super::catalogue::DEFAULT_TTS_MODEL
738                || opts.model.trim().is_empty()
739            {
740                manifest.model_id.clone()
741            } else {
742                opts.model.clone()
743            };
744            let (voice_canonical, voice_internal) = resolve_pack_voice(&manifest, &opts.voice)?;
745            validate_speaking_rate(opts.speaking_rate)?;
746            resolve_sample_rate(opts.sample_rate_hz, manifest.sample_rate_hz)?;
747            let trust = manifest.trust;
748            let adapter = manifest.adapter_id.clone();
749            // Registry pin held for the full synthesis (JOE-1646).
750            let lease = self
751                .ensure_loaded_from_pack_pin(&pack_dir, opts.allow_unverified)
752                .await?
753                .0;
754
755            let timeout = Duration::from_millis(if opts.timeout_ms == 0 {
756                super::validate::DEFAULT_TIMEOUT_MS
757            } else {
758                opts.timeout_ms
759            });
760            let op = OpContext::from_optional_cancel(opts.cancel.clone())
761                .with_deadline_from_now(timeout);
762            op.check()?;
763            let text_owned = prepared.text.clone();
764            let text_chars = prepared.text_chars;
765            let opts_owned = opts.clone();
766            let op_for_worker = op.clone();
767            let gov = Arc::clone(&self.governor);
768            let join = tokio::task::spawn_blocking(move || {
769                // Move registry pin into the worker so soft outer deadlines cannot
770                // release residency while ONNX is still running (JOE-1646).
771                let _lease = lease;
772                let _permit = gov.acquire_tts(0, Some(&op_for_worker))?;
773                op_for_worker.check()?;
774                synthesize_with_pack(
775                    _lease.value().as_ref(),
776                    &text_owned,
777                    &opts_owned,
778                    &voice_internal,
779                    &voice_canonical,
780                    &model_canonical,
781                    text_chars,
782                    &op_for_worker,
783                    &adapter,
784                    trust,
785                    "local_pack",
786                )
787            })
788            .await
789            .map_err(|e| {
790                crate::error::TranscriptionError::internal(format!("TTS pack synth join: {e}"))
791            })??;
792            return Ok(join);
793        }
794
795        let (model_info, voice_info) = resolve_voice_for_model(&opts.model, &opts.voice)?;
796        // Canonical IDs for honesty JSON / metadata.
797        opts.model = model_info.id.to_string();
798        opts.voice = voice_info.id.to_string();
799        validate_speaking_rate(opts.speaking_rate)?;
800        resolve_sample_rate(opts.sample_rate_hz, model_info.sample_rate_hz)?;
801
802        let local_only = opts.local_only || self.local_only;
803        let lease = self.ensure_loaded_pin(&opts.model, local_only).await?;
804
805        let timeout = Duration::from_millis(if opts.timeout_ms == 0 {
806            super::validate::DEFAULT_TIMEOUT_MS
807        } else {
808            opts.timeout_ms
809        });
810        let op =
811            OpContext::from_optional_cancel(opts.cancel.clone()).with_deadline_from_now(timeout);
812        op.check()?;
813
814        let text_owned = prepared.text.clone();
815        let text_chars = prepared.text_chars;
816        let opts_owned = opts.clone();
817        let voice_internal = voice_info.internal_key.to_string();
818        let voice_canonical = voice_info.id.to_string();
819        let model_canonical = model_info.id.to_string();
820        let op_for_worker = op.clone();
821        let adapter = model_info.adapter.to_string();
822        let gov = Arc::clone(&self.governor);
823
824        // Hold the TTS + blocking permit and registry pin for the entire native
825        // job lifetime — even if the caller soft-times out (JOE-1600 / JOE-1646).
826        let join = tokio::task::spawn_blocking(move || {
827            let _lease = lease;
828            let _permit = gov.acquire_tts(0, Some(&op_for_worker))?;
829            op_for_worker.check()?;
830            synthesize_with_pack(
831                _lease.value().as_ref(),
832                &text_owned,
833                &opts_owned,
834                &voice_internal,
835                &voice_canonical,
836                &model_canonical,
837                text_chars,
838                &op_for_worker,
839                &adapter,
840                TrustMode::Builtin,
841                "builtin",
842            )
843        });
844
845        // Pin the join handle so a soft deadline can detach a reaper without
846        // abandoning the permit while native work still runs (JOE-1600).
847        let join = join;
848        tokio::select! {
849            join_res = join => {
850                match join_res {
851                    Ok(result) => result,
852                    Err(e) => Err(crate::error::TranscriptionError::internal(format!(
853                        "TTS synth join: {e}"
854                    ))),
855                }
856            }
857            _ = tokio::time::sleep(timeout) => {
858                // Soft deadline: cooperative cancel so chunk loops exit when
859                // possible. The blocking task retains the ResourceGovernor permit
860                // until it returns; we cannot safely abort in-process ONNX.
861                op.cancel.cancel();
862                // Note: `join` was moved into select; when this arm wins, the
863                // JoinHandle is dropped. Tokio does not cancel spawn_blocking
864                // work on JoinHandle drop — the worker continues and drops the
865                // permit when the closure returns. That keeps concurrency honest
866                // (permits stay occupied) without a reaper task.
867                Err(ProviderError::DeadlineExceeded.into())
868            }
869        }
870    }
871
872    async fn preload(&self, model: &str, voice: &str) -> Result<()> {
873        // Same contract as synthesize: model-scoped voice validation first.
874        let (model_info, _) = resolve_voice_for_model(model, voice)?;
875        let pin = self
876            .ensure_loaded_pin(model_info.id, self.local_only)
877            .await?;
878        drop(pin);
879        Ok(())
880    }
881}
882
883/// Convenience: synthesize with a one-shot provider.
884pub async fn synthesize_local(
885    cache_dir: impl Into<PathBuf>,
886    text: &str,
887    opts: &SynthesisOptions,
888) -> Result<SynthesisResult> {
889    let provider = LocalTtsProvider::new(cache_dir.into()).with_progress(false);
890    provider.synthesize(text, opts).await
891}
892
893#[cfg(test)]
894mod tests {
895    use super::*;
896
897    #[tokio::test]
898    async fn empty_text_user_error() {
899        let dir = tempfile::tempdir().unwrap();
900        let p = LocalTtsProvider::new(dir.path().to_path_buf()).with_local_only(true);
901        let err = p
902            .synthesize("  ", &SynthesisOptions::default())
903            .await
904            .unwrap_err();
905        assert_eq!(err.exit_code(), 2);
906    }
907
908    #[tokio::test]
909    async fn missing_pack_local_only() {
910        let dir = tempfile::tempdir().unwrap();
911        let p = LocalTtsProvider::new(dir.path().to_path_buf()).with_local_only(true);
912        let err = p
913            .synthesize("Hello", &SynthesisOptions::default())
914            .await
915            .unwrap_err();
916        assert!(matches!(err.exit_code(), 2 | 4));
917    }
918
919    #[tokio::test]
920    async fn unsupported_language_user_error() {
921        let dir = tempfile::tempdir().unwrap();
922        let p = LocalTtsProvider::new(dir.path().to_path_buf()).with_local_only(true);
923        let opts = SynthesisOptions {
924            language: "fr".into(),
925            ..Default::default()
926        };
927        let err = p.synthesize("Hello", &opts).await.unwrap_err();
928        assert_eq!(err.exit_code(), 2);
929        assert!(err.to_string().contains("unsupported TTS language"));
930    }
931
932    #[tokio::test]
933    async fn non_native_sample_rate_rejected() {
934        let dir = tempfile::tempdir().unwrap();
935        let p = LocalTtsProvider::new(dir.path().to_path_buf()).with_local_only(true);
936        let opts = SynthesisOptions {
937            sample_rate_hz: Some(16_000),
938            ..Default::default()
939        };
940        let err = p.synthesize("Hello", &opts).await.unwrap_err();
941        assert_eq!(err.exit_code(), 2);
942        assert!(err.to_string().contains("sample rate"));
943    }
944
945    #[tokio::test]
946    async fn invalid_speaking_rate_rejected() {
947        let dir = tempfile::tempdir().unwrap();
948        let p = LocalTtsProvider::new(dir.path().to_path_buf()).with_local_only(true);
949        let opts = SynthesisOptions {
950            speaking_rate: 9.0,
951            ..Default::default()
952        };
953        let err = p.synthesize("Hello", &opts).await.unwrap_err();
954        assert_eq!(err.exit_code(), 2);
955    }
956
957    #[tokio::test]
958    async fn oversized_text_rejected_not_truncated() {
959        let dir = tempfile::tempdir().unwrap();
960        let p = LocalTtsProvider::new(dir.path().to_path_buf())
961            .with_local_only(true)
962            .with_max_chars(10);
963        let err = p
964            .synthesize(
965                "this text is definitely longer than ten",
966                &SynthesisOptions::default(),
967            )
968            .await
969            .unwrap_err();
970        assert_eq!(err.exit_code(), 2);
971        assert!(err.to_string().contains("too long"));
972    }
973
974    #[test]
975    fn clear_sessions_is_safe_when_empty() {
976        let dir = tempfile::tempdir().unwrap();
977        let p = LocalTtsProvider::new(dir.path().to_path_buf()).with_local_only(true);
978        p.clear_sessions();
979        p.clear_sessions();
980    }
981
982    /// Real Kokoro synthesis smoke (network-free with pre-seeded cache).
983    ///
984    /// ```text
985    /// AURUM_KOKORO_INTEGRATION=1 AURUM_TTS_CACHE=/path/to/cache \
986    ///   cargo test -p aurum-core --lib kokoro_real_synth -- --ignored
987    /// ```
988    #[tokio::test]
989    #[ignore]
990    async fn kokoro_real_synth_from_cache() {
991        if std::env::var("AURUM_KOKORO_INTEGRATION").ok().as_deref() != Some("1") {
992            return;
993        }
994        let cache = std::env::var("AURUM_TTS_CACHE")
995            .map(std::path::PathBuf::from)
996            .unwrap_or_else(|_| {
997                PathBuf::from(std::env::var("HOME").unwrap_or_else(|_| "/tmp".into()))
998                    .join(".cache/aurum")
999            });
1000        let p = LocalTtsProvider::new(cache).with_local_only(true);
1001        let opts = SynthesisOptions {
1002            model: crate::tts::catalogue::KOKORO_TTS_MODEL.into(),
1003            voice: crate::tts::catalogue::KOKORO_DEFAULT_VOICE.into(),
1004            ..Default::default()
1005        };
1006        let r = p
1007            .synthesize("Hello from Kokoro.", &opts)
1008            .await
1009            .expect("kokoro synth");
1010        assert_eq!(r.sample_rate_hz, 24_000);
1011        assert!(!r.pcm_i16_mono.is_empty());
1012        assert_eq!(r.adapter.as_deref(), Some(ADAPTER_KOKORO_ONNX_V0));
1013        assert!(r.duration_ms > 0);
1014    }
1015}