ensure samples don't get dropped (#1043)

* ensure samples don't get dropped

* format

src-tauri/src/audio_toolkit/audio/recorder.rs

@@ -1,6 +1,9 @@
 use std::{
     io::Error,
-    sync::{mpsc, Arc, Mutex},
+    sync::{
+        atomic::{AtomicBool, Ordering},
+        mpsc, Arc, Mutex,
+    },
     time::Duration,
 };
 
@@ -22,6 +25,11 @@ enum Cmd {
     Shutdown,
 }
 
+enum AudioChunk {
+    Samples(Vec<f32>),
+    EndOfStream,
+}
+
 pub struct AudioRecorder {
     device: Option<Device>,
     cmd_tx: Option<mpsc::Sender<Cmd>>,
@@ -59,7 +67,7 @@ impl AudioRecorder {
             return Ok(()); // already open
         }
 
-        let (sample_tx, sample_rx) = mpsc::channel::<Vec<f32>>();
+        let (sample_tx, sample_rx) = mpsc::channel::<AudioChunk>();
         let (cmd_tx, cmd_rx) = mpsc::channel::<Cmd>();
         let (init_tx, init_rx) = mpsc::sync_channel::<Result<(), String>>(1);
 
@@ -77,6 +85,8 @@ impl AudioRecorder {
         let level_cb = self.level_cb.clone();
 
         let worker = std::thread::spawn(move || {
+            let stop_flag = Arc::new(AtomicBool::new(false));
+            let stop_flag_for_stream = stop_flag.clone();
             let init_result = (|| -> Result<(cpal::Stream, u32), String> {
                 let config = AudioRecorder::get_preferred_config(&thread_device)
                     .map_err(|e| format!("Failed to fetch preferred config: {e}"))?;
@@ -98,6 +108,7 @@ impl AudioRecorder {
                         &config,
                         sample_tx,
                         channels,
+                        stop_flag_for_stream,
                     )
                     .map_err(|e| format!("Failed to build input stream: {e}"))?,
                     cpal::SampleFormat::I8 => AudioRecorder::build_stream::<i8>(
@@ -105,6 +116,7 @@ impl AudioRecorder {
                         &config,
                         sample_tx,
                         channels,
+                        stop_flag_for_stream,
                     )
                     .map_err(|e| format!("Failed to build input stream: {e}"))?,
                     cpal::SampleFormat::I16 => AudioRecorder::build_stream::<i16>(
@@ -112,6 +124,7 @@ impl AudioRecorder {
                         &config,
                         sample_tx,
                         channels,
+                        stop_flag_for_stream,
                     )
                     .map_err(|e| format!("Failed to build input stream: {e}"))?,
                     cpal::SampleFormat::I32 => AudioRecorder::build_stream::<i32>(
@@ -119,6 +132,7 @@ impl AudioRecorder {
                         &config,
                         sample_tx,
                         channels,
+                        stop_flag_for_stream,
                     )
                     .map_err(|e| format!("Failed to build input stream: {e}"))?,
                     cpal::SampleFormat::F32 => AudioRecorder::build_stream::<f32>(
@@ -126,6 +140,7 @@ impl AudioRecorder {
                         &config,
                         sample_tx,
                         channels,
+                        stop_flag_for_stream,
                     )
                     .map_err(|e| format!("Failed to build input stream: {e}"))?,
                     sample_format => {
@@ -144,7 +159,7 @@ impl AudioRecorder {
                 Ok((stream, sample_rate)) => {
                     let _ = init_tx.send(Ok(()));
                     // Keep the stream alive while we process samples.
-                    run_consumer(sample_rate, vad, sample_rx, cmd_rx, level_cb);
+                    run_consumer(sample_rate, vad, sample_rx, cmd_rx, level_cb, stop_flag);
                     drop(stream);
                 }
                 Err(error_message) => {
@@ -209,23 +224,32 @@ impl AudioRecorder {
     fn build_stream<T>(
         device: &cpal::Device,
         config: &cpal::SupportedStreamConfig,
-        sample_tx: mpsc::Sender<Vec<f32>>,
+        sample_tx: mpsc::Sender<AudioChunk>,
         channels: usize,
+        stop_flag: Arc<AtomicBool>,
     ) -> Result<cpal::Stream, cpal::BuildStreamError>
     where
         T: Sample + SizedSample + Send + 'static,
         f32: cpal::FromSample<T>,
     {
         let mut output_buffer = Vec::new();
+        let mut eos_sent = false;
 
         let stream_cb = move |data: &[T], _: &cpal::InputCallbackInfo| {
+            if stop_flag.load(Ordering::Relaxed) {
+                if !eos_sent {
+                    let _ = sample_tx.send(AudioChunk::EndOfStream);
+                    eos_sent = true;
+                }
+                return;
+            }
+            eos_sent = false;
+
             output_buffer.clear();
 
             if channels == 1 {
-                // Direct conversion without intermediate Vec
                 output_buffer.extend(data.iter().map(|&sample| sample.to_sample::<f32>()));
             } else {
-                // Convert to mono directly
                 let frame_count = data.len() / channels;
                 output_buffer.reserve(frame_count);
 
@@ -239,7 +263,10 @@ impl AudioRecorder {
                 }
             }
 
-            if sample_tx.send(output_buffer.clone()).is_err() {
+            if sample_tx
+                .send(AudioChunk::Samples(output_buffer.clone()))
+                .is_err()
+            {
                 log::error!("Failed to send samples");
             }
         };
@@ -326,9 +353,10 @@ mod tests {
 fn run_consumer(
     in_sample_rate: u32,
     vad: Option<Arc<Mutex<Box<dyn vad::VoiceActivityDetector>>>>,
-    sample_rx: mpsc::Receiver<Vec<f32>>,
+    sample_rx: mpsc::Receiver<AudioChunk>,
     cmd_rx: mpsc::Receiver<Cmd>,
     level_cb: Option<Arc<dyn Fn(Vec<f32>) + Send + Sync + 'static>>,
+    stop_flag: Arc<AtomicBool>,
 ) {
     let mut frame_resampler = FrameResampler::new(
         in_sample_rate as usize,
@@ -372,11 +400,16 @@ fn run_consumer(
     }
 
     loop {
-        let raw = match sample_rx.recv() {
+        let chunk = match sample_rx.recv() {
-            Ok(s) => s,
+            Ok(c) => c,
             Err(_) => break, // stream closed
         };
 
+        let raw = match chunk {
+            AudioChunk::Samples(s) => s,
+            AudioChunk::EndOfStream => continue,
+        };
+
         // ---------- spectrum processing ---------------------------------- //
         if let Some(buckets) = visualizer.feed(&raw) {
             if let Some(cb) = &level_cb {
@@ -393,21 +426,35 @@ fn run_consumer(
         while let Ok(cmd) = cmd_rx.try_recv() {
             match cmd {
                 Cmd::Start => {
+                    stop_flag.store(false, Ordering::Relaxed);
                     processed_samples.clear();
                     recording = true;
-                    visualizer.reset(); // Reset visualization buffer
+                    visualizer.reset();
                     if let Some(v) = &vad {
                         v.lock().unwrap().reset();
                     }
                 }
                 Cmd::Stop(reply_tx) => {
                     recording = false;
+                    stop_flag.store(true, Ordering::Relaxed);
 
-                    // Drain any audio chunks that were captured but not yet consumed
+                    // Drain all remaining audio until the producer confirms end-of-stream.
-                    while let Ok(remaining) = sample_rx.try_recv() {
+                    // The cpal callback sees the stop flag, sends EndOfStream, and goes
-                        frame_resampler.push(&remaining, &mut |frame: &[f32]| {
+                    // silent — guaranteeing every captured sample is in the channel
-                            handle_frame(frame, true, &vad, &mut processed_samples)
+                    // ahead of the sentinel.
-                        });
+                    loop {
+                        match sample_rx.recv_timeout(Duration::from_secs(2)) {
+                            Ok(AudioChunk::Samples(remaining)) => {
+                                frame_resampler.push(&remaining, &mut |frame: &[f32]| {
+                                    handle_frame(frame, true, &vad, &mut processed_samples)
+                                });
+                            }
+                            Ok(AudioChunk::EndOfStream) => break,
+                            Err(_) => {
+                                log::warn!("Timed out waiting for EndOfStream from audio callback");
+                                break;
+                            }
+                        }
                     }
 
                     frame_resampler.finish(&mut |frame: &[f32]| {
@@ -415,8 +462,15 @@ fn run_consumer(
                     });
 
                     let _ = reply_tx.send(std::mem::take(&mut processed_samples));
+
+                    // Resume the audio callback so the consumer loop can continue
+                    // receiving chunks (important for always-on microphone mode).
+                    stop_flag.store(false, Ordering::Relaxed);
+                }
+                Cmd::Shutdown => {
+                    stop_flag.store(true, Ordering::Relaxed);
+                    return;
                 }
-                Cmd::Shutdown => return,
             }
         }
     }