aurum_core/providers/
local.rs

1//! Local transcription via whisper.cpp (through the `whisper-rs` bindings).
2//!
3//! Maintains a cache of loaded `WhisperContext` values keyed by model path so
4//! repeated calls do not reload ggml. By default providers share a
5//! **process-global** pool; [`AurumEngine`](crate::AurumEngine) owns an
6//! isolated pool per engine (JOE-1784).
7//!
8//! Concurrent cold starts coalesce via singleflight (JOE-1597). Resident weight
9//! and eviction use the model registry (JOE-1598). Inference threads and
10//! blocking work are admitted by a [`ResourceGovernor`] (JOE-1596/1599).
11//!
12//! Embedders (mic hosts) should prefer:
13//! - [`LocalWhisperProvider::preload`] at startup
14//! - [`LocalWhisperProvider::transcribe_pcm`] or [`AudioInput::from_pcm`] for buffers
15//! - [`crate::pcm::PcmBuffer`] to accumulate capture chunks
16
17use super::{
18    BackendKind, Segment, TranscriptionOptions, TranscriptionProvider, TranscriptionResult,
19};
20use crate::audio::{AudioInput, WHISPER_SAMPLE_RATE};
21use crate::error::{ProviderError, Result};
22use crate::model::{self, DownloadProgressCallback, EnsureModelOptions};
23use crate::postprocess;
24use crate::runtime::{
25    LoadKey, ModelRegistry, OpContext, PermitKind, RegistryConfig, RegistryPin, ResidencyWeight,
26    ResourceGovernor, Singleflight,
27};
28use async_trait::async_trait;
29use once_cell::sync::Lazy;
30use std::path::{Path, PathBuf};
31use std::sync::Arc;
32use whisper_rs::{FullParams, SamplingStrategy, WhisperContext, WhisperContextParameters};
33
34static LOGGING_HOOKS: once_cell::sync::OnceCell<()> = once_cell::sync::OnceCell::new();
35
36/// STT context residency + singleflight loader (JOE-1784).
37///
38/// Own one per [`crate::AurumEngine`], or share via [`process_global_stt_pool`].
39pub struct SttContextPool {
40    flight: Singleflight<WhisperContext>,
41    registry: ModelRegistry<WhisperContext>,
42}
43
44impl Default for SttContextPool {
45    fn default() -> Self {
46        Self::new()
47    }
48}
49
50impl SttContextPool {
51    pub fn new() -> Self {
52        Self {
53            flight: Singleflight::default(),
54            registry: ModelRegistry::new(RegistryConfig::default()),
55        }
56    }
57
58    /// Number of resident (loaded) contexts in this pool.
59    pub fn resident_len(&self) -> usize {
60        self.registry.len()
61    }
62
63    fn load_key(model_path: &Path, model_name: &str) -> LoadKey {
64        LoadKey::stt(model_name, model_path.display().to_string())
65    }
66
67    /// Conservative weight: on-disk size + 50% runtime overhead headroom.
68    fn weight_for(model_path: &Path) -> ResidencyWeight {
69        let disk = std::fs::metadata(model_path).map(|m| m.len()).unwrap_or(0);
70        let bytes = disk.saturating_add(disk / 2).max(16 * 1024 * 1024);
71        ResidencyWeight { bytes }
72    }
73
74    /// Load or reuse a context and return an **active registry pin** held for the
75    /// whole operation (JOE-1646). No unpinned Arc is returned to callers.
76    pub fn get_or_load_pin(
77        &self,
78        model_path: &Path,
79        model_name: &str,
80        gov: &Arc<ResourceGovernor>,
81    ) -> Result<RegistryPin<WhisperContext>> {
82        let key = Self::load_key(model_path, model_name);
83        // Atomic lookup+pin under registry lock.
84        if let Some(pin) = self.registry.get_and_pin(&key) {
85            return Ok(pin);
86        }
87
88        let weight = Self::weight_for(model_path);
89        if weight.bytes > self.registry.config().max_resident_bytes {
90            return Err(ProviderError::Overload {
91                reason: format!(
92                    "model weight {} exceeds residency budget {}",
93                    weight.bytes,
94                    self.registry.config().max_resident_bytes
95                ),
96            }
97            .into());
98        }
99
100        // Coalesce concurrent cold loads; LeaderGuard is panic-safe.
101        loop {
102            if let Some(pin) = self.registry.get_and_pin(&key) {
103                return Ok(pin);
104            }
105            match self.flight.begin_or_wait_guard(key.clone()) {
106                Ok(None) => {
107                    // Leader finished — re-check registry.
108                    continue;
109                }
110                Err(message) => {
111                    return Err(ProviderError::ModelLoad {
112                        model: model_name.to_string(),
113                        reason: message,
114                    }
115                    .into());
116                }
117                Ok(Some(leader)) => {
118                    let path_owned = model_path.to_path_buf();
119                    let name_owned = model_name.to_string();
120                    let gov = Arc::clone(gov);
121
122                    let load_result = (|| -> Result<Arc<WhisperContext>> {
123                        let _permit = gov.acquire(PermitKind::ModelLoad, None)?;
124                        LOGGING_HOOKS.get_or_init(|| {
125                            whisper_rs::install_logging_hooks();
126                        });
127                        let params = WhisperContextParameters::default();
128                        let ctx = WhisperContext::new_with_params(
129                            path_owned.to_string_lossy().as_ref(),
130                            params,
131                        )
132                        .map_err(|e| ProviderError::ModelLoad {
133                            model: name_owned.clone(),
134                            reason: e.to_string(),
135                        })?;
136                        Ok(Arc::new(ctx))
137                    })();
138
139                    match load_result {
140                        Ok(ctx) => {
141                            match self.registry.insert_and_pin(
142                                key.clone(),
143                                Arc::clone(&ctx),
144                                weight,
145                            ) {
146                                Ok(pin) => {
147                                    leader.success();
148                                    return Ok(pin);
149                                }
150                                Err(e) => {
151                                    leader.fail(e.to_string());
152                                    drop(ctx);
153                                    return Err(e);
154                                }
155                            }
156                        }
157                        Err(e) => {
158                            leader.fail(e.to_string());
159                            return Err(e);
160                        }
161                    }
162                }
163            }
164        }
165    }
166
167    pub fn get_or_load(
168        &self,
169        model_path: &Path,
170        model_name: &str,
171        gov: &Arc<ResourceGovernor>,
172    ) -> Result<Arc<WhisperContext>> {
173        Ok(Arc::clone(
174            self.get_or_load_pin(model_path, model_name, gov)?.value(),
175        ))
176    }
177
178    /// Drop idle STT contexts only; active decode pins keep residency (JOE-1646).
179    pub fn clear(&self) {
180        let _ = self.registry.clear_idle();
181    }
182
183    pub fn contains(&self, model_path: &Path, model_name: &str) -> bool {
184        let key = Self::load_key(model_path, model_name);
185        self.registry.get(&key).is_some()
186    }
187}
188
189static PROCESS_STT_POOL: Lazy<Arc<SttContextPool>> = Lazy::new(|| Arc::new(SttContextPool::new()));
190
191/// Process-global STT pool used by default providers / CLI (JOE-1784).
192pub fn process_global_stt_pool() -> Arc<SttContextPool> {
193    Arc::clone(&PROCESS_STT_POOL)
194}
195
196/// Drop idle contexts in the **process-global** STT pool (call before exit when idle).
197///
198/// Prefer [`crate::AurumEngine::clear_model_caches`] for engine-owned pools.
199pub fn clear_context_cache() {
200    PROCESS_STT_POOL.clear();
201}
202
203/// Local whisper.cpp provider.
204pub struct LocalWhisperProvider {
205    cache_dir: PathBuf,
206    show_progress: bool,
207    /// When true, never download models — fail if missing from cache.
208    local_only: bool,
209    on_download_progress: Option<DownloadProgressCallback>,
210    /// Model residency pool (process-global by default; engine-local when injected).
211    pool: Arc<SttContextPool>,
212    /// Admission governor for load/STT permits.
213    governor: Arc<ResourceGovernor>,
214}
215
216impl LocalWhisperProvider {
217    pub fn new(cache_dir: PathBuf) -> Self {
218        Self::with_runtime(
219            cache_dir,
220            process_global_stt_pool(),
221            ResourceGovernor::process_global(),
222        )
223    }
224
225    /// Construct with an explicit STT pool and governor (JOE-1784 / engine path).
226    pub fn with_runtime(
227        cache_dir: PathBuf,
228        pool: Arc<SttContextPool>,
229        governor: Arc<ResourceGovernor>,
230    ) -> Self {
231        Self {
232            cache_dir,
233            show_progress: true,
234            local_only: false,
235            on_download_progress: None,
236            pool,
237            governor,
238        }
239    }
240
241    pub fn with_progress(mut self, show: bool) -> Self {
242        self.show_progress = show;
243        self
244    }
245
246    /// Fail closed if the model is not already on disk (no network).
247    pub fn with_local_only(mut self, local_only: bool) -> Self {
248        self.local_only = local_only;
249        self
250    }
251
252    /// Structured download progress for UI hosts (in addition to optional CLI bar).
253    pub fn with_download_progress(mut self, cb: DownloadProgressCallback) -> Self {
254        self.on_download_progress = Some(cb);
255        self
256    }
257
258    pub fn cache_dir(&self) -> &Path {
259        &self.cache_dir
260    }
261
262    pub fn local_only(&self) -> bool {
263        self.local_only
264    }
265
266    /// Whether the ggml file is present and passes basic integrity checks.
267    pub fn is_model_cached(&self, model: &str) -> bool {
268        model::is_model_cached(&self.cache_dir, model)
269    }
270
271    /// Whether a WhisperContext for this model is already loaded in this provider's pool.
272    pub fn is_model_loaded(&self, model: &str) -> bool {
273        let Ok(info) = model::lookup_model(model) else {
274            return false;
275        };
276        let path = model::model_path(&self.cache_dir, info);
277        self.pool.contains(&path, model)
278    }
279
280    pub fn pool(&self) -> &Arc<SttContextPool> {
281        &self.pool
282    }
283
284    pub fn governor(&self) -> &Arc<ResourceGovernor> {
285        &self.governor
286    }
287
288    fn ensure_opts(&self) -> EnsureModelOptions {
289        let mut opts = EnsureModelOptions::new()
290            .local_only(self.local_only)
291            .show_progress(self.show_progress);
292        if let Some(cb) = &self.on_download_progress {
293            opts = opts.on_progress(Arc::clone(cb));
294        }
295        opts
296    }
297
298    /// Download (unless `local_only`) and load the model into this provider's pool.
299    pub async fn preload(&self, model: &str) -> Result<PathBuf> {
300        let path =
301            model::ensure_model_with_options(&self.cache_dir, model, self.ensure_opts()).await?;
302        let model_name = model.to_string();
303        let path_clone = path.clone();
304        let pool = Arc::clone(&self.pool);
305        let gov = Arc::clone(&self.governor);
306        tokio::task::spawn_blocking(move || pool.get_or_load(&path_clone, &model_name, &gov))
307            .await
308            .map_err(|e| ProviderError::TranscriptionFailed {
309                reason: format!("preload worker panicked: {e}"),
310            })??;
311        Ok(path)
312    }
313
314    /// Transcribe a mono PCM buffer (must be [`WHISPER_SAMPLE_RATE`] Hz).
315    ///
316    /// Preferred entry point for mic hosts — no temp files, no ffmpeg.
317    pub async fn transcribe_pcm(
318        &self,
319        samples: &[f32],
320        options: &TranscriptionOptions,
321    ) -> Result<TranscriptionResult> {
322        let input = AudioInput::from_pcm_slice(samples, WHISPER_SAMPLE_RATE)?;
323        self.transcribe(&input, options).await
324    }
325
326    /// Transcribe an owned/shared PCM buffer already wrapped as [`AudioInput`].
327    pub async fn transcribe_audio(
328        &self,
329        input: &AudioInput,
330        options: &TranscriptionOptions,
331    ) -> Result<TranscriptionResult> {
332        self.transcribe(input, options).await
333    }
334}
335
336#[async_trait]
337impl TranscriptionProvider for LocalWhisperProvider {
338    fn name(&self) -> &'static str {
339        "local"
340    }
341
342    fn backend_kind(&self) -> BackendKind {
343        BackendKind::Asr
344    }
345
346    async fn transcribe(
347        &self,
348        input: &AudioInput,
349        options: &TranscriptionOptions,
350    ) -> Result<TranscriptionResult> {
351        if input.sample_rate() != WHISPER_SAMPLE_RATE {
352            return Err(crate::error::UserError::UnsupportedSampleRate {
353                got: input.sample_rate(),
354                need: WHISPER_SAMPLE_RATE,
355            }
356            .into());
357        }
358
359        let op = OpContext::from_optional_cancel(options.cancel.clone());
360        op.check()?;
361        op.emit("stt", "ensure_model");
362
363        let model_name = options.model.clone();
364        let language = options.language.clone();
365        let timestamps = options.timestamps;
366        let cancel = op.cancel.clone();
367        let samples: Arc<[f32]> = Arc::clone(input.samples());
368        let duration_secs = input.duration_secs();
369        let cache_dir = self.cache_dir.clone();
370        let opts = self.ensure_opts();
371
372        let model_path = model::ensure_model_with_options(&cache_dir, &model_name, opts).await?;
373
374        op.check()?;
375        op.emit("stt", "inference");
376
377        let pool = Arc::clone(&self.pool);
378        let gov = Arc::clone(&self.governor);
379        let result = tokio::task::spawn_blocking(move || {
380            // Admission inside the worker so queue wait does not occupy a Tokio worker.
381            let n_threads = gov.recommend_stt_threads();
382            let op_inner = OpContext::with_cancel(cancel.clone());
383            let _permit = gov.acquire_stt(n_threads, 0, Some(&op_inner))?;
384            // Re-check after queue wait so late work does not start after cancel.
385            op_inner.check()?;
386            run_whisper(
387                &pool,
388                &gov,
389                &model_path,
390                &model_name,
391                &samples,
392                duration_secs,
393                &language,
394                timestamps,
395                Some(&cancel),
396                n_threads as i32,
397            )
398        })
399        .await
400        .map_err(|e| ProviderError::TranscriptionFailed {
401            reason: format!("worker thread panicked: {e}"),
402        })??;
403
404        Ok(postprocess::normalize_result(result))
405    }
406}
407
408unsafe extern "C" fn abort_callback_atomic(user_data: *mut std::ffi::c_void) -> bool {
409    if user_data.is_null() {
410        return false;
411    }
412    let flag = &*(user_data as *const std::sync::atomic::AtomicBool);
413    flag.load(std::sync::atomic::Ordering::SeqCst)
414}
415
416#[allow(clippy::too_many_arguments)]
417fn run_whisper(
418    pool: &SttContextPool,
419    gov: &Arc<ResourceGovernor>,
420    model_path: &Path,
421    model_name: &str,
422    samples: &[f32],
423    duration_secs: f64,
424    language: &str,
425    timestamps: bool,
426    cancel: Option<&crate::cancel::CancelFlag>,
427    n_threads: i32,
428) -> Result<TranscriptionResult> {
429    if cancel.is_some_and(|c| c.is_cancelled()) {
430        return Err(ProviderError::Cancelled.into());
431    }
432
433    // Atomic registry lease held for the entire decode (JOE-1646).
434    let lease = pool.get_or_load_pin(model_path, model_name, gov)?;
435    let ctx = Arc::clone(lease.value());
436
437    let mut state = ctx
438        .create_state()
439        .map_err(|e| ProviderError::TranscriptionFailed {
440            reason: format!("failed to create whisper state: {e}"),
441        })?;
442
443    let mut params = FullParams::new(SamplingStrategy::Greedy { best_of: 1 });
444
445    params.set_print_special(false);
446    params.set_print_progress(false);
447    params.set_print_realtime(false);
448    params.set_print_timestamps(false);
449    params.set_suppress_blank(true);
450    params.set_suppress_nst(true);
451    params.set_token_timestamps(timestamps);
452    params.set_no_speech_thold(0.6);
453
454    // Keep Arc alive across full() so the C abort callback's pointer stays valid.
455    let _cancel_keep: Option<std::sync::Arc<std::sync::atomic::AtomicBool>> =
456        cancel.map(|f| f.clone_arc());
457    if let Some(ref arc) = _cancel_keep {
458        let ptr = std::sync::Arc::as_ptr(arc) as *mut std::ffi::c_void;
459        unsafe {
460            params.set_abort_callback(Some(abort_callback_atomic));
461            params.set_abort_callback_user_data(ptr);
462        }
463    }
464
465    params.set_n_threads(n_threads.clamp(1, 8));
466
467    let lang = language.trim().to_ascii_lowercase();
468    let auto = lang.is_empty() || lang == "auto";
469    if auto {
470        params.set_language(None);
471    } else {
472        params.set_language(Some(lang.as_str()));
473    }
474
475    let full_res = state.full(params, samples);
476    if cancel.is_some_and(|c| c.is_cancelled()) {
477        return Err(ProviderError::Cancelled.into());
478    }
479    full_res.map_err(|e| ProviderError::TranscriptionFailed {
480        reason: e.to_string(),
481    })?;
482
483    let mut segments = Vec::new();
484    let mut full_text = String::new();
485
486    for segment in state.as_iter() {
487        let start = segment.start_timestamp() as f64 / 100.0;
488        let end = segment.end_timestamp() as f64 / 100.0;
489        let text = segment.to_string();
490        let text = text.trim();
491        if text.is_empty() {
492            continue;
493        }
494        if !full_text.is_empty() {
495            full_text.push(' ');
496        }
497        full_text.push_str(text);
498        segments.push(Segment::from_parts_unchecked(start, end, text.to_string()));
499    }
500
501    let lang_id = state.full_lang_id_from_state();
502    let detected = if lang_id >= 0 {
503        whisper_rs::get_lang_str(lang_id).map(|s| s.to_string())
504    } else {
505        None
506    };
507
508    let result = TranscriptionResult::local(
509        full_text,
510        segments,
511        detected.or_else(|| {
512            if lang != "auto" && !lang.is_empty() {
513                Some(lang)
514            } else {
515                None
516            }
517        }),
518        model_name.to_string(),
519        duration_secs,
520    );
521    // Explicitly hold the registry lease until after decode completes.
522    drop(lease);
523    Ok(result)
524}
aurum_core/providers/local.rs

aurum_core/providers/
local.rs
