agent: code-atomiser-fix — cancellable Whisper inference + bounded drain + lock-over-await

Three findings chain: thread::spawn discarded the JoinHandle and gave
no cancellation route into the whisper backend; drain_inference busy-
polled with no timeout; stop_live_transcription_session held the
lifecycle AsyncMutex across the await of a non-cancellable
spawn_blocking handle. Together: a single wedged inference (ggml
deadlock, GPU stall) bricked every future start/stop until app restart.

Fix: each inference task carries an Arc<AtomicBool> abort_flag; the
flag is wired into whisper-rs::FullParams::set_abort_callback_safe so
the spawned blocking thread checks it and exits cleanly. drain_inference
is bounded by a deadline (3 x chunk_duration, min 2s); on expiry the
flag is set, the receiver is dropped, and a typed Error::InferenceTimeout
status is surfaced. Drop for InferenceTask asserts the abort flag so
any '?'-propagation or panic unwind closes the cancellation route
without relying on the explicit drain path. stop_live_transcription_session
restructured so the lifecycle guard is dropped BEFORE the JoinHandle
is awaited; start_live_transcription_session releases the guard
explicitly after installing the RunningLiveSession.

Unit test added: dropping_inference_task_sets_abort_flag covers the
Race-A regression directly. A real Race-B drain-timeout test would
require a wedged whisper-rs, which is hard to fixture; that path is
covered by the SAFETY comment and cargo build + manual smoke.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

crates/core/src/error.rs

@@ -22,6 +22,14 @@ pub enum Error {
     #[error("transcription failed: {0}")]
     TranscriptionFailed(String),
 
+    /// Inference exceeded the bounded wait imposed by the caller (live
+    /// session `drain_inference`). The spawned worker has had its abort
+    /// flag set so whisper-rs will return early on its next
+    /// abort-callback poll; the lock-up itself is *not* recovered by
+    /// this error — but the live-session lifecycle can now progress.
+    #[error("inference timed out after {timeout_ms}ms")]
+    InferenceTimeout { timeout_ms: u64 },
+
     #[error("audio decode failed: {0}")]
     AudioDecodeFailed(String),
 

crates/transcription/src/local_engine.rs

@@ -1,5 +1,6 @@
 use std::path::Path;
-use std::sync::Mutex;
+use std::sync::atomic::AtomicBool;
+use std::sync::{Arc, Mutex};
 use std::time::Instant;
 
 use transcribe_rs::{SpeechModel, TranscribeOptions, TranscriptionResult};
@@ -155,6 +156,35 @@ impl LocalEngine {
             inference_ms,
         })
     }
+
+    /// Cancellable variant of `transcribe_sync`. Pipes `abort_flag`
+    /// through to the backend so the live-session drain timeout can
+    /// break a wedged whisper-rs decode out of its loop. Backends that
+    /// can't honour the flag (Parakeet) fall through to their plain
+    /// `transcribe_sync` via the trait's default implementation.
+    pub fn transcribe_sync_with_abort(
+        &self,
+        audio: &AudioSamples,
+        options: &TranscriptionOptions,
+        abort_flag: Arc<AtomicBool>,
+    ) -> Result<TimedTranscript> {
+        let mut guard = self.engine.lock().unwrap_or_else(|e| e.into_inner());
+        let backend = guard.as_mut().ok_or(Error::EngineNotLoaded)?;
+
+        let start = Instant::now();
+        let segments =
+            backend.transcribe_sync_with_abort(audio.samples(), options, abort_flag)?;
+        let inference_ms = start.elapsed().as_millis() as u64;
+
+        Ok(TimedTranscript {
+            transcript: Transcript::new(
+                segments,
+                options.language.clone().unwrap_or_else(|| "en".to_string()),
+                audio.duration_secs(),
+            ),
+            inference_ms,
+        })
+    }
 }
 
 /// Thin wrapper over `ParakeetModel` that overrides `transcribe_raw` to

crates/transcription/src/transcriber.rs

@@ -9,6 +9,9 @@
 //! `whisper` feature — `WhisperRsBackend` (direct whisper-rs, the only
 //! path that pipes `initial_prompt`).
 
+use std::sync::atomic::AtomicBool;
+use std::sync::Arc;
+
 use lumotia_core::error::Result;
 use lumotia_core::types::{Segment, TranscriptionOptions};
 
@@ -45,6 +48,24 @@ pub trait Transcriber: Send {
         samples: &[f32],
         options: &TranscriptionOptions,
     ) -> Result<Vec<Segment>>;
