Lumotia/src-tauri/src/commands/transcription.rs

// Tauri command handlers must match the frontend's invoke() parameter lists,
// so the argument counts are dictated by the Svelte code.
#![allow(clippy::too_many_arguments)]

use std::path::Path;
use std::sync::Arc;

use crate::commands::build_initial_prompt;
use crate::commands::models::{default_model_id_for_engine, ensure_model_loaded};
use crate::commands::security::ensure_main_window;
use crate::AppState;
use lumotia_ai_formatting::{post_process_segments, FormatMode, PostProcessOptions};
use lumotia_core::constants::WHISPER_SAMPLE_RATE;
use lumotia_core::types::{AudioSamples, Segment, Transcript, TranscriptionOptions};

const PARAKEET_CHUNK_THRESHOLD_SECS: usize = 18;
const PARAKEET_CHUNK_SECS: usize = 15;
const PARAKEET_CHUNK_OVERLAP_SECS: usize = 1;
const FILE_CHUNK_THRESHOLD_SECS: usize = 8 * 60;
const FILE_CHUNK_SECS: usize = 3 * 60;
const FILE_CHUNK_OVERLAP_SECS: usize = 2;
const MAX_FILE_TRANSCRIPTION_SECS: f64 = 2.0 * 60.0 * 60.0;

/// Trust-5 (conf 80, code-atomiser-fix 2026-05-12): defence-in-depth
/// against arbitrary-path / arbitrary-blob feeds into the audio decoder.
/// The OS file picker already constrains the user's typical path, but
/// the IPC surface itself accepted any string — we add an extension
/// allowlist and a byte-size cap so a compromised webview can't point
/// the decoder at e.g. a 50 GiB sparse file or a `.so` payload chosen
/// to provoke a parser bug in symphonia.
///
/// The allowlist is the Symphonia-supported subset Lumotia ships with;
/// users who want to transcribe something exotic can convert it first.
const ALLOWED_AUDIO_EXTENSIONS: &[&str] = &[
    "wav", "mp3", "m4a", "mp4", "flac", "ogg", "opus", "webm", "aac",
];

/// 1 GiB ceiling on the input file. Two hours of 48kHz WAV stereo is
/// ~1.4 GiB, but we bound transcription to 2 hours via
/// `MAX_FILE_TRANSCRIPTION_SECS` which most realistic compressed
/// formats (MP3, M4A, Opus) clear by an order of magnitude. A 1 GiB
/// cap leaves headroom for lossless WAV at typical mono-16kHz capture
/// rates while still rejecting absurd inputs that would OOM the
/// decoder before our duration check fires.
const MAX_TRANSCRIBE_BYTES: u64 = 1024 * 1024 * 1024;

/// Pure validator for Trust-5. Pulled out of `transcribe_file` so the
/// rule can be unit-tested without spawning Tauri or hitting disk for
/// the decoder. Order of checks matters: extension first (cheap,
/// rejects most obvious attacks), then size (one stat call), then
/// path canonicalisation (slowest, only meaningful once the cheaper
/// gates pass).
pub(crate) fn validate_transcribe_input(path: &Path, metadata_len: u64) -> Result<(), String> {
    let ext = path
        .extension()
        .and_then(|s| s.to_str())
        .map(|s| s.to_ascii_lowercase())
        .ok_or_else(|| {
            format!(
                "Refusing to transcribe {}: file has no extension. \
                 Supported: {}.",
                path.display(),
                ALLOWED_AUDIO_EXTENSIONS.join(", ")
            )
        })?;

    if !ALLOWED_AUDIO_EXTENSIONS
        .iter()
        .any(|allowed| *allowed == ext)
    {
        return Err(format!(
            "Refusing to transcribe {}: extension '.{}' is not in the \
             allowlist. Supported: {}.",
            path.display(),
            ext,
            ALLOWED_AUDIO_EXTENSIONS.join(", ")
        ));
    }

    if metadata_len > MAX_TRANSCRIBE_BYTES {
        return Err(format!(
            "Refusing to transcribe {}: file is {} bytes, limit is {} bytes \
             (1 GiB).",
            path.display(),
            metadata_len,
            MAX_TRANSCRIBE_BYTES
        ));
    }

    Ok(())
}

struct ChunkingStrategy {
    chunk_samples: usize,
    overlap_samples: usize,
}

fn pick_engine(
    state: &AppState,
    engine: &str,
) -> Result<Arc<lumotia_transcription::LocalEngine>, String> {
    match engine {
        "whisper" => Ok(state.whisper_engine.clone()),
        "parakeet" => Ok(state.parakeet_engine.clone()),
        other => Err(format!("Unknown engine: {other}")),
    }
}

