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, stage_verified_for_load};
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        let cache_root = self.cache_dir.clone();
274        let onnx_sha = Some(info.onnx.sha256.to_string());
275        let voices_sha = Some(info.voices.sha256.to_string());
276        let onnx_leaf = onnx
277            .file_name()
278            .and_then(|s| s.to_str())
279            .unwrap_or("model.onnx")
280            .to_string();
281        let voices_leaf = voices_file
282            .file_name()
283            .and_then(|s| s.to_str())
284            .unwrap_or("voices.npz")
285            .to_string();
286
287        let pin = tokio::task::spawn_blocking(move || {
288            pool.get_or_load_pin(key_for_load, weight, &gov, || {
289                let onnx_load =
290                    stage_verified_for_load(&onnx, onnx_sha.as_deref(), &cache_root, &onnx_leaf)?;
291                let voices_load = stage_verified_for_load(
292                    &voices_file,
293                    voices_sha.as_deref(),
294                    &cache_root,
295                    &voices_leaf,
296                )?;
297                load_pack(
298                    &onnx_load,
299                    &voices_load,
300                    sample_rate,
301                    speed_priors,
302                    catalogue_max,
303                    &model_id,
304                )
305            })
306        })
307        .await
308        .map_err(|e| crate::error::TranscriptionError::internal(format!("TTS load join: {e}")))??;
309
310        Ok(pin)
311    }
312
313    /// Load a verified (or explicitly unverified) local model pack for a supported
314    /// adapter (JOE-1619). Cache key is isolated from built-in catalogue identities.
315    async fn ensure_loaded_from_pack_pin(
316        &self,
317        pack_dir: &Path,
318        allow_unverified: bool,
319    ) -> Result<(RegistryPin<LoadedPack>, super::adapter::ModelPackManifest)> {
320        let (root, manifest) = load_pack_dir(pack_dir, allow_unverified)?;
321        if manifest.adapter_id != ADAPTER_KITTEN_ONNX_V1
322            && manifest.adapter_id != ADAPTER_KOKORO_ONNX_V0
323        {
324            return Err(UserError::InvalidConfig {
325                reason: format!(
326                    "local pack override synthesis currently supports adapters \
327                     '{ADAPTER_KITTEN_ONNX_V1}' and '{ADAPTER_KOKORO_ONNX_V0}' \
328                     (got '{}'); use `aurum tts inspect` / conformance for others",
329                    manifest.adapter_id
330                ),
331            }
332            .into());
333        }
334        let onnx_name = manifest
335            .artifact("onnx")
336            .map(|a| a.filename.clone())
337            .ok_or_else(|| UserError::InvalidConfig {
338                reason: "pack missing onnx artifact".into(),
339            })?;
340        let voices_name = manifest
341            .artifact("voices")
342            .map(|a| a.filename.clone())
343            .ok_or_else(|| UserError::InvalidConfig {
344                reason: "pack missing voices artifact".into(),
345            })?;
346        let config_name = manifest
347            .artifact("config")
348            .map(|a| a.filename.clone())
349            .unwrap_or_else(|| "config.json".into());
350        let onnx = root.join(&onnx_name);
351        let voices_file = root.join(&voices_name);
352        let onnx_sha = manifest.artifact("onnx").and_then(|a| a.sha256.clone());
353        let voices_sha = manifest.artifact("voices").and_then(|a| a.sha256.clone());
354        let sample_rate = manifest.sample_rate_hz;
355        let catalogue_max = manifest.max_phoneme_tokens;
356        let model_id = manifest.model_id.clone();
357        // Isolate local packs from built-in cache keys.
358        let key = LoadKey::tts(
359            format!("local-pack:{}", manifest.model_id),
360            root.display().to_string(),
361        );
362        if let Some(pin) = self.pool.registry.get_and_pin(&key) {
363            return Ok((pin, manifest));
364        }
365        let key_for_load = key.clone();
366        let speed_priors = load_speed_priors_from_path(&root.join(&config_name));
367        let weight = TtsSessionPool::weight_for(&onnx);
368        let pool = Arc::clone(&self.pool);
369        let gov = Arc::clone(&self.governor);
370        let cache_root = self.cache_dir.clone();
371
372        let pin = tokio::task::spawn_blocking(move || {
373            pool.get_or_load_pin(key_for_load, weight, &gov, || {
374                // JOE-1918 re-open: stage verified digests into private snap paths
375                // so ORT/NPZ open process-owned bytes, not the mutable pack path.
376                let onnx_load =
377                    stage_verified_for_load(&onnx, onnx_sha.as_deref(), &cache_root, &onnx_name)?;
378                let voices_load = stage_verified_for_load(
379                    &voices_file,
380                    voices_sha.as_deref(),
381                    &cache_root,
382                    &voices_name,
383                )?;
384                load_pack(
385                    &onnx_load,
386                    &voices_load,
387                    sample_rate,
388                    speed_priors,
389                    catalogue_max,
390                    &model_id,
391                )
392            })
393        })
394        .await
395        .map_err(|e| {
396            crate::error::TranscriptionError::internal(format!("TTS pack load join: {e}"))
397        })??;
398
399        Ok((pin, manifest))
400    }
401}
402
403#[allow(clippy::too_many_arguments)]
404fn synthesize_with_pack(
405    pack: &LoadedPack,
406    text: &str,
407    opts: &SynthesisOptions,
408    voice_internal: &str,
409    voice_canonical: &str,
410    model_canonical: &str,
411    text_chars: usize,
412    op: &OpContext,
413    adapter: &str,
414    trust: TrustMode,
415    provenance: &str,
416) -> Result<SynthesisResult> {
417    op.check()?;
418
419    let rate = validate_speaking_rate(opts.speaking_rate)?;
420    let sample_rate = resolve_sample_rate(opts.sample_rate_hz, pack.sample_rate_hz)?;
421
422    let voice_mat = pack
423        .voices
424        .get(voice_internal)
425        .ok_or_else(|| UserError::Other {
426            message: format!(
427                "voice embedding '{voice_internal}' missing from pack; available: {:?}",
428                pack.voices.keys().collect::<Vec<_>>()
429            ),
430        })?;
431
432    // Voice matrix has one style embedding per supported sequence length.
433    // Reserve one row because the token vector includes start/end pads.
434    let pack_max = voice_mat.nrows.saturating_sub(1);
435    let max_tokens = pack_max.min(pack.max_phoneme_tokens);
436    if max_tokens <= 2 {
437        return Err(ProviderError::ModelLoad {
438            model: model_canonical.to_string(),
439            reason: "voice embedding has no usable sequence rows".into(),
440        }
441        .into());
442    }
443
444    let codec = phoneme_codec_for_adapter(adapter);
445    let chunks = prepare_tts_chunks_with(text, max_tokens, codec)?;
446    let chunk_count = chunks.len();
447    let effective_speed = rate
448        * pack
449            .speed_priors
450            .get(voice_internal)
451            .copied()
452            .unwrap_or(1.0);
453    let pause_samples = (sample_rate as u64)
454        .saturating_mul(CHUNK_PAUSE_MS)
455        .checked_div(1_000)
456        .unwrap_or(0) as usize;
457
458    let mut pcm_f32: Vec<f32> = Vec::new();
459    let trim_policy = TailTrimPolicy::default();
460
461    for (index, chunk) in chunks.iter().enumerate() {
462        op.check()?;
463        let chunk_audio =
464            synthesize_chunk_f32(pack, voice_mat, chunk, effective_speed).map_err(|err| {
465                ProviderError::Other {
466                    message: format!(
467                        "TTS chunk {}/{} failed near {:?}: {err}",
468                        index + 1,
469                        chunks.len(),
470                        chunk.text.chars().take(80).collect::<String>()
471                    ),
472                }
473            })?;
474        let trimmed = trim_trailing_silence(&chunk_audio, trim_policy);
475        if index > 0 {
476            pcm_f32.resize(pcm_f32.len().saturating_add(pause_samples), 0.0);
477        }
478        pcm_f32.extend_from_slice(trimmed);
479    }
480
481    validate_raw_pcm(&pcm_f32, sample_rate)?;
482    let pcm = peak_guard_f32_to_i16(&pcm_f32, PEAK_LIMIT);
483    if pcm.is_empty() {
484        return Err(ProviderError::Other {
485            message: "synthesis produced empty audio after validation".into(),
486        }
487        .into());
488    }
489
490    let duration_ms = duration_ms_from_pcm(pcm.len(), sample_rate);
491    let language = normalize_tts_language(&opts.language).unwrap_or_else(|_| "en".into());
492
493    Ok(SynthesisResult {
494        pcm_i16_mono: pcm,
495        sample_rate_hz: sample_rate,
496        channels: 1,
497        backend_kind: BackendKind::Local,
498        provider: "local".into(),
499        model: model_canonical.to_string(),
500        voice: voice_canonical.to_string(),
501        language,
502        duration_ms,
503        text_chars,
504        text_truncated: false,
505        chunk_count,
506        synthesized_chars: text_chars,
507        adapter: Some(adapter.into()),
508        trust: Some(trust.as_str().into()),
509        provenance: Some(provenance.into()),
510    })
511}
512
513fn phoneme_codec_for_adapter(adapter: &str) -> PhonemeCodec {
514    if adapter == ADAPTER_KOKORO_ONNX_V0 {
515        PhonemeCodec::Kokoro
516    } else {
517        PhonemeCodec::Kitten
518    }
519}
520
521fn synthesize_chunk_f32(
522    pack: &LoadedPack,
523    voice_mat: &super::npz::VoiceMatrix,
524    chunk: &TtsChunk,
525    effective_speed: f32,
526) -> Result<Vec<f32>> {
527    let seq_len = chunk.ids.len();
528    // Kitten and Kokoro both index style embeddings by phoneme sequence length.
529    let style = voice_mat.style_row(seq_len).to_vec();
530    let style_dim = style.len();
531    let t_ids = Tensor::<i64>::from_array(([1usize, seq_len], chunk.ids.clone())).map_err(|e| {
532        ProviderError::Other {
533            message: format!("tokens/input_ids tensor: {e}"),
534        }
535    })?;
536    let t_style = Tensor::<f32>::from_array(([1usize, style_dim], style)).map_err(|e| {
537        ProviderError::Other {
538            message: format!("style tensor: {e}"),
539        }
540    })?;
541    let t_speed = Tensor::<f32>::from_array(([1usize], vec![effective_speed])).map_err(|e| {
542        ProviderError::Other {
543            message: format!("speed tensor: {e}"),
544        }
545    })?;
546
547    let mut session = pack
548        .session
549        .lock()
550        .map_err(|_| crate::error::TranscriptionError::internal("ORT session mutex poisoned"))?;
551    // Positional order matches both Kitten and Kokoro graphs: tokens, style, speed.
552    let outputs = session
553        .run(ort::inputs![t_ids, t_style, t_speed])
554        .map_err(|e| ProviderError::Other {
555            message: format!("ONNX inference failed: {e}"),
556        })?;
557    let (_shape, audio_data) =
558        outputs[0]
559            .try_extract_tensor::<f32>()
560            .map_err(|e| ProviderError::Other {
561                message: format!("extract audio tensor: {e}"),
562            })?;
563    Ok(audio_data.to_vec())
564}
565
566fn load_pack(
567    onnx: &Path,
568    voices_file: &Path,
569    sample_rate_hz: u32,
570    speed_priors: HashMap<String, f32>,
571    catalogue_max_tokens: usize,
572    model_label: &str,
573) -> Result<LoadedPack> {
574    let session = Session::builder()
575        .map_err(|e| ProviderError::ModelLoad {
576            model: model_label.to_string(),
577            reason: format!("ORT session builder: {e}"),
578        })?
579        .commit_from_file(onnx)
580        .map_err(|e| ProviderError::ModelLoad {
581            model: model_label.to_string(),
582            reason: format!("load ONNX: {e}"),
583        })?;
584    let voices = load_voices_npz(voices_file)?;
585    // Derive capacity from the first voice matrix (all Kitten voices share shape).
586    let matrix_max = voices
587        .values()
588        .map(|m| m.nrows.saturating_sub(1))
589        .min()
590        .unwrap_or(0);
591    let max_phoneme_tokens = if matrix_max == 0 {
592        catalogue_max_tokens
593    } else {
594        matrix_max.min(catalogue_max_tokens)
595    };
596    Ok(LoadedPack {
597        session: Mutex::new(session),
598        voices,
599        sample_rate_hz,
600        max_phoneme_tokens,
601        speed_priors,
602    })
603}
604
605fn load_speed_priors(
606    cache_dir: &Path,
607    info: &super::catalogue::TtsModelInfo,
608) -> HashMap<String, f32> {
609    load_speed_priors_from_path(&super::catalogue::config_path(cache_dir, info))
610}
611
612fn load_speed_priors_from_path(path: &Path) -> HashMap<String, f32> {
613    let Ok(bytes) = std::fs::read(path) else {
614        return HashMap::new();
615    };
616    let Ok(v) = serde_json::from_slice::<serde_json::Value>(&bytes) else {
617        return HashMap::new();
618    };
619    let mut out = HashMap::new();
620    if let Some(map) = v.get("speed_priors").and_then(|x| x.as_object()) {
621        for (k, val) in map {
622            if let Some(f) = val.as_f64() {
623                out.insert(k.clone(), f as f32);
624            }
625        }
626    }
627    out
628}
629
630/// Synthesize via the fake-sine adapter (conformance / custom packs; no ONNX).
631fn synthesize_fake_adapter(
632    opts: &SynthesisOptions,
633    model_id: &str,
634    trust: TrustMode,
635    provenance: &str,
636    text_chars: usize,
637) -> Result<SynthesisResult> {
638    // Map text length to a short bounded duration (deterministic, no network).
639    let duration_ms = ((text_chars as u64).saturating_mul(40)).clamp(50, 2_000);
640    let pcm_f32 = synthesize_fake_sine_ms(duration_ms).map_err(|e| ProviderError::Other {
641        message: format!("fake-sine synth: {e}"),
642    })?;
643    validate_raw_pcm(&pcm_f32, 24_000)?;
644    let pcm = peak_guard_f32_to_i16(&pcm_f32, PEAK_LIMIT);
645    let sample_rate = resolve_sample_rate(opts.sample_rate_hz, 24_000)?;
646    let language = normalize_tts_language(&opts.language).unwrap_or_else(|_| "en".into());
647    let voice = if opts.voice.trim().is_empty() {
648        "Tone".into()
649    } else {
650        opts.voice.clone()
651    };
652    let out_duration_ms = duration_ms_from_pcm(pcm.len(), sample_rate);
653    Ok(SynthesisResult {
654        pcm_i16_mono: pcm,
655        sample_rate_hz: sample_rate,
656        channels: 1,
657        backend_kind: BackendKind::Local,
658        provider: "local".into(),
659        model: model_id.into(),
660        voice,
661        language,
662        duration_ms: out_duration_ms,
663        text_chars,
664        text_truncated: false,
665        chunk_count: 1,
666        synthesized_chars: text_chars,
667        adapter: Some(ADAPTER_FAKE_SINE_V1.into()),
668        trust: Some(trust.as_str().into()),
669        provenance: Some(provenance.into()),
670    })
671}
672
673/// Resolve a voice for a local pack: prefer matching catalogue model voices, else pack manifest.
674fn resolve_pack_voice(
675    manifest: &super::adapter::ModelPackManifest,
676    requested: &str,
677) -> Result<(String, String)> {
678    let req = requested.trim();
679    let default_voice = if manifest.adapter_id == ADAPTER_KOKORO_ONNX_V0 {
680        super::catalogue::KOKORO_DEFAULT_VOICE
681    } else {
682        super::catalogue::DEFAULT_TTS_VOICE
683    };
684    let req = if req.is_empty() { default_voice } else { req };
685    if let Ok((_, v)) = resolve_voice_for_model(&manifest.model_id, req) {
686        return Ok((v.id.to_string(), v.internal_key.to_string()));
687    }
688    if let Ok((_, v)) = resolve_voice_for_model(super::catalogue::DEFAULT_TTS_MODEL, req) {
689        return Ok((v.id.to_string(), v.internal_key.to_string()));
690    }
691    if let Some(v) = manifest
692        .voices
693        .iter()
694        .find(|v| v.id.eq_ignore_ascii_case(req) || v.internal_key.eq_ignore_ascii_case(req))
695    {
696        return Ok((v.id.clone(), v.internal_key.clone()));
697    }
698    let available: Vec<_> = manifest.voices.iter().map(|v| v.id.as_str()).collect();
699    Err(UserError::Other {
700        message: format!(
701            "voice '{req}' not found in pack '{}'; available: {available:?}",
702            manifest.model_id
703        ),
704    }
705    .into())
706}
707
708#[async_trait]
709impl SynthesisProvider for LocalTtsProvider {
710    fn name(&self) -> &'static str {
711        "local"
712    }
713
714    async fn synthesize(&self, text: &str, opts: &SynthesisOptions) -> Result<SynthesisResult> {
715        let prepared = prepare_text(text, self.max_chars)?;
716        let mut opts = opts.clone();
717        opts.language = normalize_tts_language(&opts.language)?;
718
719        // Local pack override path (JOE-1619): never hits network, never shadows
720        // built-in cache identity. Bare ONNX is rejected by load_pack_dir.
721        if let Some(pack_dir) = opts.pack_dir.clone() {
722            let (root, manifest) = load_pack_dir(&pack_dir, opts.allow_unverified)?;
723            let _ = root;
724            if manifest.adapter_id == ADAPTER_FAKE_SINE_V1 {
725                let trust = manifest.trust;
726                // Prefer pack identity over the built-in default when the caller
727                // did not name a distinct custom model id.
728                let model_id = if opts.model.trim().is_empty()
729                    || opts.model == super::catalogue::DEFAULT_TTS_MODEL
730                {
731                    manifest.model_id.clone()
732                } else {
733                    opts.model.clone()
734                };
735                if opts.voice.trim().is_empty() || opts.voice == super::catalogue::DEFAULT_TTS_VOICE
736                {
737                    if let Some(v) = manifest.voices.first() {
738                        opts.voice = v.id.clone();
739                    } else {
740                        opts.voice = "Tone".into();
741                    }
742                }
743                return synthesize_fake_adapter(
744                    &opts,
745                    &model_id,
746                    trust,
747                    "local_pack",
748                    prepared.text_chars,
749                );
750            }
751            if manifest.adapter_id != ADAPTER_KITTEN_ONNX_V1
752                && manifest.adapter_id != ADAPTER_KOKORO_ONNX_V0
753            {
754                return Err(UserError::UnsupportedCapability {
755                    provider: "tts".into(),
756                    model: manifest.model_id,
757                    reason: format!(
758                        "adapter '{}' is not enabled for local pack synthesis",
759                        manifest.adapter_id
760                    ),
761                    hint:
762                        "use kitten-onnx-v1, kokoro-onnx-v0, or fake-sine-v1 packs; see `aurum tts adapters`"
763                            .into(),
764                }
765                .into());
766            }
767            // Prefer pack-declared model id for honesty when caller left default.
768            let model_canonical = if opts.model == super::catalogue::DEFAULT_TTS_MODEL
769                || opts.model.trim().is_empty()
770            {
771                manifest.model_id.clone()
772            } else {
773                opts.model.clone()
774            };
775            let (voice_canonical, voice_internal) = resolve_pack_voice(&manifest, &opts.voice)?;
776            validate_speaking_rate(opts.speaking_rate)?;
777            resolve_sample_rate(opts.sample_rate_hz, manifest.sample_rate_hz)?;
778            let trust = manifest.trust;
779            let adapter = manifest.adapter_id.clone();
780            // Registry pin held for the full synthesis (JOE-1646).
781            let lease = self
782                .ensure_loaded_from_pack_pin(&pack_dir, opts.allow_unverified)
783                .await?
784                .0;
785
786            let timeout = Duration::from_millis(if opts.timeout_ms == 0 {
787                super::validate::DEFAULT_TIMEOUT_MS
788            } else {
789                opts.timeout_ms
790            });
791            // Same soft-deadline contract as built-in path (JOE-1830): outer
792            // timeout returns DeadlineExceeded while the blocking job retains
793            // the permit/pin until native work returns.
794            let op = OpContext::from_optional_cancel(opts.cancel.clone())
795                .with_deadline_from_now(timeout);
796            op.check()?;
797            op.emit("tts", "pack_synth");
798            let text_owned = prepared.text.clone();
799            let text_chars = prepared.text_chars;
800            let opts_owned = opts.clone();
801            let op_for_worker = op.clone();
802            let gov = Arc::clone(&self.governor);
803            let join = tokio::task::spawn_blocking(move || {
804                // Move registry pin into the worker so soft outer deadlines cannot
805                // release residency while ONNX is still running (JOE-1646).
806                let _lease = lease;
807                let _permit = gov.acquire_tts(0, Some(&op_for_worker))?;
808                op_for_worker.check()?;
809                synthesize_with_pack(
810                    _lease.value().as_ref(),
811                    &text_owned,
812                    &opts_owned,
813                    &voice_internal,
814                    &voice_canonical,
815                    &model_canonical,
816                    text_chars,
817                    &op_for_worker,
818                    &adapter,
819                    trust,
820                    "local_pack",
821                )
822            });
823
824            return tokio::select! {
825                join_res = join => {
826                    match join_res {
827                        Ok(result) => result,
828                        Err(e) => Err(crate::error::TranscriptionError::internal(format!(
829                            "TTS pack synth join: {e}"
830                        ))),
831                    }
832                }
833                _ = tokio::time::sleep(timeout) => {
834                    // Soft deadline: cooperative cancel so chunk loops exit when
835                    // possible. JoinHandle drop does not abort spawn_blocking;
836                    // the worker keeps the permit until the closure returns.
837                    op.cancel.cancel();
838                    Err(ProviderError::DeadlineExceeded.into())
839                }
840            };
841        }
842
843        let (model_info, voice_info) = resolve_voice_for_model(&opts.model, &opts.voice)?;
844        // Canonical IDs for honesty JSON / metadata.
845        opts.model = model_info.id.to_string();
846        opts.voice = voice_info.id.to_string();
847        validate_speaking_rate(opts.speaking_rate)?;
848        resolve_sample_rate(opts.sample_rate_hz, model_info.sample_rate_hz)?;
849
850        let local_only = opts.local_only || self.local_only;
851        let lease = self.ensure_loaded_pin(&opts.model, local_only).await?;
852
853        let timeout = Duration::from_millis(if opts.timeout_ms == 0 {
854            super::validate::DEFAULT_TIMEOUT_MS
855        } else {
856            opts.timeout_ms
857        });
858        let op =
859            OpContext::from_optional_cancel(opts.cancel.clone()).with_deadline_from_now(timeout);
860        op.check()?;
861        op.emit("tts", "builtin_synth");
862
863        let text_owned = prepared.text.clone();
864        let text_chars = prepared.text_chars;
865        let opts_owned = opts.clone();
866        let voice_internal = voice_info.internal_key.to_string();
867        let voice_canonical = voice_info.id.to_string();
868        let model_canonical = model_info.id.to_string();
869        let op_for_worker = op.clone();
870        let adapter = model_info.adapter.to_string();
871        let gov = Arc::clone(&self.governor);
872
873        // Hold the TTS + blocking permit and registry pin for the entire native
874        // job lifetime — even if the caller soft-times out (JOE-1600 / JOE-1646 /
875        // JOE-1830 uniform soft-deadline contract).
876        let join = tokio::task::spawn_blocking(move || {
877            let _lease = lease;
878            let _permit = gov.acquire_tts(0, Some(&op_for_worker))?;
879            op_for_worker.check()?;
880            synthesize_with_pack(
881                _lease.value().as_ref(),
882                &text_owned,
883                &opts_owned,
884                &voice_internal,
885                &voice_canonical,
886                &model_canonical,
887                text_chars,
888                &op_for_worker,
889                &adapter,
890                TrustMode::Builtin,
891                "builtin",
892            )
893        });
894
895        // Soft deadline: caller timeout is distinguished from still-running
896        // native work. JoinHandle drop does not abort spawn_blocking; the
897        // worker retains the permit until return (JOE-1600 / JOE-1830).
898        tokio::select! {
899            join_res = join => {
900                match join_res {
901                    Ok(result) => result,
902                    Err(e) => Err(crate::error::TranscriptionError::internal(format!(
903                        "TTS synth join: {e}"
904                    ))),
905                }
906            }
907            _ = tokio::time::sleep(timeout) => {
908                op.cancel.cancel();
909                Err(ProviderError::DeadlineExceeded.into())
910            }
911        }
912    }
913
914    async fn preload(&self, model: &str, voice: &str) -> Result<()> {
915        // Same contract as synthesize: model-scoped voice validation first.
916        let (model_info, _) = resolve_voice_for_model(model, voice)?;
917        let pin = self
918            .ensure_loaded_pin(model_info.id, self.local_only)
919            .await?;
920        drop(pin);
921        Ok(())
922    }
923}
924
925/// Convenience: synthesize with a one-shot provider.
926pub async fn synthesize_local(
927    cache_dir: impl Into<PathBuf>,
928    text: &str,
929    opts: &SynthesisOptions,
930) -> Result<SynthesisResult> {
931    let provider = LocalTtsProvider::new(cache_dir.into()).with_progress(false);
932    provider.synthesize(text, opts).await
933}
934
935#[cfg(test)]
936mod tests {
937    use super::*;
938
939    #[tokio::test]
940    async fn empty_text_user_error() {
941        let dir = tempfile::tempdir().unwrap();
942        let p = LocalTtsProvider::new(dir.path().to_path_buf()).with_local_only(true);
943        let err = p
944            .synthesize("  ", &SynthesisOptions::default())
945            .await
946            .unwrap_err();
947        assert_eq!(err.exit_code(), 2);
948    }
949
950    #[tokio::test]
951    async fn missing_pack_local_only() {
952        let dir = tempfile::tempdir().unwrap();
953        let p = LocalTtsProvider::new(dir.path().to_path_buf()).with_local_only(true);
954        let err = p
955            .synthesize("Hello", &SynthesisOptions::default())
956            .await
957            .unwrap_err();
958        assert!(matches!(err.exit_code(), 2 | 4));
959    }
960
961    #[tokio::test]
962    async fn unsupported_language_user_error() {
963        let dir = tempfile::tempdir().unwrap();
964        let p = LocalTtsProvider::new(dir.path().to_path_buf()).with_local_only(true);
965        let opts = SynthesisOptions {
966            language: "fr".into(),
967            ..Default::default()
968        };
969        let err = p.synthesize("Hello", &opts).await.unwrap_err();
970        assert_eq!(err.exit_code(), 2);
971        assert!(err.to_string().contains("unsupported TTS language"));
972    }
973
974    #[tokio::test]
975    async fn non_native_sample_rate_rejected() {
976        let dir = tempfile::tempdir().unwrap();
977        let p = LocalTtsProvider::new(dir.path().to_path_buf()).with_local_only(true);
978        let opts = SynthesisOptions {
979            sample_rate_hz: Some(16_000),
980            ..Default::default()
981        };
982        let err = p.synthesize("Hello", &opts).await.unwrap_err();
983        assert_eq!(err.exit_code(), 2);
984        assert!(err.to_string().contains("sample rate"));
985    }
986
987    #[tokio::test]
988    async fn invalid_speaking_rate_rejected() {
989        let dir = tempfile::tempdir().unwrap();
990        let p = LocalTtsProvider::new(dir.path().to_path_buf()).with_local_only(true);
991        let opts = SynthesisOptions {
992            speaking_rate: 9.0,
993            ..Default::default()
994        };
995        let err = p.synthesize("Hello", &opts).await.unwrap_err();
996        assert_eq!(err.exit_code(), 2);
997    }
998
999    #[tokio::test]
1000    async fn oversized_text_rejected_not_truncated() {
1001        let dir = tempfile::tempdir().unwrap();
1002        let p = LocalTtsProvider::new(dir.path().to_path_buf())
1003            .with_local_only(true)
1004            .with_max_chars(10);
1005        let err = p
1006            .synthesize(
1007                "this text is definitely longer than ten",
1008                &SynthesisOptions::default(),
1009            )
1010            .await
1011            .unwrap_err();
1012        assert_eq!(err.exit_code(), 2);
1013        assert!(err.to_string().contains("too long"));
1014    }
1015
1016    #[test]
1017    fn clear_sessions_is_safe_when_empty() {
1018        let dir = tempfile::tempdir().unwrap();
1019        let p = LocalTtsProvider::new(dir.path().to_path_buf()).with_local_only(true);
1020        p.clear_sessions();
1021        p.clear_sessions();
1022    }
1023
1024    /// Real Kokoro synthesis smoke (network-free with pre-seeded cache).
1025    ///
1026    /// ```text
1027    /// AURUM_KOKORO_INTEGRATION=1 AURUM_TTS_CACHE=/path/to/cache \
1028    ///   cargo test -p aurum-core --lib kokoro_real_synth -- --ignored
1029    /// ```
1030    #[tokio::test]
1031    #[ignore]
1032    async fn kokoro_real_synth_from_cache() {
1033        if std::env::var("AURUM_KOKORO_INTEGRATION").ok().as_deref() != Some("1") {
1034            return;
1035        }
1036        let cache = std::env::var("AURUM_TTS_CACHE")
1037            .map(std::path::PathBuf::from)
1038            .unwrap_or_else(|_| {
1039                PathBuf::from(std::env::var("HOME").unwrap_or_else(|_| "/tmp".into()))
1040                    .join(".cache/aurum")
1041            });
1042        let p = LocalTtsProvider::new(cache).with_local_only(true);
1043        let opts = SynthesisOptions {
1044            model: crate::tts::catalogue::KOKORO_TTS_MODEL.into(),
1045            voice: crate::tts::catalogue::KOKORO_DEFAULT_VOICE.into(),
1046            ..Default::default()
1047        };
1048        let r = p
1049            .synthesize("Hello from Kokoro.", &opts)
1050            .await
1051            .expect("kokoro synth");
1052        assert_eq!(r.sample_rate_hz, 24_000);
1053        assert!(!r.pcm_i16_mono.is_empty());
1054        assert_eq!(r.adapter.as_deref(), Some(ADAPTER_KOKORO_ONNX_V0));
1055        assert!(r.duration_ms > 0);
1056    }
1057}