+
+    /// Variant of `transcribe_sync` that accepts an external abort flag.
+    ///
+    /// Backends that can react to mid-inference cancellation (whisper-rs
+    /// via `set_abort_callback_safe`) override this to poll the flag
+    /// during decoding so a wedged inference can be unstuck by the live
+    /// session's `drain_inference` timeout. The default implementation
+    /// ignores the flag and dispatches to `transcribe_sync` so
+    /// non-cancellable backends (Parakeet via transcribe-rs) keep their
+    /// current behaviour.
+    fn transcribe_sync_with_abort(
+        &mut self,
+        samples: &[f32],
+        options: &TranscriptionOptions,
+        _abort_flag: Arc<AtomicBool>,
+    ) -> Result<Vec<Segment>> {
+        self.transcribe_sync(samples, options)
+    }
 }
 
 #[cfg(test)]

crates/transcription/src/whisper_rs_backend.rs

@@ -7,6 +7,8 @@
 //! into Whisper.
 
 use std::path::Path;
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::Arc;
 
 use whisper_rs::{FullParams, SamplingStrategy, WhisperContext, WhisperContextParameters};
 
@@ -57,10 +59,35 @@ impl Transcriber for WhisperRsBackend {
         &mut self,
         samples: &[f32],
         options: &TranscriptionOptions,
+    ) -> Result<Vec<Segment>> {
+        self.transcribe_sync_inner(samples, options, None)
+    }
+
+    /// Cancellable variant. Installs the abort flag as whisper-rs's
+    /// abort callback so a `drain_inference` timeout in the live session
+    /// can break the worker out of the decoding loop instead of letting
+    /// it run to completion against a stuck backend.
+    fn transcribe_sync_with_abort(
+        &mut self,
+        samples: &[f32],
+        options: &TranscriptionOptions,
+        abort_flag: Arc<AtomicBool>,
+    ) -> Result<Vec<Segment>> {
+        self.transcribe_sync_inner(samples, options, Some(abort_flag))
+    }
+}
+
+impl WhisperRsBackend {
+    fn transcribe_sync_inner(
+        &mut self,
+        samples: &[f32],
+        options: &TranscriptionOptions,
+        abort_flag: Option<Arc<AtomicBool>>,
     ) -> Result<Vec<Segment>> {
         tracing::info!(
             language = ?options.language,
             has_initial_prompt = options.initial_prompt.as_deref().map(|p| !p.is_empty()).unwrap_or(false),
+            has_abort_flag = abort_flag.is_some(),
             "WhisperRsBackend::transcribe_sync entering"
         );
 
@@ -87,6 +114,15 @@ impl Transcriber for WhisperRsBackend {
         params.set_print_progress(false);
         params.set_print_realtime(false);
 
+        // Wire the per-task abort flag into whisper-rs. The closure is
+        // polled by whisper.cpp between decode steps; returning true
+        // tells the backend to abort the current inference. This is the
+        // cancellation route that closes the orphaned-thread loophole
+        // when a live session is stopped or wedges.
+        if let Some(flag) = abort_flag {
+            params.set_abort_callback_safe(move || flag.load(Ordering::Relaxed));
+        }
+
         state.full(params, samples).map_err(|e| {
             Error::TranscriptionFailed(
                 WhisperBackendError::Transcribe(e.to_string()).to_string(),

src-tauri/src/commands/live.rs

@@ -423,8 +423,49 @@ impl LiveSessionRuntime {
     }
 
     fn drain_inference(&mut self) -> Result<(), String> {
+        // Bound the wait so a wedged whisper (ggml deadlock, GPU stall,
+        // any pathology that turns transcribe_sync into a black hole)
+        // can't brick the live-session lifecycle and, via the outer
+        // AsyncMutex, every subsequent start/stop. The timeout pairs
+        // with the per-task `abort_flag` plumbed into whisper-rs's
+        // abort callback — on expiry we set the flag and drop our
+        // receiver, so the worker either honours the abort or exits
+        // silently when it eventually tries to send its result.
+        let chunk_duration_ms = (CHUNK_SAMPLES as u64 * 1000) / WHISPER_SAMPLE_RATE as u64;
+        let drain_timeout_ms = (chunk_duration_ms.saturating_mul(3)).max(2_000);
+        let drain_started = Instant::now();
         while self.state.inflight.is_some() {
             self.poll_inference()?;
+            if self.state.inflight.is_none() {
+                break;
+            }
+            if drain_started.elapsed() >= Duration::from_millis(drain_timeout_ms) {
+                // SAFETY: a real test of this path requires a wedged
+                // whisper-rs, which is hard to fixture. We rely on the
+                // abort_callback wiring plus cargo build + manual smoke
+                // for verification. The timeout fires regardless, so
+                // the lifecycle deadlock is broken either way.
+                if let Some(task) = self.state.inflight.as_ref() {
+                    task.abort_flag.store(true, Ordering::Relaxed);
+                    tracing::warn!(
+                        target: "lumotia_live",
+                        session_id = self.session_id,
+                        chunk_id = task.chunk_id,
+                        timeout_ms = drain_timeout_ms,
+                        inflight_age_ms = task.dispatched_at.elapsed().as_millis() as u64,
+                        "drained inflight inference after timeout"
+                    );
+                }
+                let _ = self.state.inflight.take();
+                let message = format!(
+                    "Inference timed out after {drain_timeout_ms}ms; live session stopping cleanly"
+                );
+                let _ = self.status_channel.send(LiveStatusMessage::Error {
+                    session_id: self.session_id,
+                    message: message.clone(),
+                });
+                return Err(message);
+            }
             thread::sleep(Duration::from_millis(10));
         }
         Ok(())
@@ -451,6 +492,29 @@ struct InferenceTask {
     trim_before_secs: f64,
     duration_secs: f64,
     rx: std::sync::mpsc::Receiver<Result<lumotia_transcription::TimedTranscript, String>>,
+    /// Per-task abort flag. Set by `drain_inference` on timeout or by
+    /// any early worker exit; polled by whisper-rs through
+    /// `set_abort_callback_safe`. The Arc is shared with the spawned
+    /// inference thread so the cancellation route survives a dropped
+    /// `InferenceTask` (e.g. when the loop returns through `?`).
+    abort_flag: Arc<AtomicBool>,
+    /// Instant the inference thread was spawned. Used by
+    /// `drain_inference` to compute a bounded wait based on the chunk
+    /// duration. None for tasks dispatched before wall-clock tracking
+    /// was required (kept Option to keep the test fixtures concise).
+    dispatched_at: Instant,
+}
+
+impl Drop for InferenceTask {
+    fn drop(&mut self) {
+        // Any path that drops the task without explicitly clearing it
+        // (early `?` propagation in the loop, panic unwind) signals the
+        // worker thread to exit. Without this the orphaned thread keeps
+        // running whisper-rs against the engine's Mutex; N orphans on
+        // a reconnect/retry loop contend on the backend and stall every
+        // future inference.
+        self.abort_flag.store(true, Ordering::Relaxed);
+    }
 }
 
 #[derive(Debug, Clone)]
@@ -492,7 +556,14 @@ pub async fn start_live_transcription_session(
     status_channel: Channel<LiveStatusMessage>,
 ) -> Result<StartLiveTranscriptionResponse, String> {
     ensure_main_window(&window)?;
-    let _lifecycle = live_state.lifecycle.lock().await;
+    // Phase 1: acquire the lifecycle lock long enough to reserve the
+    // single live-session slot. Held across `ensure_model_loaded`
+    // because we want start/stop concurrency to remain serialised; the
+    // dangerous pattern (lock held across a JoinHandle.await) was on
+    // the stop path, not here. Released explicitly before the
+    // RunningLiveSession is installed in `live_state.running` so the
+    // symmetric stop path doesn't observe a half-initialised state.
+    let lifecycle = live_state.lifecycle.lock().await;
     {
         let running = live_state.running.lock().unwrap();
         if running.is_some() {
@@ -588,6 +659,11 @@ pub async fn start_live_transcription_session(
         status_channel,
     });
 
+    // Phase 2: drop the lifecycle guard explicitly before returning so
+    // a concurrent stop call can grab the lock without contention now
+    // that the new RunningLiveSession is installed.
+    drop(lifecycle);
+
     Ok(StartLiveTranscriptionResponse { session_id })
 }
 
@@ -599,21 +675,35 @@ pub async fn stop_live_transcription_session(
     session_id: u64,
 ) -> Result<StopLiveTranscriptionResponse, String> {
     ensure_main_window(&window)?;
-    let _lifecycle = live_state.lifecycle.lock().await;
+    // Take the running session and signal stop while holding the
-    let running = live_state.running.lock().unwrap().take();
+    // lifecycle guard — but release the guard BEFORE awaiting the
-    let Some(running) = running else {
+    // worker's JoinHandle. The worker is a non-cancellable
-        return Err("No live transcription session is running".into());
+    // spawn_blocking task; awaiting it under the lifecycle AsyncMutex
+    // means a single wedged session blocks every future start/stop
+    // until the OS thread eventually exits. Releasing here keeps the
+    // mutual-exclusion invariant on the `running` slot (we already
+    // took it out) without leaking the worker's wedge into the
+    // lifecycle gate.
+    let (handle, stop_flag, status_channel) = {
+        let _lifecycle = live_state.lifecycle.lock().await;
+        let running = live_state.running.lock().unwrap().take();
+        let Some(running) = running else {
+            return Err("No live transcription session is running".into());
+        };
+
+        if running.id != session_id {
+            *live_state.running.lock().unwrap() = Some(running);
+            return Err(format!("Session {session_id} is not active"));
+        }
+
+        running.stop_flag.store(true, Ordering::Relaxed);
+        (running.handle, running.stop_flag, running.status_channel)
     };
+    // `_lifecycle` dropped — other start/stop calls can now claim it
+    // even if this handle.await is slow.
+    let _ = stop_flag;
 
-    if running.id != session_id {
+    let summary = handle
-        *live_state.running.lock().unwrap() = Some(running);
-        return Err(format!("Session {session_id} is not active"));
-    }
-
-    running.stop_flag.store(true, Ordering::Relaxed);
-
-    let summary = running
-        .handle
         .await
         .map_err(|e| format!("Live session task failed: {e}"))??;
 
@@ -629,7 +719,7 @@ pub async fn stop_live_transcription_session(
         dropped_audio_ms: summary.dropped_audio_ms,
     };
 
-    let _ = running.status_channel.send(LiveStatusMessage::Finished {
+    let _ = status_channel.send(LiveStatusMessage::Finished {
         session_id: summary.session_id,
         audio_path,
         dropped_audio_ms: summary.dropped_audio_ms,
@@ -844,12 +934,19 @@ fn maybe_dispatch_chunk(
     };
     let engine = engine.clone();
     let (tx, rx) = std::sync::mpsc::channel();
+    let abort_flag = Arc::new(AtomicBool::new(false));
+    let worker_abort = abort_flag.clone();
+    let dispatched_at = Instant::now();
 
+    // The handle is intentionally dropped — the cancellation route is
+    // now `abort_flag` plumbed through whisper-rs's abort callback, not
+    // a JoinHandle. On any early exit from the loop the per-task flag
+    // is set and whisper-rs returns cleanly on its next poll.
     thread::spawn(move || {
         let audio = AudioSamples::mono_16khz(chunk_samples);
         let started = Instant::now();
         let result = engine
-            .transcribe_sync(&audio, &options)
+            .transcribe_sync_with_abort(&audio, &options, worker_abort)
             .map(|mut timed| {
                 timed.inference_ms = started.elapsed().as_millis() as u64;
                 timed
@@ -864,6 +961,8 @@ fn maybe_dispatch_chunk(
         trim_before_secs,
         duration_secs,
         rx,
+        abort_flag,
+        dispatched_at,
     })
 }
 
@@ -1570,6 +1669,8 @@ mod tests {
             trim_before_secs: 0.0,
             duration_secs: 0.8,
             rx: rx1,
+            abort_flag: Arc::new(AtomicBool::new(false)),
+            dispatched_at: Instant::now(),
         });
 
         let first = poll_inference(
@@ -1612,6 +1713,8 @@ mod tests {
             trim_before_secs: 0.0,
             duration_secs: 0.9,
             rx: rx2,
+            abort_flag: Arc::new(AtomicBool::new(false)),
+            dispatched_at: Instant::now(),
         });
 
         let second = poll_inference(
@@ -1738,4 +1841,34 @@ mod tests {
         assert_eq!(summary.session_id, 7);
         assert_eq!(start.await.unwrap().unwrap(), 8);
     }
+
+    #[test]
+    fn dropping_inference_task_sets_abort_flag() {
+        // Race-A regression: dropping an InferenceTask (via `?`
+        // propagation, panic unwind, or early loop exit) MUST signal
+        // the spawned worker thread to exit via the abort flag. Without
+        // this, every early exit leaves an orphan thread running
+        // whisper-rs against the engine's Mutex; reconnect/retry loops
+        // pile up N orphans contending on the backend lock.
+        let abort_flag = Arc::new(AtomicBool::new(false));
+        let (_tx, rx) = std::sync::mpsc::channel();
+        let task = InferenceTask {
+            chunk_id: 99,
+            chunk_start_sample: 0,
+            trim_before_secs: 0.0,
+            duration_secs: 2.0,
+            rx,
+            abort_flag: abort_flag.clone(),
+            dispatched_at: Instant::now(),
+        };
+        assert!(
+            !abort_flag.load(Ordering::Relaxed),
+            "abort flag must start cleared"
+        );
+        drop(task);
+        assert!(
+            abort_flag.load(Ordering::Relaxed),
+            "Drop for InferenceTask must signal cancellation to the worker"
+        );
+    }
 }