fn pick_chunking_strategy(engine_name: &str, sample_count: usize) -> Option<ChunkingStrategy> {
    let samples_per_second = WHISPER_SAMPLE_RATE as usize;
    match engine_name {
        "parakeet" if sample_count > PARAKEET_CHUNK_THRESHOLD_SECS * samples_per_second => {
            Some(ChunkingStrategy {
                chunk_samples: PARAKEET_CHUNK_SECS * samples_per_second,
                overlap_samples: PARAKEET_CHUNK_OVERLAP_SECS * samples_per_second,
            })
        }
        _ if sample_count > FILE_CHUNK_THRESHOLD_SECS * samples_per_second => {
            Some(ChunkingStrategy {
                chunk_samples: FILE_CHUNK_SECS * samples_per_second,
                overlap_samples: FILE_CHUNK_OVERLAP_SECS * samples_per_second,
            })
        }
        _ => None,
    }
}

fn trim_overlap_segments(segments: &mut Vec<Segment>, trim_before_secs: f64) {
    if trim_before_secs <= 0.0 {
        return;
    }

    segments.retain(|segment| segment.end > trim_before_secs);
    for segment in segments.iter_mut() {
        if segment.start < trim_before_secs {
            segment.start = trim_before_secs;
        }
    }
}

fn transcribe_samples_sync(
    engine: Arc<lumotia_transcription::LocalEngine>,
    engine_name: &str,
    samples: Vec<f32>,
    options: TranscriptionOptions,
) -> Result<lumotia_transcription::TimedTranscript, String> {
    let Some(strategy) = pick_chunking_strategy(engine_name, samples.len()) else {
        let audio = AudioSamples::mono_16khz(samples);
        return engine
            .transcribe_sync(&audio, &options)
            .map_err(|e| e.to_string());
    };

    let total_duration_secs = samples.len() as f64 / WHISPER_SAMPLE_RATE as f64;
    let stride = strategy
        .chunk_samples
        .saturating_sub(strategy.overlap_samples)
        .max(1);
    let chunk_count = ((samples.len().saturating_sub(1)) / stride) + 1;
    tracing::info!(
        engine = engine_name,
        duration_secs = total_duration_secs,
        chunk_count,
        "chunking audio for transcription"
    );

    let mut all_segments = Vec::new();
    let mut total_inference_ms = 0u64;
    let mut chunk_start = 0usize;

    while chunk_start < samples.len() {
        let chunk_end = (chunk_start + strategy.chunk_samples).min(samples.len());
        let chunk_audio = AudioSamples::mono_16khz(samples[chunk_start..chunk_end].to_vec());
        let timed = engine
            .transcribe_sync(&chunk_audio, &options)
            .map_err(|e| e.to_string())?;
        total_inference_ms = total_inference_ms.saturating_add(timed.inference_ms);

        let mut chunk_segments = timed.transcript.segments().to_vec();
        if chunk_start > 0 {
            trim_overlap_segments(
                &mut chunk_segments,
                strategy.overlap_samples as f64 / WHISPER_SAMPLE_RATE as f64,
            );
        }

        let chunk_offset_secs = chunk_start as f64 / WHISPER_SAMPLE_RATE as f64;
        for segment in &mut chunk_segments {
            segment.start += chunk_offset_secs;
            segment.end += chunk_offset_secs;
        }
        all_segments.extend(chunk_segments);

        if chunk_end >= samples.len() {
            break;
        }
        chunk_start = chunk_end.saturating_sub(strategy.overlap_samples);
    }

    Ok(lumotia_transcription::TimedTranscript {
        transcript: Transcript::new(
            all_segments,
            options.language.clone().unwrap_or_else(|| "en".to_string()),
            total_duration_secs,
        ),
        inference_ms: total_inference_ms,
    })
}

fn join_segment_text(segments: &[Segment]) -> String {
    segments
        .iter()
        .map(|segment| segment.text.trim())
        .filter(|segment| !segment.is_empty())
        .collect::<Vec<_>>()
        .join(" ")
}

