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