/// Transcribe an audio file by path. Decodes, resamples to 16kHz, runs Whisper.
#[tauri::command]
pub async fn transcribe_file(
    window: tauri::WebviewWindow,
    state: tauri::State<'_, AppState>,
    path: String,
    engine: Option<String>,
    model_id: Option<String>,
    language: String,
    initial_prompt: String,
    remove_fillers: bool,
    british_english: bool,
    anti_hallucination: bool,
    format_mode: String,
    profile_id: Option<String>,
) -> Result<serde_json::Value, String> {
    ensure_main_window(&window)?;

    // Trust-5: extension + size gate before we hand the path to the
    // audio decoder. `std::fs::metadata` resolves symlinks so the size
    // check sees the actual blob, not a symlink-target lie.
    let metadata = std::fs::metadata(&path).map_err(|e| format!("Cannot stat {path}: {e}"))?;
    validate_transcribe_input(Path::new(&path), metadata.len())?;

    let resolved_profile_id =
        profile_id.unwrap_or_else(|| lumotia_storage::DEFAULT_PROFILE_ID.to_string());

    let profile = lumotia_storage::database::get_profile(&state.db, &resolved_profile_id)
        .await
        .map_err(|e| e.to_string())?
        .ok_or_else(|| format!("Profile {resolved_profile_id} not found"))?;

    let profile_terms: Vec<String> =
        lumotia_storage::database::list_profile_terms(&state.db, &resolved_profile_id)
            .await
            .map_err(|e| e.to_string())?
            .into_iter()
            .map(|t| t.term)
            .collect();

    let engine_name = engine.unwrap_or_else(|| "whisper".to_string());
    let model_id =
        model_id.unwrap_or_else(|| default_model_id_for_engine(&engine_name).to_string());
    // None: transcribe paths don't enforce sequential-GPU mode. That's
    // owned by the Settings-level load flows (see models.rs).
    ensure_model_loaded(&state, &engine_name, &model_id, None).await?;

    let engine = pick_engine(&state, &engine_name)?;
    let options = TranscriptionOptions {
        language: Some(language),
        initial_prompt: build_initial_prompt(
            &initial_prompt,
            &profile.initial_prompt,
            &profile_terms,
        ),
    };
    let engine_name_for_worker = engine_name.clone();
    let path_for_probe = Path::new(&path);
    if let Some(duration_secs) =
        lumotia_audio::probe_audio_duration_secs(path_for_probe).map_err(|e| e.to_string())?
    {
        if duration_secs > MAX_FILE_TRANSCRIPTION_SECS {
            return Err(format!(
                "File is {:.1} hours long. Lumotia imports up to 2 hours at a time.",
                duration_secs / 3600.0
            ));
        }
    }

    let timed = tokio::task::spawn_blocking(move || {
        let audio = lumotia_audio::decode_audio_file_limited(
            Path::new(&path),
            Some(MAX_FILE_TRANSCRIPTION_SECS),
        )
        .map_err(|e| e.to_string())?;
        let resampled = lumotia_audio::resample_to_16khz(&audio).map_err(|e| e.to_string())?;
        transcribe_samples_sync(
            engine,
            &engine_name_for_worker,
            resampled.into_samples(),
            options,
        )
    })
    .await
    .map_err(|e| e.to_string())??;

    let dictionary_terms = profile_terms.clone();

    let mut segments: Vec<Segment> = timed.transcript.segments().to_vec();
    let raw_text = join_segment_text(&segments);
    post_process_segments(
        &mut segments,
        &PostProcessOptions {
            remove_fillers,
            british_english,
            anti_hallucination,
            format_mode: FormatMode::parse(&format_mode),
            dictionary_terms,
        },
        Some(state.llm_engine.as_ref()),
    );

    Ok(serde_json::json!({
        "engine": engine_name,
        "modelId": model_id,
        "segments": segments,
        "language": timed.transcript.language(),
        "duration": timed.transcript.duration(),
        "inference_ms": timed.inference_ms,
        "raw_text": raw_text,
    }))
}

#[cfg(test)]
mod tests_trust5 {
    use super::{validate_transcribe_input, MAX_TRANSCRIBE_BYTES};
    use std::path::Path;

    #[test]
    fn accepts_supported_wav() {
        assert!(validate_transcribe_input(Path::new("/tmp/clip.wav"), 1024).is_ok());
    }

    #[test]
    fn accepts_supported_mp3_case_insensitive() {
        assert!(validate_transcribe_input(Path::new("/tmp/clip.MP3"), 1024).is_ok());
    }

    #[test]
    fn accepts_each_allowed_extension() {
        for ext in [
            "wav", "mp3", "m4a", "mp4", "flac", "ogg", "opus", "webm", "aac",
        ] {
            let path_string = format!("/tmp/clip.{ext}");
            let path = Path::new(&path_string);
            assert!(
                validate_transcribe_input(path, 1024).is_ok(),
                "extension '{ext}' should be allowed",
            );
        }
    }

    #[test]
    fn rejects_unsupported_extension() {
        let result = validate_transcribe_input(Path::new("/tmp/payload.so"), 1024);
        let err = result.expect_err("must reject .so");
        assert!(
            err.contains("not in the allowlist"),
            "unexpected error: {err}"
        );
    }

    #[test]
    fn rejects_no_extension() {
        let result = validate_transcribe_input(Path::new("/tmp/payload"), 1024);
        let err = result.expect_err("must reject missing extension");
        assert!(err.contains("no extension"), "unexpected error: {err}");
    }

    #[test]
    fn rejects_oversize_file() {
        let result =
            validate_transcribe_input(Path::new("/tmp/huge.wav"), MAX_TRANSCRIBE_BYTES + 1);
        let err = result.expect_err("must reject oversize file");
        assert!(err.contains("1 GiB"), "unexpected error: {err}");
    }

    #[test]
    fn accepts_file_exactly_at_size_cap() {
        // Exactly-at-cap must pass — the rule is "greater than cap rejects".
        assert!(
            validate_transcribe_input(Path::new("/tmp/edge.wav"), MAX_TRANSCRIBE_BYTES).is_ok()
        );
    }

    #[test]
    fn rejects_traversal_with_disallowed_extension() {
        // A "../../etc/passwd" payload would also fail the extension
        // gate even if the OS dialog somehow let it through.
        let result = validate_transcribe_input(Path::new("../../etc/passwd"), 1024);
        assert!(result.is_err());
    }
}